diff --git a/pkg/cli/docker_images.go b/pkg/cli/docker_images.go index 7b2159414cb..d4eefb47c04 100644 --- a/pkg/cli/docker_images.go +++ b/pkg/cli/docker_images.go @@ -43,6 +43,7 @@ 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) @@ -50,6 +51,7 @@ type dockerPullState struct { 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) } } diff --git a/pkg/cli/docker_images_test.go b/pkg/cli/docker_images_test.go index 3d4c39f8128..b77de597188 100644 --- a/pkg/cli/docker_images_test.go +++ b/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,6 +517,74 @@ 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() @@ -523,7 +592,8 @@ func TestStartDockerImageDownload_NilContext(t *testing.T) { 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") } diff --git a/pkg/cli/docker_images_test_helpers_test.go b/pkg/cli/docker_images_test_helpers_test.go index 2539f3d0839..de9504d715a 100644 --- a/pkg/cli/docker_images_test_helpers_test.go +++ b/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