Merge pull request #25043 from stevvooe/forked-pull-context

container/controller: avoid cancellation with forked pull context

daemon/cluster/executor/container/controller.go

@@ -24,6 +24,10 @@ type controller struct {
 	adapter *containerAdapter
 	closed  chan struct{}
 	err     error
+
+	pulled     chan struct{} // closed after pull
+	cancelPull func()        // cancels pull context if not nil
+	pullErr    error         // pull error, only read after pulled closed
 }
 
 var _ exec.Controller = &controller{}
@@ -86,24 +90,40 @@ func (r *controller) Prepare(ctx context.Context) error {
 	}
 
 	if os.Getenv("DOCKER_SERVICE_PREFER_OFFLINE_IMAGE") != "1" {
-		if err := r.adapter.pullImage(ctx); err != nil {
-			cause := errors.Cause(err)
-			if cause == context.Canceled || cause == context.DeadlineExceeded {
-				return err
-			}
+		if r.pulled == nil {
+			// Fork the pull to a different context to allow pull to continue
+			// on re-entrant calls to Prepare. This ensures that Prepare can be
+			// idempotent and not incur the extra cost of pulling when
+			// cancelled on updates.
+			var pctx context.Context
 
-			// NOTE(stevvooe): We always try to pull the image to make sure we have
-			// the most up to date version. This will return an error, but we only
-			// log it. If the image truly doesn't exist, the create below will
-			// error out.
-			//
-			// This gives us some nice behavior where we use up to date versions of
-			// mutable tags, but will still run if the old image is available but a
-			// registry is down.
-			//
-			// If you don't want this behavior, lock down your image to an
-			// immutable tag or digest.
-			log.G(ctx).WithError(err).Error("pulling image failed")
+			r.pulled = make(chan struct{})
+			pctx, r.cancelPull = context.WithCancel(context.Background()) // TODO(stevvooe): Bind a context to the entire controller.
+
+			go func() {
+				defer close(r.pulled)
+				r.pullErr = r.adapter.pullImage(pctx) // protected by closing r.pulled
+			}()
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-r.pulled:
+			if r.pullErr != nil {
+				// NOTE(stevvooe): We always try to pull the image to make sure we have
+				// the most up to date version. This will return an error, but we only
+				// log it. If the image truly doesn't exist, the create below will
+				// error out.
+				//
+				// This gives us some nice behavior where we use up to date versions of
+				// mutable tags, but will still run if the old image is available but a
+				// registry is down.
+				//
+				// If you don't want this behavior, lock down your image to an
+				// immutable tag or digest.
+				log.G(ctx).WithError(r.pullErr).Error("pulling image failed")
+			}
 		}
 	}
 
@@ -252,6 +272,10 @@ func (r *controller) Shutdown(ctx context.Context) error {
 		return err
 	}
 
+	if r.cancelPull != nil {
+		r.cancelPull()
+	}
+
 	if err := r.adapter.shutdown(ctx); err != nil {
 		if isUnknownContainer(err) || isStoppedContainer(err) {
 			return nil
@@ -269,6 +293,10 @@ func (r *controller) Terminate(ctx context.Context) error {
 		return err
 	}
 
+	if r.cancelPull != nil {
+		r.cancelPull()
+	}
+
 	if err := r.adapter.terminate(ctx); err != nil {
 		if isUnknownContainer(err) {
 			return nil
@@ -286,6 +314,10 @@ func (r *controller) Remove(ctx context.Context) error {
 		return err
 	}
 
+	if r.cancelPull != nil {
+		r.cancelPull()
+	}
+
 	// It may be necessary to shut down the task before removing it.
 	if err := r.Shutdown(ctx); err != nil {
 		if isUnknownContainer(err) {
@@ -320,6 +352,10 @@ func (r *controller) Close() error {
 	case <-r.closed:
 		return r.err
 	default:
+		if r.cancelPull != nil {
+			r.cancelPull()
+		}
+
 		r.err = exec.ErrControllerClosed
 		close(r.closed)
 	}