container/controller: avoid cancellation with forked pull context

Context cancellations were previously causing `Prepare` to fail
completely on re-entrant calls. To prevent this, we filtered out cancels
and deadline errors. While this allowed the service to proceed without
errors, it had the possibility of interrupting long pulls, causing the
pull to happen twice.

This PR forks the context of the pull to match the lifetime of
`Controller`, ensuring that for each task, the pull is only performed
once. It also ensures that multiple calls to `Prepare` are re-entrant,
ensuring that the pull resumes from its original position.

Signed-off-by: Stephen J Day <stephen.day@docker.com>

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
+
+			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
+			}()
+		}
 
-			// 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")
+		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)
 	}