fix(L7): use shutdown context in Snowflake async continuation (#296) - #311
fix(L7): use shutdown context in Snowflake async continuation (#296)#311mvanhorn wants to merge 1 commit into
Conversation
…-data-and-ai#296) `startSnowflakeAsyncContinuation` launched a background goroutine with `context.Background()`, so a manager shutdown via SIGINT/SIGTERM left the goroutine writing to a Snowflake user channel + the data store cache long after the controller manager had stopped — exactly the "leaks on shutdown" failure mode reported in redhat-data-and-ai#296. Plumb a single `shutdownCtx := ctrl.SetupSignalHandler()` from main through `preloadCache` into `startSnowflakeAsyncContinuation` so the async goroutine derives its working context from the same signal context the manager uses. Cancellation now propagates into both `FetchRemainingUsersAsync` and the per-user `dataStore.User.SetBackend` calls so the goroutine winds down cleanly when shutdown fires instead of racing the process exit. `ctrl.SetupSignalHandler()` can only be called once per process; the prior call site at `mgr.Start(ctrl.SetupSignalHandler())` is replaced with `mgr.Start(shutdownCtx)` so we keep the same one-call contract. The completion/error reporting half of redhat-data-and-ai#296 (health-probe surfacing) is intentionally left for a follow-up so each PR stays focused on a single concern. Verification: go build ./... # ok go vet ./... # ok go test $(go list ./... | grep -v e2e | grep -v internal/controller) # ok (controller # envtest needs # kubebuilder bins, # not exercised here)
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a resource leak issue where background goroutines responsible for Snowflake async continuations persisted after the controller manager received a shutdown signal. By plumbing a unified shutdown context from the application entry point down to the asynchronous tasks, the system now ensures that all long-lived background operations are properly canceled when the process receives SIGINT or SIGTERM. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. The signal arrives at the gate, To stop all the work, not too late. With context now passed, The leak will not last, And shutdown proceeds in a state. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors context propagation in cmd/main.go to pass a shared shutdown context to preloadCache and the subsequent Snowflake async continuation, ensuring background goroutines stop cleanly on SIGINT/SIGTERM. The review feedback highlights critical issues regarding context and logger traceability: shadowing the context with the errgroup context causes the request ID to be lost, and passing shutdownCtx directly to the async continuation prevents it from accessing the request ID. The reviewer suggests using a separate preloadCtx to preserve both the request ID and the shutdown signal, and simplifying startSnowflakeAsyncContinuation to accept this context directly.
| // Derive from the shutdown context so any background continuation kicked | ||
| // off after preload (e.g. Snowflake async, #296) gets canceled on | ||
| // shutdown. preloadCache itself is synchronous so this only matters for | ||
| // long-lived children launched below. | ||
| ctx := logger.WithRequestId(shutdownCtx, types.UID(uuid.New().String())) | ||
|
|
||
| // Use errgroup for concurrent processing with proper error handling | ||
| g, ctx := errgroup.WithContext(ctx) |
There was a problem hiding this comment.
By shadowing ctx with the errgroup context, the original context (which contains the request ID and is derived from shutdownCtx) is lost for any code after the errgroup initialization.
Instead of shadowing ctx, we can define a separate preloadCtx that holds the long-lived context with the request ID, and then derive the errgroup context ctx from it. This allows us to pass preloadCtx to the Snowflake async continuation so it retains both the request ID and the shutdown cancellation signal without being prematurely canceled when g.Wait() returns.
| // Derive from the shutdown context so any background continuation kicked | |
| // off after preload (e.g. Snowflake async, #296) gets canceled on | |
| // shutdown. preloadCache itself is synchronous so this only matters for | |
| // long-lived children launched below. | |
| ctx := logger.WithRequestId(shutdownCtx, types.UID(uuid.New().String())) | |
| // Use errgroup for concurrent processing with proper error handling | |
| g, ctx := errgroup.WithContext(ctx) | |
| // Derive from the shutdown context so any background continuation kicked | |
| // off after preload (e.g. Snowflake async, #296) gets canceled on | |
| // shutdown. preloadCache itself is synchronous so this only matters for | |
| // long-lived children launched below. | |
| preloadCtx := logger.WithRequestId(shutdownCtx, types.UID(uuid.New().String())) | |
| // Use errgroup for concurrent processing with proper error handling | |
| g, ctx := errgroup.WithContext(preloadCtx) |
| for _, state := range snowflakeStates { | ||
| if state.lastUser != "" { | ||
| startSnowflakeAsyncContinuation(ctx, state, dataStore, cacheMutex) | ||
| startSnowflakeAsyncContinuation(shutdownCtx, state, dataStore, cacheMutex) | ||
| } | ||
| } |
There was a problem hiding this comment.
Pass the long-lived preloadCtx (which contains the request ID and is tied to shutdownCtx) instead of shutdownCtx. This ensures the Snowflake async continuation has access to the request ID for log traceability.
| for _, state := range snowflakeStates { | |
| if state.lastUser != "" { | |
| startSnowflakeAsyncContinuation(ctx, state, dataStore, cacheMutex) | |
| startSnowflakeAsyncContinuation(shutdownCtx, state, dataStore, cacheMutex) | |
| } | |
| } | |
| for _, state := range snowflakeStates { | |
| if state.lastUser != "" { | |
| startSnowflakeAsyncContinuation(preloadCtx, state, dataStore, cacheMutex) | |
| } | |
| } |
| func startSnowflakeAsyncContinuation( | ||
| originalCtx context.Context, | ||
| shutdownCtx context.Context, | ||
| state *snowflakeAsyncState, | ||
| dataStore *store.Store, | ||
| cacheMutex *sync.RWMutex, | ||
| ) { | ||
| // Create a fresh context since the errgroup context is canceled after g.Wait() returns | ||
| // Transfer the logger (with request ID) from original context for traceability | ||
| asyncCtx := context.Background() | ||
| if entry := logger.Logger(originalCtx); entry != nil { | ||
| asyncCtx = context.WithValue(asyncCtx, logger.RequestIdKey, entry) | ||
| // Derive a child context from the manager's shutdown context so SIGINT/SIGTERM | ||
| // propagates into FetchRemainingUsersAsync and the cache writes below. | ||
| // Transfer the request-id logger value from the originating context for | ||
| // traceability. | ||
| asyncCtx := shutdownCtx | ||
| if entry := logger.Logger(shutdownCtx); entry != nil { | ||
| asyncCtx = context.WithValue(shutdownCtx, logger.RequestIdKey, entry) | ||
| } |
There was a problem hiding this comment.
Since we are now passing a long-lived context (preloadCtx) that is already derived from shutdownCtx and contains the request ID, we do not need to manually recreate the context or transfer the logger. We can simply accept asyncCtx directly and use it.
Additionally, the current implementation has a bug where logger.Logger(shutdownCtx) is called on shutdownCtx (which does not contain the request ID), resulting in the request ID being completely lost.
| func startSnowflakeAsyncContinuation( | |
| originalCtx context.Context, | |
| shutdownCtx context.Context, | |
| state *snowflakeAsyncState, | |
| dataStore *store.Store, | |
| cacheMutex *sync.RWMutex, | |
| ) { | |
| // Create a fresh context since the errgroup context is canceled after g.Wait() returns | |
| // Transfer the logger (with request ID) from original context for traceability | |
| asyncCtx := context.Background() | |
| if entry := logger.Logger(originalCtx); entry != nil { | |
| asyncCtx = context.WithValue(asyncCtx, logger.RequestIdKey, entry) | |
| // Derive a child context from the manager's shutdown context so SIGINT/SIGTERM | |
| // propagates into FetchRemainingUsersAsync and the cache writes below. | |
| // Transfer the request-id logger value from the originating context for | |
| // traceability. | |
| asyncCtx := shutdownCtx | |
| if entry := logger.Logger(shutdownCtx); entry != nil { | |
| asyncCtx = context.WithValue(shutdownCtx, logger.RequestIdKey, entry) | |
| } | |
| func startSnowflakeAsyncContinuation( | |
| asyncCtx context.Context, | |
| state *snowflakeAsyncState, | |
| dataStore *store.Store, | |
| cacheMutex *sync.RWMutex, | |
| ) { |
References
- Always pass context.Context as the first parameter to functions that may block. (link)
1 participant
Summary
startSnowflakeAsyncContinuationlaunched a background goroutine withcontext.Background(), so a manager shutdown via SIGINT / SIGTERM left the goroutine running, writing to a Snowflake user channel and the data-store cache long after the controller manager had stopped. That is the "leaks on shutdown" failure mode reported in #296.This PR plumbs a single shutdown context from
maininto the async continuation so cancellation propagates into bothFetchRemainingUsersAsyncand the per-userdataStore.User.SetBackendcalls.What changed
cmd/main.go: captureshutdownCtx := ctrl.SetupSignalHandler()at the top of the lifecycle and reuse it for bothpreloadCache(...)andmgr.Start(shutdownCtx). The priormgr.Start(ctrl.SetupSignalHandler())is removed becauseSetupSignalHandlercan only be called once per process.preloadCache: takes the shutdown context as its first arg; derives its internal errgroup context from it; passes it (not the errgroup ctx, which gets canceled wheng.Waitreturns) intostartSnowflakeAsyncContinuationso the goroutine survivespreloadCachereturning while still listening for shutdown.startSnowflakeAsyncContinuation: receives the shutdown context and derivesasyncCtxfrom it (preserving the request-id logger value the prior call attached).FetchRemainingUsersAsyncand the per-user cache writes now see SIGINT / SIGTERM as a normal context cancellation.Scope
This addresses the cancelable-context leg of #296. The companion concern — wiring completion / error reporting into the health probe so the goroutine's terminal status surfaces externally — is intentionally left for a follow-up so each PR stays focused on a single concern.
Verification
Closes the cancelable-context half of #296.