Add join handle for StartDockerImageDownload to bound goroutine lifecycle

pkg/cli/docker_images.go

@@ -43,13 +43,15 @@ const (
 type dockerPullState struct {
 	mu                  sync.RWMutex
 	downloading         map[string]bool // image -> is currently downloading
+	inflight            map[string]chan struct{}
 	mockAvailable       map[string]bool // for testing: override IsDockerImageAvailable
 	mockAvailableInUse  bool            // for testing: whether to use mockAvailable
 	mockDockerAvailable bool            // for testing: override IsDockerAvailable (default true)
 }
 
 var pullState = &dockerPullState{
 	downloading:         make(map[string]bool),
+	inflight:            make(map[string]chan struct{}),
 	mockAvailable:       make(map[string]bool),
 	mockDockerAvailable: true,
 }
@@ -130,34 +132,44 @@ func IsDockerAvailable(ctx context.Context) bool {
 // StartDockerImageDownload starts downloading a Docker image in the background
 // Returns true if download was started, false if already downloading or available
 // The download can be cancelled by cancelling the provided context
-func StartDockerImageDownload(ctx context.Context, image string) bool {
+// The returned join function blocks until the corresponding download goroutine exits.
+func StartDockerImageDownload(ctx context.Context, image string) (bool, func()) {
 	ctx = normalizeDockerContext(ctx)
 
 	// Check availability and downloading status atomically under lock
 	pullState.mu.Lock()
 	defer pullState.mu.Unlock()
 
-	// Check if already available (inside lock for atomicity)
-	if isDockerImageAvailableUnlocked(ctx, image) {
-		dockerImagesLog.Printf("Image %s is already available", image)
-		return false
-	}
-
 	// Check if already downloading
 	if pullState.downloading[image] {
 		dockerImagesLog.Printf("Image %s is already downloading", image)
-		return false
+		done := pullState.inflight[image]
+		return false, func() {
+			if done != nil {
+				<-done
+			}
+		}
 	}
 
+	// Check if already available (inside lock for atomicity)
+	if isDockerImageAvailableUnlocked(ctx, image) {
+		dockerImagesLog.Printf("Image %s is already available", image)
+		return false, func() {}
+	}
+
+	done := make(chan struct{})
 	pullState.downloading[image] = true
+	pullState.inflight[image] = done
 
 	// Start the download in a goroutine with retry logic
 	go func() {
+		defer close(done)
 		defer func() {
 			func() {
 				pullState.mu.Lock()
 				defer pullState.mu.Unlock()
 				delete(pullState.downloading, image)
+				delete(pullState.inflight, image)
 			}()
 			if r := recover(); r != nil {
 				dockerImagesLog.Printf("Panic in docker image download for %s (recovered): %v", image, r)
@@ -218,7 +230,9 @@ func StartDockerImageDownload(ctx context.Context, image string) bool {
 		dockerImagesLog.Printf("Failed to download image %s after %d attempts: %v\nOutput: %s", image, maxAttempts, lastErr, string(lastOutput))
 	}()
 
-	return true
+	return true, func() {
+		<-done
+	}
 }
 
 // CheckAndPrepareDockerImages checks if required Docker images are available
@@ -319,7 +333,7 @@ func CheckAndPrepareDockerImages(ctx context.Context, useZizmor, usePoutine, use
 			downloadingImages = append(downloadingImages, img.name)
 		} else {
 			// Start download
-			StartDockerImageDownload(ctx, img.image)
+			_, _ = StartDockerImageDownload(ctx, img.image)
 			missingImages = append(missingImages, img.name)
 		}
 	}

pkg/cli/docker_images_test.go

@@ -351,7 +351,7 @@ func TestStartDockerImageDownload_ConcurrentCalls(t *testing.T) {
 	for i := range numGoroutines {
 		go func(index int) {
 			<-startChan // Wait for the signal to start
-			started[index] = StartDockerImageDownload(context.Background(), testImage)
+			started[index], _ = StartDockerImageDownload(context.Background(), testImage)
 			doneChan <- index
 		}(i)
 	}
@@ -407,7 +407,7 @@ func TestStartDockerImageDownload_ConcurrentCallsWithAvailableImage(t *testing.T
 	for i := range numGoroutines {
 		go func(index int) {
 			<-startChan
-			started[index] = StartDockerImageDownload(context.Background(), testImage)
+			started[index], _ = StartDockerImageDownload(context.Background(), testImage)
 			doneChan <- index
 		}(i)
 	}
@@ -459,7 +459,8 @@ func TestStartDockerImageDownload_RaceWithExternalDownload(t *testing.T) {
 
 	for range numGoroutines {
 		go func() {
-			results <- StartDockerImageDownload(context.Background(), testImage)
+			started, _ := StartDockerImageDownload(context.Background(), testImage)
+			results <- started
 		}()
 	}
 
@@ -491,7 +492,7 @@ func TestStartDockerImageDownload_ContextCancellation(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 
 	// Start the download
-	started := StartDockerImageDownload(ctx, testImage)
+	started, join := StartDockerImageDownload(ctx, testImage)
 	if !started {
 		t.Fatal("Expected download to start")
 	}
@@ -504,8 +505,8 @@ func TestStartDockerImageDownload_ContextCancellation(t *testing.T) {
 	// Cancel the context immediately
 	cancel()
 
-	// Wait a bit for the goroutine to notice the cancellation
-	time.Sleep(100 * time.Millisecond)
+	// Join the download goroutine and ensure cleanup is complete
+	join()
 
 	// The image should no longer be marked as downloading after cancellation
 	if IsDockerImageDownloading(testImage) {
@@ -516,14 +517,83 @@ func TestStartDockerImageDownload_ContextCancellation(t *testing.T) {
 	ResetDockerPullState()
 }
 
+func TestStartDockerImageDownload_JoinPointForExistingDownload(t *testing.T) {
+	ResetDockerPullState()
+
+	testImage := "test/join-existing:v1.0.0"
+	SetMockImageAvailable(testImage, false)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	startedFirst, joinFirst := StartDockerImageDownload(ctx, testImage)
+	if !startedFirst {
+		t.Fatal("Expected first call to start download")
+	}
+
+	startedSecond, joinSecond := StartDockerImageDownload(ctx, testImage)
+	if startedSecond {
+		t.Fatal("Expected second call to observe existing download")
+	}
+
+	secondJoined := make(chan struct{})
+	go func() {
+		defer close(secondJoined)
+		joinSecond()
+	}()
+
+	select {
+	case <-secondJoined:
+		t.Fatal("Expected second join to block while shared download is still running")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	cancel()
+	joinSecond()
+	joinFirst()
+
+	if IsDockerImageDownloading(testImage) {
+		t.Error("Expected image to not be marked as downloading after joined cancellation")
+	}
+
+	ResetDockerPullState()
+}
+
+func TestStartDockerImageDownload_JoinPointNoopWhenImageAvailable(t *testing.T) {
+	ResetDockerPullState()
+
+	testImage := "test/join-noop:v1.0.0"
+	SetMockImageAvailable(testImage, true)
+
+	started, join := StartDockerImageDownload(context.Background(), testImage)
+	if started {
+		t.Fatal("Expected download not to start for already-available image")
+	}
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		join()
+	}()
+
+	select {
+	case <-done:
+		// expected: join is a no-op when no goroutine was started
+	case <-time.After(2 * time.Second):
+		t.Fatal("Expected join to return immediately when image is available")
+	}
+
+	ResetDockerPullState()
+}
+
 func TestStartDockerImageDownload_NilContext(t *testing.T) {
 	ResetDockerPullState()
 
 	testImage := "test/nil-context-download:v1.0.0"
 	SetMockImageAvailable(testImage, true)
 
 	//nolint:staticcheck // Intentionally validating nil context handling behavior.
-	if StartDockerImageDownload(nil, testImage) {
+	started, _ := StartDockerImageDownload(nil, testImage)
+	if started {
 		t.Error("Expected download not to start for available image with nil context")
 	}
 

pkg/cli/docker_images_test_helpers_test.go

@@ -7,6 +7,7 @@ func ResetDockerPullState() {
 	pullState.mu.Lock()
 	defer pullState.mu.Unlock()
 	pullState.downloading = make(map[string]bool)
+	pullState.inflight = make(map[string]chan struct{})
 	pullState.mockAvailable = make(map[string]bool)
 	pullState.mockAvailableInUse = false
 	pullState.mockDockerAvailable = true