How the Keyed Package Manages Goroutine Lifecycle with Context Propagation

The keyed package orchestrates per-key goroutines by deriving child contexts from a user-supplied root context, enabling cancellation propagation, automatic restarts with back-off, and coordinated lifecycle management through SetContext and the runningRoutine state machine.

The aperturerobotics/util repository provides a robust keyed package that solves the complex problem of managing multiple long-running goroutines keyed by unique identifiers. By leveraging Go's context.Context for cancellation propagation, the package ensures that updating or cancelling a root context gracefully cascades to all active routines while supporting individual retry logic and configurable release delays.

Root Context Management in keyed/keyed.go

The lifecycle of all managed goroutines originates from a single root context stored in Keyed.ctx. According to the source code in keyed/keyed.go (lines 35-42), this root context holds the overall cancellation signal for the entire manager.

When callers invoke SetContext(ctx, restart), the method acquires the manager's mutex and delegates to setContextLocked (lines 88-108). This internal method updates the root context and iterates through all runningRoutine entries. For each routine, it clears the child context reference (rr.ctx = nil) and invokes the stored ctxCancel function to immediately terminate the existing goroutine.

If the restart parameter is true, or if the routine previously exited with an error, setContextLocked triggers runningRoutine.start with the new root context. This design ensures that context updates can selectively restart failed routines while leaving successfully completed ones untouched unless explicitly requested.

Per-Routine Lifecycle Control

Each logical key maps to a runningRoutine instance defined in keyed/routine.go (lines 11-38). This struct tracks the routine's state, including its derived child context (ctx), cancellation function (ctxCancel), exit channel, error status, success flag, and optional back-off configuration.

Context Derivation and Routine Initialization

The runningRoutine.start method (lines 68-79) handles the actual goroutine spawning. It first checks if the routine already succeeded or is still running, respecting the forceRestart flag. When starting, it derives a fresh child context using context.WithCancel(ctx) from the current root context stored in the manager. This child context is what the user-provided Routine function receives as its sole argument.

If a previous child context exists, it is cancelled before spawning the new goroutine. The method then launches runningRoutine.execute in a separate goroutine, passing the newly created child context.

Execution Loop and Cancellation Handling

Inside runningRoutine.execute (lines 99-155), the routine optionally waits on a waitCh used during retry delays, then invokes the user function r.routine(ctx). Upon return—whether success, failure, or cancellation—the method immediately calls cancel() to clean up the child context, closes the exit channel, and records the error state.

The method sets a success flag when err == nil and marks the routine as exited. This state determines whether the routine qualifies for automatic restart when the root context changes or if it should remain dormant.

Automatic Restart and Back-off Policies

The keyed package supports resilient operation through configurable retry mechanisms implemented in keyed/routine.go (lines 30-48 and 130-148). When execute detects an error and a back-off factory is configured via WithRetry or WithBackoff, it schedules a retry using time.AfterFunc.

The retry logic reacquires the manager lock and calls start with forceRestart=true, creating a brand new child context from the current root. Successful runs reset the back-off timer, while consecutive failures use progressively longer delays according to the configured policy.

Cleanup Strategies and Release Delays

Removing a key triggers runningRoutine.remove (lines 58-90). If configured with WithReleaseDelay from keyed/keyed-opts.go, the routine enters a delayed deletion state. Instead of immediate cancellation, time.AfterFunc postpones the actual map deletion for the specified duration, allowing in-flight operations to complete. Without a release delay, the routine is cancelled and removed immediately from the internal map.

All state transitions are protected by Keyed.mtx, ensuring thread-safe operations during concurrent additions, removals, and context updates.

Practical Implementation Examples

The following examples demonstrate how to leverage context propagation and lifecycle management in real-world scenarios.

Basic Context Cancellation

package main

import (
    "context"
    "fmt"
    "time"
    
    "github.com/aperturerobotics/util/keyed"
)

func workerRoutine(ctx context.Context) error {
    select {
    case <-ctx.Done():
        fmt.Println("Worker received cancellation signal")
        return ctx.Err()
    case <-time.After(10 * time.Second):
        return nil
    }
}

func main() {
    km := keyed.NewKeyed[string, struct{}](
        func(key string) (keyed.Routine, struct{}) {
            return workerRoutine, struct{}{}
        },
    )
    
    km.SetKey("worker-1", true)
    
    rootCtx, cancel := context.WithCancel(context.Background())
    km.SetContext(rootCtx, false)
    
    // Cancel all routines after 2 seconds
    time.Sleep(2 * time.Second)
    cancel()
    
    // Cleanup
    km.ClearContext()
}

In keyed/keyed.go, SetContext derives child contexts for each key. When cancel() is invoked, the cancellation propagates through the child context to workerRoutine.

Retry with Exponential Back-off

var attemptCount int

func flakyRoutine(ctx context.Context) error {
    attemptCount++
    if attemptCount < 3 {
        return fmt.Errorf("attempt %d failed", attemptCount)
    }
    fmt.Println("Success on attempt", attemptCount)
    return nil
}

km := keyed.NewKeyed[int, struct{}](
    func(key int) (keyed.Routine, struct{}) {
        return flakyRoutine, struct{}{}
    },
    keyed.WithBackoff[int, struct{}](func(_ int) backoff.BackOff {
        return backoff.NewExponentialBackOff()
    }),
)

km.SetKey(1, true)
km.SetContext(context.Background(), false)

// The manager automatically retries using the back-off policy
// defined in keyed/routine.go lines 130-148

Graceful Shutdown with Release Delay

km := keyed.NewKeyed[string, struct{}](
    func(key string) (keyed.Routine, struct{}) {
        return workerRoutine, struct{}{}
    },
    keyed.WithReleaseDelay[string, struct{}](5*time.Second),
)

km.SetKey("temp", true)
km.SetContext(context.Background(), false)

// Routine continues for 5 seconds after removal
km.RemoveKey("temp")

Summary

• The keyed package manages per-key goroutines through a hierarchical context structure where a root context drives cancellation for all child routines.
• SetContext in keyed/keyed.go coordinates mass updates by cancelling existing child contexts and optionally restarting routines based on their error state.
• Each runningRoutine in keyed/routine.go maintains its own cancellable child context, enabling fine-grained lifecycle control while respecting the parent root.
• Automatic retry mechanisms use time.AfterFunc to schedule restarts with configurable back-off policies, resetting on successful execution.
• The WithReleaseDelay option allows routines to complete in-flight work before removal, while the mutex-protected state machine ensures thread-safe concurrent access.

Frequently Asked Questions

How does changing the root context affect running goroutines?

When SetContext is called, the manager invokes setContextLocked (lines 88-108 in keyed/keyed.go), which iterates through all active routines and cancels their child contexts by calling the stored ctxCancel function. If the restart parameter is true, or if a routine previously failed, it immediately starts a new goroutine with a fresh child context derived from the new root. Successfully completed routines are only restarted if explicitly requested.

Can individual routines have different back-off policies?

Yes. The WithBackoff option accepts a factory function that receives the key as an argument, allowing you to return different BackOff implementations per key. As implemented in keyed/keyed-opts.go, this factory is stored in the routine configuration and invoked when runningRoutine.execute schedules a retry via time.AfterFunc (lines 130-148 in keyed/routine.go).

What happens if a routine ignores context cancellation?

The runningRoutine.execute method (lines 99-155 in keyed/routine.go) calls the user routine and then immediately invokes cancel() on the child context after the function returns, regardless of whether the routine honored the cancellation. However, if the user routine blocks indefinitely without checking ctx.Done(), the goroutine will leak until the function returns. The manager cannot forcefully kill goroutines; it relies on the user code respecting context cancellation.

How does the release delay interact with context cancellation?

The WithReleaseDelay option creates a grace period between RemoveKey being called and the actual deletion of the routine from the internal map. During this delay, the routine continues running with its existing child context. If the root context is cancelled during the delay, the normal cancellation propagation still occurs through setContextLocked, which will cancel the routine's context before the release timer expires. The delay only affects map retention, not context lifecycle.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}