fix: Fixing provider cache test flake (gruntwork-io#6056)

Author: yhakbar

docs/src/data/changelog/v1.0.4/cache-server-graceful-shutdown.mdx

@@ -0,0 +1,10 @@
+---
+version: "v1.0.4"
+category: "bug-fixes"
+---
+
+#### Provider Cache Server now drains pending downloads on shutdown
+
+When Terragrunt finished a run with the Provider Cache Server enabled, the cache server stopped instantly rather than letting in-flight provider downloads complete. Downloads still being served to OpenTofu/Terraform were cut off mid-response, and successful runs ended with a `context canceled` error in the logs.
+
+Terragrunt now waits for those downloads to finish before stopping the cache server, and a graceful shutdown is no longer reported as an error.

internal/providercache/providercache_test.go

@@ -109,164 +109,72 @@ func TestProviderCache(t *testing.T) {
 		t.Run(fmt.Sprintf("testCase-%d", i), func(t *testing.T) {
 			t.Parallel()
 
-			// TODO: Remove this once we can invest time in figuring out why this test is so flaky.
-			//
-			// It's a pain, but it's not worth the time to fix it.
-			maxRetries := 3
+			ctx, cancel := context.WithCancel(t.Context())
+			defer cancel()
 
-			var lastErr error
+			errGroup, ctx := errgroup.WithContext(ctx)
+			l := logger.CreateLogger()
 
-			for attempt := 1; attempt <= maxRetries; attempt++ {
-				if attempt > 1 {
-					t.Logf("Retry attempt %d/%d for test case %d", attempt, maxRetries, i)
-				}
+			providerService := services.NewProviderService(providerCacheDir, pluginCacheDir, nil, l)
+			providerHandler := handlers.NewDirectProviderHandler(l, new(cliconfig.ProviderInstallationDirect), nil)
+			proxyProviderHandler := handlers.NewProxyProviderHandler(l, nil)
 
-				// Create a new context for each test case to avoid interference
-				//
-				//nolint:usetesting
-				ctx := context.Background()
+			tc.opts = append(tc.opts,
+				cache.WithProviderService(providerService),
+				cache.WithProviderHandlers(providerHandler),
+				cache.WithProxyProviderHandler(proxyProviderHandler),
+			)
 
-				ctx, cancel := context.WithCancel(ctx)
-				defer cancel()
+			server := cache.NewServer(tc.opts...)
 
-				errGroup, ctx := errgroup.WithContext(ctx)
-				logger := logger.CreateLogger()
+			ln, err := server.Listen(t.Context())
+			require.NoError(t, err)
 
-				providerService := services.NewProviderService(providerCacheDir, pluginCacheDir, nil, logger)
-				providerHandler := handlers.NewDirectProviderHandler(logger, new(cliconfig.ProviderInstallationDirect), nil)
-				proxyProviderHandler := handlers.NewProxyProviderHandler(logger, nil)
+			defer ln.Close()
 
-				tc.opts = append(tc.opts,
-					cache.WithProviderService(providerService),
-					cache.WithProviderHandlers(providerHandler),
-					cache.WithProxyProviderHandler(proxyProviderHandler),
-				)
+			errGroup.Go(func() error {
+				return server.Run(ctx, ln)
+			})
 
-				server := cache.NewServer(tc.opts...)
+			urlPath := server.ProviderController.URL()
+			urlPath.Path += tc.relURLPath
 
-				ln, err := server.Listen(t.Context())
-				if err != nil {
-					lastErr = err
-
-					if attempt < maxRetries {
-						continue
-					}
-
-					require.NoError(t, err)
-				}
-				defer ln.Close()
-
-				errGroup.Go(func() error {
-					return server.Run(ctx, ln)
-				})
-
-				urlPath := server.ProviderController.URL()
-				urlPath.Path += tc.relURLPath
-
-				if tc.fullURLPath != "" {
-					urlPath.Path = tc.fullURLPath
-				}
-
-				req, err := http.NewRequestWithContext(ctx, http.MethodGet, urlPath.String(), nil)
-				if err != nil {
-					lastErr = err
-
-					if attempt < maxRetries {
-						continue
-					}
-
-					require.NoError(t, err)
-				}
-
-				req.Header.Set("Authorization", "Bearer "+token)
-
-				resp, err := http.DefaultClient.Do(req)
-				if err != nil {
-					lastErr = err
-
-					if attempt < maxRetries {
-						continue
-					}
-
-					require.NoError(t, err)
-				}
-				defer resp.Body.Close()
-
-				if resp.StatusCode != tc.expectedStatusCode {
-					lastErr = fmt.Errorf("expected status code %d, got %d", tc.expectedStatusCode, resp.StatusCode)
-
-					if attempt < maxRetries {
-						continue
-					}
-
-					assert.Equal(t, tc.expectedStatusCode, resp.StatusCode)
-				}
-
-				if tc.expectedBodyReg != nil {
-					body, err := io.ReadAll(resp.Body)
-					if err != nil {
-						lastErr = err
-
-						if attempt < maxRetries {
-							continue
-						}
-
-						require.NoError(t, err)
-					}
-
-					if !tc.expectedBodyReg.MatchString(string(body)) {
-						lastErr = fmt.Errorf("body did not match expected regex: %s", tc.expectedBodyReg.String())
-
-						if attempt < maxRetries {
-							continue
-						}
-
-						assert.Regexp(t, tc.expectedBodyReg, string(body))
-					}
-				}
-
-				// Skip WaitForCacheReady for unauthorized test cases since they don't trigger background operations,
-				// and we cancel context at the end of the test.
-				if tc.expectedStatusCode != http.StatusUnauthorized {
-					_, err = providerService.WaitForCacheReady("")
-					if err != nil {
-						lastErr = err
-
-						if attempt < maxRetries {
-							continue
-						}
+			if tc.fullURLPath != "" {
+				urlPath.Path = tc.fullURLPath
+			}
 
-						require.NoError(t, err)
-					}
-				}
+			req, err := http.NewRequestWithContext(ctx, http.MethodGet, urlPath.String(), nil)
+			require.NoError(t, err)
 
-				if tc.expectedCachePath != "" {
-					if !assert.FileExists(t, filepath.Join(providerCacheDir, tc.expectedCachePath)) {
-						lastErr = fmt.Errorf("expected cache file does not exist: %s", tc.expectedCachePath)
+			req.Header.Set("Authorization", "Bearer "+token)
 
-						if attempt < maxRetries {
-							continue
-						}
-					}
-				}
+			resp, err := http.DefaultClient.Do(req)
+			require.NoError(t, err)
 
-				cancel()
+			defer resp.Body.Close()
 
-				err = errGroup.Wait()
-				if err != nil {
-					lastErr = err
+			assert.Equal(t, tc.expectedStatusCode, resp.StatusCode)
 
-					if attempt < maxRetries {
-						continue
-					}
+			if tc.expectedBodyReg != nil {
+				body, err := io.ReadAll(resp.Body)
+				require.NoError(t, err)
+				assert.Regexp(t, tc.expectedBodyReg, string(body))
+			}
 
-					require.NoError(t, err)
-				}
+			// Skip WaitForCacheReady for unauthorized test cases since they don't trigger background operations,
+			// and we cancel context at the end of the test.
+			if tc.expectedStatusCode != http.StatusUnauthorized {
+				_, err = providerService.WaitForCacheReady("")
+				require.NoError(t, err)
+			}
 
-				return
+			if tc.expectedCachePath != "" {
+				assert.FileExists(t, filepath.Join(providerCacheDir, tc.expectedCachePath))
 			}
 
-			t.Fatalf("Test case %d failed after %d attempts. Last error: %v", i, maxRetries, lastErr)
+			cancel()
+
+			require.NoError(t, errGroup.Wait())
 		})
 	}
 }

internal/tf/cache/server.go

@@ -103,10 +103,13 @@ func (server *Server) Run(ctx context.Context, ln net.Listener) error {
 		<-ctx.Done()
 		server.logger.Infof("Shutting down Terragrunt Cache server...")
 
-		ctx, cancel := context.WithTimeout(ctx, server.shutdownTimeout)
+		// The parent ctx is by definition already cancelled here; detach from
+		// its cancellation (preserving any values) so http.Server.Shutdown gets
+		// the full shutdownTimeout to drain in-flight requests.
+		shutdownCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), server.shutdownTimeout)
 		defer cancel()
 
-		if err := server.Shutdown(ctx); err != nil {
+		if err := server.Shutdown(shutdownCtx); err != nil {
 			return errors.New(err)
 		}
 
@@ -119,5 +122,9 @@ func (server *Server) Run(ctx context.Context, ln net.Listener) error {
 
 	defer server.logger.Infof("Terragrunt Cache server stopped")
 
-	return errGroup.Wait()
+	if err := errGroup.Wait(); err != nil && !errors.Is(err, context.Canceled) {
+		return err
+	}
+
+	return nil
 }