wip

aws · Apr 23, 2024 · 3bb91b8 · 3bb91b8
1 parent 251d73b
commit 3bb91b8
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
diff --git a/agent/dockerclient/dockerapi/docker_client.go b/agent/dockerclient/dockerapi/docker_client.go
@@ -373,6 +373,10 @@ func (dg *dockerGoClient) PullImageManifest(
 	})
 
 	if err != nil {
+		if errors.Is(err, context.DeadlineExceeded) {
+			timeoutErr := &DockerTimeoutError{time.Since(startTime), "MANIFEST_PULLED"}
+			return registry.DistributionInspect{}, timeoutErr
+		}
 		return registry.DistributionInspect{}, wrapManifestPullErrorAsNamedError(imageRef, err)
 	}
 	if ctxErr := ctx.Err(); ctxErr != nil {

diff --git a/agent/engine/engine_integ_test.go b/agent/engine/engine_integ_test.go
@@ -645,7 +645,9 @@ func TestPullContainerManifestInteg(t *testing.T) {
 	tcs := []struct {
 		name           string
 		image          string
+		setConfig      func(c *config.Config)
 		expectedDigest string
+		assertError    func(t *testing.T, err error)
 	}{
 		{
 			name:           "digest available in image reference",
@@ -662,6 +664,22 @@ func TestPullContainerManifestInteg(t *testing.T) {
 			image:          "127.0.0.1:51670/busybox",
 			expectedDigest: "sha256:51de9138b0cc394c813df84f334d638499333cac22edd05d0300b2c9a2dc80dd",
 		},
+		{
+			name:      "manifest pull can timeout",
+			image:     "127.0.0.1:51670/busybox",
+			setConfig: func(c *config.Config) { c.ManifestPullTimeout = 0 },
+			assertError: func(t *testing.T, err error) {
+				assert.ErrorContains(t, err, "Could not transition to MANIFEST_PULLED; timed out")
+			},
+		},
+		{
+			name:      "manifest pull can timeout - non-zero timeout",
+			image:     "127.0.0.1:51670/busybox",
+			setConfig: func(c *config.Config) { c.ManifestPullTimeout = 100 * time.Microsecond },
+			assertError: func(t *testing.T, err error) {
+				assert.ErrorContains(t, err, "Could not transition to MANIFEST_PULLED; timed out")
+			},
+		},
 	}
 	imagePullBehaviors := []config.ImagePullBehaviorType{
 		config.ImagePullDefaultBehavior, config.ImagePullAlwaysBehavior,
@@ -672,7 +690,10 @@ func TestPullContainerManifestInteg(t *testing.T) {
 			t.Run(fmt.Sprintf("%s - %d", tc.name, imagePullBehavior), func(t *testing.T) {
 				cfg := defaultTestConfigIntegTest()
 				cfg.ImagePullBehavior = imagePullBehavior
-				cfg.ManifestPullTimeout = 0
+
+				if tc.setConfig != nil {
+					tc.setConfig(cfg)
+				}
 
 				taskEngine, done, _ := setup(cfg, nil, t)
 				defer done()
@@ -681,8 +702,12 @@ func TestPullContainerManifestInteg(t *testing.T) {
 				task := &apitask.Task{Containers: []*apicontainer.Container{container}}
 
 				res := taskEngine.(*DockerTaskEngine).pullContainerManifest(task, container)
-				require.NoError(t, res.Error)
-				assert.Equal(t, tc.expectedDigest, container.GetImageDigest())
+				if tc.assertError == nil {
+					require.NoError(t, res.Error)
+					assert.Equal(t, tc.expectedDigest, container.GetImageDigest())
+				} else {
+					tc.assertError(t, res.Error)
+				}
 			})
 		}
 	}