Rework checkpoint/restore code based on review feedback

Signed-off-by: Saied Kazemi <saied@google.com>