From facbaa0e798fb8048de7a84bb896da982b120995 Mon Sep 17 00:00:00 2001
From: Lantao Liu <lantaol@google.com>
Date: Thu, 5 Dec 2019 14:51:02 -0800
Subject: [PATCH] Better handle unknown state.

Signed-off-by: Lantao Liu <lantaol@google.com>
---
 pkg/server/events.go               | 6 ++++++
 pkg/server/helpers.go              | 1 +
 pkg/server/restart.go              | 4 +++-
 pkg/store/container/status.go      | 6 ++++++
 pkg/store/container/status_test.go | 8 ++++++++
 5 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/pkg/server/events.go b/pkg/server/events.go
index e4536a559..fcbf0a3ca 100644
--- a/pkg/server/events.go
+++ b/pkg/server/events.go
@@ -333,6 +333,12 @@ func handleContainerExit(ctx context.Context, e *eventtypes.TaskExit, cntr conta
 		status.Pid = 0
 		status.FinishedAt = e.ExitedAt.UnixNano()
 		status.ExitCode = int32(e.ExitStatus)
+		// Unknown state can only transit to EXITED state, so we need
+		// to handle unknown state here.
+		if status.Unknown {
+			logrus.Debugf("Container %q transited from UNKNOWN to EXITED", cntr.ID)
+			status.Unknown = false
+		}
 		return status, nil
 	})
 	if err != nil {
diff --git a/pkg/server/helpers.go b/pkg/server/helpers.go
index 98fe09d49..163629f8a 100644
--- a/pkg/server/helpers.go
+++ b/pkg/server/helpers.go
@@ -359,6 +359,7 @@ func unknownContainerStatus() containerstore.Status {
 		FinishedAt: 0,
 		ExitCode:   unknownExitCode,
 		Reason:     unknownExitReason,
+		Unknown:    true,
 	}
 }
 
diff --git a/pkg/server/restart.go b/pkg/server/restart.go
index 80578dd02..d1058ca48 100644
--- a/pkg/server/restart.go
+++ b/pkg/server/restart.go
@@ -308,7 +308,9 @@ func (c *criService) loadContainer(ctx context.Context, cntr containerd.Containe
 	}()
 	if err != nil {
 		log.G(ctx).WithError(err).Errorf("Failed to load container status for %q", id)
-		status = unknownContainerStatus()
+		// Only set the unknown field in this case, because other fields may
+		// contain useful information loaded from the checkpoint.
+		status.Unknown = true
 	}
 	opts := []containerstore.Opts{
 		containerstore.WithStatus(status, containerDir),
diff --git a/pkg/store/container/status.go b/pkg/store/container/status.go
index c46325bbe..9aaef6970 100644
--- a/pkg/store/container/status.go
+++ b/pkg/store/container/status.go
@@ -94,10 +94,16 @@ type Status struct {
 	// Removing indicates that the container is in removing state.
 	// This field doesn't need to be checkpointed.
 	Removing bool `json:"-"`
+	// Unknown indicates that the container status is not fully loaded.
+	// This field doesn't need to be checkpointed.
+	Unknown bool `json:"-"`
 }
 
 // State returns current state of the container based on the container status.
 func (s Status) State() runtime.ContainerState {
+	if s.Unknown {
+		return runtime.ContainerState_CONTAINER_UNKNOWN
+	}
 	if s.FinishedAt != 0 {
 		return runtime.ContainerState_CONTAINER_EXITED
 	}
diff --git a/pkg/store/container/status_test.go b/pkg/store/container/status_test.go
index 00fad9b62..16db8b0d6 100644
--- a/pkg/store/container/status_test.go
+++ b/pkg/store/container/status_test.go
@@ -36,6 +36,12 @@ func TestContainerState(t *testing.T) {
 		state  runtime.ContainerState
 	}{
 		"unknown state": {
+			status: Status{
+				Unknown: true,
+			},
+			state: runtime.ContainerState_CONTAINER_UNKNOWN,
+		},
+		"unknown state because there is no timestamp set": {
 			status: Status{},
 			state:  runtime.ContainerState_CONTAINER_UNKNOWN,
 		},
@@ -76,6 +82,7 @@ func TestStatusEncodeDecode(t *testing.T) {
 		Message:    "test-message",
 		Removing:   true,
 		Starting:   true,
+		Unknown:    true,
 	}
 	assert := assertlib.New(t)
 	data, err := s.encode()
@@ -84,6 +91,7 @@ func TestStatusEncodeDecode(t *testing.T) {
 	assert.NoError(newS.decode(data))
 	s.Removing = false // Removing should not be encoded.
 	s.Starting = false // Starting should not be encoded.
+	s.Unknown = false  // Unknown should not be encoded.
 	assert.Equal(s, newS)
 
 	unsupported, err := json.Marshal(&versionedStatus{