Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Integrate restart policy into the container object #4160

Merged
merged 1 commit into from
May 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions agent/api/container/container.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
69 changes: 69 additions & 0 deletions agent/api/container/container_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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{
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just curious, do we/are there plans to validate RestartPolicy isn't empty if restart policy is enabled when we receive the task payload?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a good question, we may want to validate that when we initialize the restart trackers, let me sync with Fargate dataplane on how we want to do that

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍 thanks Cam! I also have no issues for this to be a follow up

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",
Expand Down
13 changes: 13 additions & 0 deletions agent/api/task/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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
Expand Down
38 changes: 38 additions & 0 deletions agent/api/task/task_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ import (
"github.com/aws/amazon-ecs-agent/agent/utils"
"github.com/aws/amazon-ecs-agent/ecs-agent/acs/model/ecsacs"
apiresource "github.com/aws/amazon-ecs-agent/ecs-agent/api/attachment/resource"
"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"
apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
Expand Down Expand Up @@ -3908,6 +3909,43 @@ func TestPostUnmarshalTaskEnvfiles(t *testing.T) {
task.Containers[0].TransitionDependenciesMap[apicontainerstatus.ContainerCreated].ResourceDependencies[0])
}

func TestPostUnmarshalTaskContainerRestartPolicy(t *testing.T) {
container := &apicontainer.Container{
Name: "containerName",
Image: "image:tag",
RestartPolicy: &restart.RestartPolicy{
Enabled: true,
},
TransitionDependenciesMap: make(map[apicontainerstatus.ContainerStatus]apicontainer.TransitionDependencySet),
}

task := &Task{
Arn: testTaskARN,
ResourcesMapUnsafe: make(map[string][]taskresource.TaskResource),
Containers: []*apicontainer.Container{container},
}

ctrl := gomock.NewController(t)
defer ctrl.Finish()

cfg := &config.Config{}
credentialsManager := mock_credentials.NewMockManager(ctrl)
resFields := &taskresource.ResourceFields{
ResourceFieldsCommon: &taskresource.ResourceFieldsCommon{
CredentialsManager: credentialsManager,
},
}

err := task.PostUnmarshalTask(cfg, credentialsManager, resFields, nil, nil)
assert.NoError(t, err)

for _, c := range task.Containers {
assert.True(t, c.RestartPolicyEnabled())
assert.NotNil(t, c.RestartTracker)
assert.Equal(t, 0, c.RestartTracker.GetRestartCount())
}
}

func TestInitializeAndGetEnvfilesResource(t *testing.T) {
envfile1 := apicontainer.EnvironmentFile{
Value: "s3://bucket/envfile1",
Expand Down

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions agent/vendor/modules.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading