Skip to content

fix(L7): use shutdown context in Snowflake async continuation (#296) #311

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mvanhorn wants to merge 1 commit into
base: main
Choose a base branch
from
+38 −15
Open

fix(L7): use shutdown context in Snowflake async continuation (#296) #311

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 38 additions & 15 deletions cmd/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,14 @@ func main() {
// Create store layer that wraps cache with prefixed keys and encapsulated operations
dataStore := store.New(cache)

if err = preloadCache(*appConf, dataStore, sharedCacheMutex); err != nil {
// Set up the shared shutdown context up front so the Snowflake async
// continuation (#296) can derive a cancelable child from it and stop
// cleanly when the manager receives SIGINT/SIGTERM. `SetupSignalHandler`
// can only be called once per process; reuse the same context for
// `mgr.Start` below.
shutdownCtx := ctrl.SetupSignalHandler()

if err = preloadCache(shutdownCtx, *appConf, dataStore, sharedCacheMutex); err != nil {
setupLog.Error(err, "failed to preload cache")
os.Exit(1)
}
Expand Down Expand Up @@ -260,7 +267,7 @@ func main() {
}
}()
setupLog.Info("starting manager")
if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
if err := mgr.Start(shutdownCtx); err != nil {
setupLog.Error(err, "problem running manager")
os.Exit(1)
}
Expand Down Expand Up @@ -294,11 +301,17 @@ type snowflakeAsyncState struct {
// This is done only once at the start of the application
// and the cache is flushed when the application is restarted
// Optimized to use goroutines for parallel processing of backends
func preloadCache(appConfig config.AppConfig, dataStore *store.Store, cacheMutex *sync.RWMutex) error {
ctx := context.Background()

// Add request ID for tracking this cache preload operation in logs
ctx = logger.WithRequestId(ctx, types.UID(uuid.New().String()))
func preloadCache(
shutdownCtx context.Context,
appConfig config.AppConfig,
dataStore *store.Store,
cacheMutex *sync.RWMutex,
) error {
// 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)
Comment on lines +310 to 317

@gemini-code-assist gemini-code-assist Bot May 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Suggested change
// 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)

Expand Down Expand Up @@ -419,10 +432,13 @@ func preloadCache(appConfig config.AppConfig, dataStore *store.Store, cacheMutex

setupLog.Info("cache preload completed for all backends")

// Start async continuation for all Snowflake backends (after all preloads done)
// Start async continuation for all Snowflake backends (after all preloads done).
// Pass the shutdownCtx (not the errgroup ctx, which gets canceled when g.Wait
// returns) so the goroutine survives preloadCache returning yet still gets
// canceled when the manager receives SIGINT/SIGTERM (#296).
for _, state := range snowflakeStates {
if state.lastUser != "" {
startSnowflakeAsyncContinuation(ctx, state, dataStore, cacheMutex)
startSnowflakeAsyncContinuation(shutdownCtx, state, dataStore, cacheMutex)
}
}
Comment on lines 439 to 443

@gemini-code-assist gemini-code-assist Bot May 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Suggested change
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)
}
}


Expand All @@ -431,17 +447,24 @@ func preloadCache(appConfig config.AppConfig, dataStore *store.Store, cacheMutex

// startSnowflakeAsyncContinuation starts a background goroutine to fetch remaining Snowflake users
// and write them to cache. This function returns immediately after starting the goroutine.
//
// shutdownCtx must be the long-lived manager shutdown context (typically
// `ctrl.SetupSignalHandler()`'s result) so the goroutine gets canceled when
// the manager receives SIGINT/SIGTERM. Passing `context.Background()` here
// leaks the goroutine across shutdown — see #296.
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)
}
Comment on lines 455 to 468

@gemini-code-assist gemini-code-assist Bot May 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Suggested change
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
  1. Always pass context.Context as the first parameter to functions that may block. (link)

userChan, errChan := state.client.FetchRemainingUsersAsync(asyncCtx, state.lastUser)

Expand Down
Loading