Integrate restart policy into the container object

agent/api/container/container.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	resourcestatus "github.com/aws/amazon-ecs-agent/agent/taskresource/status"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/container/restart"
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
 	apierrors "github.com/aws/amazon-ecs-agent/ecs-agent/api/errors"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
@@ -334,6 +335,13 @@ type Container struct {
 	ContainerPortSet map[int]struct{}
 	// ContainerPortRangeMap is a map of containerPortRange to its associated hostPortRange
 	ContainerPortRangeMap map[string]string
+
+	// RestartPolicy is an object representing the restart policy of the container
+	RestartPolicy *restart.RestartPolicy `json:"restartPolicy,omitempty"`
+	// RestartTracker tracks this container's restart policy metadata, such
+	// as restart count and last restart time. This is only initialized if the container
+	// has a restart policy defined and enabled.
+	RestartTracker *restart.RestartTracker `json:"restartTracker,omitempty"`
 }
 
 type DependsOn struct {
@@ -627,6 +635,16 @@ func (c *Container) IsEssential() bool {
 	return c.Essential
 }
 
+// RestartPolicyEnabled returns whether the restart policy is defined and enabled
+func (c *Container) RestartPolicyEnabled() bool {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+	if c.RestartPolicy == nil {
+		return false
+	}
+	return c.RestartPolicy.Enabled
+}
+
 // AWSLogAuthExecutionRole returns true if the auth is by execution role
 func (c *Container) AWSLogAuthExecutionRole() bool {
 	return c.LogsAuthStrategy == awslogsAuthExecutionRole

agent/api/container/container_test.go

@@ -23,11 +23,13 @@ import (
 	"time"
 
 	resourcestatus "github.com/aws/amazon-ecs-agent/agent/taskresource/status"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/container/restart"
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
 
 	"github.com/aws/amazon-ecs-agent/agent/utils"
 	dockercontainer "github.com/docker/docker/api/types/container"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type configPair struct {
@@ -1334,6 +1336,73 @@ func TestGetCredentialSpec(t *testing.T) {
 	}
 }
 
+func TestRestartPolicy(t *testing.T) {
+	testCases := []struct {
+		name            string
+		container       *Container
+		restartCount    int
+		expectedEnabled bool
+	}{
+		{
+			name: "nil restart policy",
+			container: &Container{
+				RestartPolicy: nil,
+			},
+			restartCount:    0,
+			expectedEnabled: false,
+		},
+		{
+			name: "not enabled restart policy",
+			container: &Container{
+				RestartPolicy: &restart.RestartPolicy{},
+			},
+			restartCount:    0,
+			expectedEnabled: false,
+		},
+		{
+			name: "explicitly not enabled restart policy",
+			container: &Container{
+				RestartPolicy: &restart.RestartPolicy{
+					Enabled: false,
+				},
+			},
+			restartCount:    0,
+			expectedEnabled: false,
+		},
+		{
+			name: "enabled restart policy",
+			container: &Container{
+				RestartPolicy: &restart.RestartPolicy{
+					Enabled: true,
+				},
+			},
+			restartCount:    0,
+			expectedEnabled: true,
+		},
+		{
+			name: "enabled restart policy, record 5 restarts",
+			container: &Container{
+				RestartPolicy: &restart.RestartPolicy{
+					Enabled: true,
+				},
+			},
+			restartCount:    5,
+			expectedEnabled: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		require.Equal(t, tc.expectedEnabled, tc.container.RestartPolicyEnabled())
+		if tc.container.RestartPolicyEnabled() {
+			tc.container.RestartTracker = restart.NewRestartTracker(*tc.container.RestartPolicy)
+			for i := 0; i < tc.restartCount; i++ {
+				tc.container.RestartTracker.RecordRestart()
+			}
+			require.Equal(t, tc.restartCount, tc.container.RestartTracker.GetRestartCount())
+		}
+	}
+}
+
 func getContainer(hostConfig string, credentialSpecs []string) *Container {
 	c := &Container{
 		Name: "c",

agent/api/task/task.go

@@ -41,6 +41,7 @@ import (
 	taskresourcevolume "github.com/aws/amazon-ecs-agent/agent/taskresource/volume"
 	"github.com/aws/amazon-ecs-agent/agent/utils"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api/container/restart"
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/model/ecs"
 	apierrors "github.com/aws/amazon-ecs-agent/ecs-agent/api/errors"
@@ -466,6 +467,8 @@ func (task *Task) PostUnmarshalTask(cfg *config.Config,
 		}
 	}
 
+	task.initRestartTrackers()
+
 	for _, opt := range options {
 		if err := opt(task); err != nil {
 			logger.Error("Could not apply task option", logger.Fields{
@@ -523,6 +526,16 @@ func (task *Task) initNetworkMode(acsTaskNetworkMode *string) {
 	}
 }
 
+// initRestartTrackers initializes the restart policy tracker for each container
+// that has a restart policy configured and enabled.
+func (task *Task) initRestartTrackers() {
+	for _, c := range task.Containers {
+		if c.RestartPolicyEnabled() {
+			c.RestartTracker = restart.NewRestartTracker(*c.RestartPolicy)
+		}
+	}
+}
+
 func (task *Task) initServiceConnectResources() error {
 	// TODO [SC]: ServiceConnectConfig will come from ACS. Adding this here for dev/testing purposes only Remove when
 	// ACS model is integrated

agent/vendor/modules.txt

@@ -16,6 +16,7 @@ github.com/aws/amazon-ecs-agent/ecs-agent/api/appnet/mocks
 github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment
 github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment/resource
 github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment/resource/mocks
+github.com/aws/amazon-ecs-agent/ecs-agent/api/container/restart
 github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status
 github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs
 github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/client