Use resolved digest for image pulls

agent/dockerclient/dockerapi/docker_client.go

@@ -92,6 +92,10 @@ const (
 	pullRetryDelayMultiplier  = 2
 	pullRetryJitterMultiplier = 0.2
 
+	// retry settings for tagging images
+	tagImageRetryAttempts = 5
+	tagImageRetryInterval = 100 * time.Millisecond
+
 	// pollStatsTimeout is the timeout for polling Docker Stats API;
 	// keeping it same as streaming stats inactivity timeout
 	pollStatsTimeout = 18 * time.Second
@@ -129,6 +133,9 @@ type DockerClient interface {
 	// PullImage pulls an image. authData should contain authentication data provided by the ECS backend.
 	PullImage(context.Context, string, *apicontainer.RegistryAuthenticationData, time.Duration) DockerContainerMetadata
 
+	// TagImage tags a local image.
+	TagImage(ctx context.Context, source string, target string) error
+
 	// CreateContainer creates a container with the provided Config, HostConfig, and name. A timeout value
 	// and a context should be provided for the request.
 	CreateContainer(context.Context, *dockercontainer.Config, *dockercontainer.HostConfig, string, time.Duration) DockerContainerMetadata
@@ -246,6 +253,7 @@ type dockerGoClient struct {
 	context                  context.Context
 	manifestPullBackoff      retry.Backoff
 	imagePullBackoff         retry.Backoff
+	imageTagBackoff          retry.Backoff
 	inactivityTimeoutHandler inactivityTimeoutHandlerFunc
 
 	_time     ttime.Time
@@ -276,6 +284,7 @@ func (dg *dockerGoClient) WithVersion(version dockerclient.DockerVersion) (Docke
 		config:              dg.config,
 		context:             dg.context,
 		manifestPullBackoff: dg.manifestPullBackoff,
+		imageTagBackoff:     dg.imageTagBackoff,
 	}
 	// Check if the version is supported
 	_, err := versionedClient.sdkDockerClient()
@@ -319,6 +328,7 @@ func NewDockerGoClient(sdkclientFactory sdkclientfactory.Factory,
 			pullRetryJitterMultiplier, pullRetryDelayMultiplier),
 		manifestPullBackoff: retry.NewExponentialBackoff(minimumPullRetryDelay, maximumPullRetryDelay,
 			pullRetryJitterMultiplier, pullRetryDelayMultiplier),
+		imageTagBackoff:          retry.NewConstantBackoff(tagImageRetryInterval),
 		inactivityTimeoutHandler: handleInactivityTimeout,
 	}, nil
 }
@@ -593,6 +603,34 @@ func getRepository(image string) string {
 	return repository
 }
 
+// TagImage tags a local image.
+func (dg *dockerGoClient) TagImage(ctx context.Context, source string, target string) error {
+	client, err := dg.sdkDockerClient()
+	if err != nil {
+		return CannotGetDockerClientError{version: dg.version, err: err}
+	}
+
+	err = retry.RetryNWithBackoffCtx(ctx, dg.imageTagBackoff, tagImageRetryAttempts, func() error {
+		if tagErr := client.ImageTag(ctx, source, target); tagErr != nil {
+			logger.Error("Attempt to tag image failed", logger.Fields{
+				"source":    source,
+				"target":    target,
+				field.Error: tagErr,
+			})
+			return tagErr
+		}
+		return nil
+	})
+
+	if err != nil {
+		return fmt.Errorf("failed to tag image '%s' as '%s': %w", source, target, err)
+	}
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+	return nil
+}
+
 func (dg *dockerGoClient) InspectImage(image string) (*types.ImageInspect, error) {
 	defer metrics.MetricsEngineGlobal.RecordDockerMetric("INSPECT_IMAGE")()
 	client, err := dg.sdkDockerClient()

agent/dockerclient/dockerapi/docker_client_test.go

@@ -2311,3 +2311,111 @@ func TestListPluginsWithFilter(t *testing.T) {
 	assert.Equal(t, 1, len(pluginNames))
 	assert.Equal(t, "name2", pluginNames[0])
 }
+
+func TestTagImage(t *testing.T) {
+	someError := errors.New("some error")
+	tcs := []struct {
+		name                      string
+		source                    string
+		target                    string
+		setSDKFactoryExpectations func(f *mock_sdkclientfactory.MockFactory, ctrl *gomock.Controller)
+		ctx                       context.Context
+		expectedError             string
+		expectedSleeps            int
+	}{
+		{
+			name: "failed to get sdkclient",
+			setSDKFactoryExpectations: func(f *mock_sdkclientfactory.MockFactory, ctrl *gomock.Controller) {
+				f.EXPECT().GetDefaultClient().Return(nil, someError)
+			},
+			expectedError:  someError.Error(),
+			expectedSleeps: 0,
+		},
+		{
+			name:   "all attempts exhausted",
+			source: "source",
+			target: "target",
+			setSDKFactoryExpectations: func(f *mock_sdkclientfactory.MockFactory, ctrl *gomock.Controller) {
+				client := mock_sdkclient.NewMockClient(ctrl)
+				client.EXPECT().
+					ImageTag(gomock.Any(), "source", "target").
+					Times(tagImageRetryAttempts).
+					Return(someError)
+				f.EXPECT().GetDefaultClient().Return(client, nil)
+			},
+			ctx:            context.Background(),
+			expectedError:  "failed to tag image 'source' as 'target': " + someError.Error(),
+			expectedSleeps: tagImageRetryAttempts - 1,
+		},
+		{
+			name:   "second attempt worked",
+			source: "source",
+			target: "target",
+			setSDKFactoryExpectations: func(f *mock_sdkclientfactory.MockFactory, ctrl *gomock.Controller) {
+				client := mock_sdkclient.NewMockClient(ctrl)
+				client.EXPECT().ImageTag(gomock.Any(), "source", "target").Return(someError)
+				client.EXPECT().ImageTag(gomock.Any(), "source", "target").Return(nil)
+				f.EXPECT().GetDefaultClient().Return(client, nil)
+			},
+			ctx:            context.Background(),
+			expectedSleeps: 1,
+		},
+		{
+			name: "canceled context",
+			setSDKFactoryExpectations: func(f *mock_sdkclientfactory.MockFactory, ctrl *gomock.Controller) {
+				client := mock_sdkclient.NewMockClient(ctrl)
+				f.EXPECT().GetDefaultClient().Return(client, nil)
+			},
+			ctx: func() context.Context {
+				c, cancel := context.WithCancel(context.Background())
+				cancel()
+				return c
+			}(),
+			expectedError: "context canceled",
+		},
+		{
+			name: "deadline exceeded",
+			setSDKFactoryExpectations: func(f *mock_sdkclientfactory.MockFactory, ctrl *gomock.Controller) {
+				client := mock_sdkclient.NewMockClient(ctrl)
+				f.EXPECT().GetDefaultClient().Return(client, nil)
+			},
+			ctx: func() context.Context {
+				c, cancel := context.WithTimeout(context.Background(), 0)
+				<-c.Done() // wait for deadline to be exceeded
+				cancel()
+				return c
+			}(),
+			expectedError: "context deadline exceeded",
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+
+			// Set up mocks
+			mockDockerSDK := mock_sdkclient.NewMockClient(ctrl)
+			mockDockerSDK.EXPECT().Ping(gomock.Any()).Return(types.Ping{}, nil)
+			sdkFactory := mock_sdkclientfactory.NewMockFactory(ctrl)
+			sdkFactory.EXPECT().GetDefaultClient().Return(mockDockerSDK, nil)
+
+			// Set up docker client for testing
+			client, err := NewDockerGoClient(sdkFactory, defaultTestConfig(), context.Background())
+			require.NoError(t, err)
+			// Make retries fast
+			client.(*dockerGoClient).imageTagBackoff = retry.NewConstantBackoff(0)
+
+			if tc.setSDKFactoryExpectations != nil {
+				tc.setSDKFactoryExpectations(sdkFactory, ctrl)
+			}
+
+			err = client.TagImage(tc.ctx, tc.source, tc.target)
+			if tc.expectedError == "" {
+				assert.NoError(t, err)
+			} else {
+				assert.EqualError(t, err, tc.expectedError)
+			}
+		})
+	}
+}

agent/dockerclient/sdkclient/interface.go

@@ -55,6 +55,7 @@ type Client interface {
 	ImagePull(ctx context.Context, refStr string, options types.ImagePullOptions) (io.ReadCloser, error)
 	ImageRemove(ctx context.Context, imageID string, options types.ImageRemoveOptions) ([]types.ImageDeleteResponseItem,
 		error)
+	ImageTag(ctx context.Context, source, target string) error
 	Ping(ctx context.Context) (types.Ping, error)
 	PluginList(ctx context.Context, filter filters.Args) (types.PluginsListResponse, error)
 	VolumeCreate(ctx context.Context, options volume.CreateOptions) (volume.Volume, error)

agent/engine/common_test.go

@@ -38,6 +38,7 @@ import (
 	mock_engine "github.com/aws/amazon-ecs-agent/agent/engine/mocks"
 	"github.com/aws/amazon-ecs-agent/agent/statechange"
 	"github.com/aws/amazon-ecs-agent/agent/utils"
+	referenceutil "github.com/aws/amazon-ecs-agent/agent/utils/reference"
 	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"
@@ -48,8 +49,11 @@ import (
 	dockercontainer "github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/registry"
 	"github.com/golang/mock/gomock"
+	"github.com/opencontainers/go-digest"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pborman/uuid"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 const (
@@ -158,15 +162,29 @@ func validateContainerRunWorkflow(t *testing.T,
 	createdContainerName chan<- string,
 	assertions func(),
 ) {
-	// Prepare mock image manifest digest for test
+	// Prepare a test digest
+	testDigest, digestParseError := digest.Parse(
+		"sha256:c5b1261d6d3e43071626931fc004f70149baeba2c8ec672bd4f27761f8e1ad6b")
+	require.NoError(t, digestParseError)
+
+	// Get expected canonical reference for the container image and test digest
+	canonicalImageRef, canonicalRefErr := referenceutil.GetCanonicalRef(
+		container.Image, testDigest.String())
+	require.NoError(t, canonicalRefErr)
+
+	// Set expectations for transition to MANIFEST_PULLED state
 	manifestPullClient := mock_dockerapi.NewMockDockerClient(ctrl)
 	client.EXPECT().WithVersion(dockerclient.Version_1_35).Return(manifestPullClient, nil)
 	manifestPullClient.EXPECT().
 		PullImageManifest(gomock.Any(), container.Image, container.RegistryAuthentication).
-		Return(registry.DistributionInspect{}, nil)
+		Return(registry.DistributionInspect{Descriptor: ocispec.Descriptor{Digest: testDigest}}, nil)
 
 	imageManager.EXPECT().AddAllImageStates(gomock.Any()).AnyTimes()
-	client.EXPECT().PullImage(gomock.Any(), container.Image, nil, gomock.Any()).Return(dockerapi.DockerContainerMetadata{})
+	client.EXPECT().
+		PullImage(gomock.Any(), canonicalImageRef.String(), nil, gomock.Any()).
+		Return(dockerapi.DockerContainerMetadata{})
+	client.EXPECT().
+		TagImage(gomock.Any(), canonicalImageRef.String(), container.Image).Return(nil)
 	imageManager.EXPECT().RecordContainerReference(container).Return(nil)
 	imageManager.EXPECT().GetImageStateFromImageName(gomock.Any()).Return(nil, false)
 	client.EXPECT().APIVersion().Return(defaultDockerClientAPIVersion, nil)