From a0a2e55a6762a4f5f39de983bd125801f19cd730 Mon Sep 17 00:00:00 2001
From: Rob Watson <rob@netflux.io>
Date: Sat, 12 Apr 2025 19:41:31 +0200
Subject: [PATCH] fixup! refactor(container): restart handling

---
 internal/container/container.go | 41 ++++++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/internal/container/container.go b/internal/container/container.go
index 30555f2..774e7ae 100644
--- a/internal/container/container.go
+++ b/internal/container/container.go
@@ -21,7 +21,6 @@ import (
 	"github.com/docker/docker/api/types/filters"
 	"github.com/docker/docker/api/types/image"
 	"github.com/docker/docker/api/types/network"
-	"github.com/docker/docker/errdefs"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
@@ -72,6 +71,7 @@ type Client struct {
 	wg           sync.WaitGroup
 	apiClient    DockerClient
 	networkID    string
+	cancelFuncs  map[string]context.CancelFunc
 	pulledImages map[string]struct{}
 	logger       *slog.Logger
 }
@@ -99,6 +99,7 @@ func NewClient(ctx context.Context, apiClient DockerClient, logger *slog.Logger)
 		cancel:       cancel,
 		apiClient:    apiClient,
 		networkID:    network.ID,
+		cancelFuncs:  make(map[string]context.CancelFunc),
 		pulledImages: make(map[string]struct{}),
 		logger:       logger,
 	}
@@ -235,6 +236,12 @@ func (a *Client) RunContainer(ctx context.Context, params RunContainerParams) (<
 
 		containerStateC <- domain.Container{ID: createResp.ID, Status: domain.ContainerStatusRunning}
 
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithCancel(ctx)
+		a.mu.Lock()
+		a.cancelFuncs[createResp.ID] = cancel
+		a.mu.Unlock()
+
 		a.runContainerLoop(
 			ctx,
 			createResp.ID,
@@ -385,21 +392,11 @@ func (a *Client) runContainerLoop(
 					containerErrC <- fmt.Errorf("container start: %w", err)
 				}
 			case err := <-errC:
-				// If this is a not found error, the container has been removed -
-				// probably by the user. This is a bit hacky, and should be more
-				// explicit, possibly by signalling to this package that the container
-				// has been removed by the user instead of just calling
-				// ContainerRemove.
-				// TODO: improve this
-				if errdefs.IsNotFound(err) {
-					a.logger.Debug("Container not found when setting ContainerWait, ignoring", "id", shortID(containerID))
-					containerRespC <- containerWaitResponse{WaitResponse: container.WaitResponse{}, restarting: false}
-				} else {
-					containerErrC <- err
-				}
+				containerErrC <- err
 				return
 			case <-ctx.Done():
-				containerErrC <- ctx.Err()
+				// This is probably because the container was stopped.
+				containerRespC <- containerWaitResponse{WaitResponse: container.WaitResponse{}, restarting: false}
 				return
 			}
 		}
@@ -519,6 +516,22 @@ func (a *Client) Close() error {
 }
 
 func (a *Client) removeContainer(ctx context.Context, id string) error {
+	a.mu.Lock()
+	cancel, ok := a.cancelFuncs[id]
+	if ok {
+		delete(a.cancelFuncs, id)
+	}
+	a.mu.Unlock()
+
+	if ok {
+		cancel()
+	} else {
+		// This shouldn't happen, but we prefer the Docker engine to be the
+		// definitive source of truth and not the cancelFuncs map so don't sweat it
+		// if it doesn't exist.
+		a.logger.Warn("removeContainer: cancelFunc not found", "id", shortID(id))
+	}
+
 	a.logger.Info("Stopping container", "id", shortID(id))
 	stopTimeout := int(stopTimeout.Seconds())
 	if err := a.apiClient.ContainerStop(ctx, id, container.StopOptions{Timeout: &stopTimeout}); err != nil {