Fix removing state recover.

Signed-off-by: Lantao Liu <lantaol@google.com>

pkg/store/container/fake_status.go

@@ -38,6 +38,10 @@ func (f *fakeStatusStorage) Get() Status {
 	return f.status
 }
 
+func (f *fakeStatusStorage) UpdateSync(u UpdateFunc) error {
+	return f.Update(u)
+}
+
 func (f *fakeStatusStorage) Update(u UpdateFunc) error {
 	f.Lock()
 	defer f.Unlock()

pkg/store/container/status.go

@@ -106,11 +106,11 @@ type UpdateFunc func(Status) (Status, error)
 type StatusStorage interface {
 	// Get a container status.
 	Get() Status
+	// UpdateSync updates the container status and the on disk checkpoint.
+	// Note that the update MUST be applied in one transaction.
+	UpdateSync(UpdateFunc) error
 	// Update the container status. Note that the update MUST be applied
 	// in one transaction.
-	// TODO(random-liu): Distinguish `UpdateSync` and `Update`, only
-	// `UpdateSync` should sync data onto disk, so that disk operation
-	// for non-critical status change could be avoided.
 	Update(UpdateFunc) error
 	// Delete the container status.
 	// Note:
@@ -167,8 +167,8 @@ func (s *statusStorage) Get() Status {
 	return s.status
 }
 
-// Update the container status.
-func (s *statusStorage) Update(u UpdateFunc) error {
+// UpdateSync updates the container status and the on disk checkpoint.
+func (s *statusStorage) UpdateSync(u UpdateFunc) error {
 	s.Lock()
 	defer s.Unlock()
 	newStatus, err := u(s.status)
@@ -186,6 +186,18 @@ func (s *statusStorage) Update(u UpdateFunc) error {
 	return nil
 }
 
+// Update the container status.
+func (s *statusStorage) Update(u UpdateFunc) error {
+	s.Lock()
+	defer s.Unlock()
+	newStatus, err := u(s.status)
+	if err != nil {
+		return err
+	}
+	s.status = newStatus
+	return nil
+}
+
 // Delete deletes the container status from disk atomically.
 func (s *statusStorage) Delete() error {
 	temp := filepath.Dir(s.path) + ".del-" + filepath.Base(s.path)

pkg/store/container/status_test.go

@@ -132,7 +132,7 @@ func TestStatus(t *testing.T) {
 	require.NoError(err)
 	assert.Equal(testStatus, loaded)
 
-	t.Logf("successful update should take effect")
+	t.Logf("successful update should take effect but not checkpoint")
 	err = s.Update(func(o Status) (Status, error) {
 		o = updateStatus
 		return o, nil
@@ -141,6 +141,33 @@ func TestStatus(t *testing.T) {
 	assert.Equal(updateStatus, s.Get())
 	loaded, err = LoadStatus(tempDir, testID)
 	require.NoError(err)
+	assert.Equal(testStatus, loaded)
+	// Recover status.
+	assert.NoError(s.Update(func(o Status) (Status, error) {
+		o = testStatus
+		return o, nil
+	}))
+
+	t.Logf("failed update sync should not take effect")
+	err = s.UpdateSync(func(o Status) (Status, error) {
+		o = updateStatus
+		return o, updateErr
+	})
+	assert.Equal(updateErr, err)
+	assert.Equal(testStatus, s.Get())
+	loaded, err = LoadStatus(tempDir, testID)
+	require.NoError(err)
+	assert.Equal(testStatus, loaded)
+
+	t.Logf("successful update sync should take effect and checkpoint")
+	err = s.UpdateSync(func(o Status) (Status, error) {
+		o = updateStatus
+		return o, nil
+	})
+	assert.NoError(err)
+	assert.Equal(updateStatus, s.Get())
+	loaded, err = LoadStatus(tempDir, testID)
+	require.NoError(err)
 	assert.Equal(updateStatus, loaded)
 
 	t.Logf("successful update should not affect existing snapshot")