From 51659a6f1396307ca5ed96f8c6810a8ff73a2e9c Mon Sep 17 00:00:00 2001
From: Onur Filiz <onfiliz@amazon.com>
Date: Sat, 2 Dec 2017 05:40:59 -0800
Subject: [PATCH] Persist container exit code in agent state file

Container exit code needs to be persisted in the agent state file so
that api.Container structs can be rebuilt after agent restarts. Go's
default JSON encoder marshals only public fields.

Rename and export KnownExitCodeUnsafe with proper JSON marking.
Add test to ensure exit code is persisted and catch any regressions.
Fix existing test panics caused by continuing after fatal nil checks.
---
 CHANGELOG.md                             |  1 +
 agent/api/container.go                   | 17 +++++++++++++----
 agent/statemanager/state_manager_test.go | 11 +++++++++--
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 055204e11fc..be535e0da29 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 * Bug - Fixed a bug where `-version` fails due to its dependency on docker client. [#1118](https://github.com/aws/amazon-ecs-agent/pull/1118)
+* Bug - Persist container exit code in agent state file [#1125](https://github.com/aws/amazon-ecs-agent/pull/1125)
 
 ## 1.16.0
 * Feature - Support pulling from Amazon ECR with specified IAM role in task definition
diff --git a/agent/api/container.go b/agent/api/container.go
index aea4ea62e14..77a8b5648d2 100644
--- a/agent/api/container.go
+++ b/agent/api/container.go
@@ -135,7 +135,14 @@ type Container struct {
 	// metadata file
 	MetadataFileUpdated bool `json:"metadataFileUpdated"`
 
-	knownExitCode     *int
+	// KnownExitCodeUnsafe specifies the exit code for the container.
+	// It is exposed outside of the package so that it's marshalled/unmarshalled in
+	// the JSON body while saving the state.
+	// NOTE: Do not access KnownExitCodeUnsafe directly. Instead, use `GetKnownExitCode`
+	// and `SetKnownExitCode`.
+	KnownExitCodeUnsafe *int `json:"KnownExitCode"`
+
+	// KnownPortBindings is an array of port bindings for the container.
 	KnownPortBindings []PortBinding
 
 	// SteadyStateStatusUnsafe specifies the steady state status for the container
@@ -237,14 +244,16 @@ func (c *Container) SetSentStatus(status ContainerStatus) {
 func (c *Container) SetKnownExitCode(i *int) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
-	c.knownExitCode = i
+
+	c.KnownExitCodeUnsafe = i
 }
 
 // GetKnownExitCode returns the container exit code
 func (c *Container) GetKnownExitCode() *int {
 	c.lock.RLock()
 	defer c.lock.RUnlock()
-	return c.knownExitCode
+
+	return c.KnownExitCodeUnsafe
 }
 
 // SetRegistryAuthCredentials sets the credentials for pulling image from ECR
@@ -363,7 +372,7 @@ func (c *Container) IsEssential() bool {
 	return c.Essential
 }
 
-// LogAuthExecutionRole returns true if the auth is by exectution role
+// AWSLogAuthExecutionRole returns true if the auth is by execution role
 func (c *Container) AWSLogAuthExecutionRole() bool {
 	return c.LogsAuthStrategy == awslogsAuthExecutionRole
 }
diff --git a/agent/statemanager/state_manager_test.go b/agent/statemanager/state_manager_test.go
index 4ddf5643fc2..01d41be1486 100644
--- a/agent/statemanager/state_manager_test.go
+++ b/agent/statemanager/state_manager_test.go
@@ -24,6 +24,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
 	"github.com/aws/amazon-ecs-agent/agent/statemanager"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestStateManagerNonexistantDirectory(t *testing.T) {
@@ -34,7 +35,7 @@ func TestStateManagerNonexistantDirectory(t *testing.T) {
 
 func TestLoadsV1DataCorrectly(t *testing.T) {
 	cleanup, err := setupWindowsTest(filepath.Join(".", "testdata", "v1", "1", "ecs_agent_data.json"))
-	assert.Nil(t, err, "Failed to set up test")
+	require.Nil(t, err, "Failed to set up test")
 	defer cleanup()
 	cfg := &config.Config{DataDir: filepath.Join(".", "testdata", "v1", "1")}
 
@@ -62,11 +63,17 @@ func TestLoadsV1DataCorrectly(t *testing.T) {
 			deadTask = task
 		}
 	}
-	assert.NotNil(t, deadTask)
+
+	require.NotNil(t, deadTask)
 	assert.Equal(t, deadTask.GetSentStatus(), api.TaskStopped)
 	assert.Equal(t, deadTask.Containers[0].SentStatusUnsafe, api.ContainerStopped)
 	assert.Equal(t, deadTask.Containers[0].DesiredStatusUnsafe, api.ContainerStopped)
 	assert.Equal(t, deadTask.Containers[0].KnownStatusUnsafe, api.ContainerStopped)
+
+	exitCode := deadTask.Containers[0].KnownExitCodeUnsafe
+	require.NotNil(t, exitCode)
+	assert.Equal(t, *exitCode, 128)
+
 	expected, _ := time.Parse(time.RFC3339, "2015-04-28T17:29:48.129140193Z")
 	assert.Equal(t, deadTask.GetKnownStatusTime(), expected)
 }