Drain Pattern — Middle Level¶

Table of Contents¶

1. Introduction
2. Prerequisites
3. The Three Phases In Practice
4. Designing The Drainable Interface
5. Idle-Wait Drains
6. Deadline-Bounded Drain Designs
7. Drain vs Hard Stop — Decision Frame
8. HTTP Server Drain
9. gRPC Drain
10. Worker Pool Drain
11. Queue Consumer Drain
12. Drain With errgroup
13. Composing Drains
14. Testing Drain
15. Observability For Drain
16. Pitfalls At The Middle Level
17. Patterns That Often Help
18. Patterns That Often Hurt
19. Comparative Languages
20. Mid-Level Cheat Sheet
21. Self-Assessment Checklist
22. Summary
23. Further Reading

Introduction¶

Focus: "I can drain one goroutine. How do I drain a real service with HTTP, workers, queues, and downstreams — and how do I prove it works?"

At the junior level you learned the recipe: stop intake, wait for in-flight, close downstream. At the middle level you learn to apply that recipe across the whole shape of a service: an HTTP API that calls a worker pool that publishes to Kafka that writes to Postgres. Each layer has its own drain, and they must compose in the right order. You also learn to test drain — not just hope it works.

After reading this file you will:

• Design a Drainable interface and use it across components.
• Build idle-wait and deadline-bounded drains correctly.
• Drain HTTP, gRPC, worker pools, queue consumers.
• Compose multiple drains with proper ordering.
• Use errgroup to coordinate startup and shutdown of N goroutines.
• Write tests that simulate SIGTERM and assert clean drain.
• Emit the right metrics so production drains are observable.
• Avoid the dozen most common mid-level pitfalls.

Prerequisites¶

• The junior page on drain.
• Comfort with context.Context and its propagation rules.
• Familiarity with sync.WaitGroup, sync.Once, and channels.
• Some experience with HTTP servers and at least one message-queue library (Kafka, NATS, RabbitMQ, Redis).
• Awareness of golang.org/x/sync/errgroup.

The Three Phases In Practice¶

The three steps — stop intake, wait, close — generalise to phases. A typical mid-size service has:

• Phase 0 (pre-drain). Flip readiness. Log the start of drain.
• Phase 1 (stop intake). Close HTTP listener, pause queue consumer, stop scheduling new jobs.
• Phase 2 (wait in-flight). Block until handlers, workers, and in-flight jobs finish — bounded by deadline.
• Phase 3 (flush). Flush async producers (Kafka, log shippers, metrics). They have their own buffers.
• Phase 4 (close downstream). Close database pools, Redis clients, file handles.
• Phase 5 (post-drain). Log the end of drain, the duration, the count of force-cancellations.

Each phase has its own time budget. The total of all phases must fit inside the orchestrator's grace period.

A worked sequence¶

type Drainable interface {
    Drain(ctx context.Context) error
}

func drainAll(ctx context.Context, in []Drainable) error {
    for _, d := range in {
        if err := d.Drain(ctx); err != nil {
            return fmt.Errorf("drain %T: %w", d, err)
        }
    }
    return nil
}

// main:
order := []Drainable{httpServer, workerPool, kafkaProducer, dbPool}
if err := drainAll(drainCtx, order); err != nil {
    log.Printf("drain error: %v", err)
}

The slice expresses the order. Reorder by editing the slice; do not entangle drain semantics into each component.

Designing The Drainable Interface¶

A small interface goes a long way:

type Drainable interface {
    Drain(ctx context.Context) error
}

What it means:

• The implementer must stop accepting new work before returning.
• It blocks until either in-flight work is done, or ctx expires.
• It returns ctx.Err() on timeout, or another error for transient failures.
• It is safe to call exactly once. Multiple calls may panic or no-op.

A larger interface for richer cases:

type StartableDrainable interface {
    Start(ctx context.Context) error
    Drain(ctx context.Context) error
}

Start is non-blocking; it spawns goroutines and returns. Drain is the cleanup. The lifecycle is symmetric.

Naming conventions¶

• Drain — graceful, may take time.
• Stop — hard cancel, returns quickly.
• Close — release resources.

Some libraries combine these. Pick names that match your team's existing code, and document the semantics in the package doc comment.

Returning errors¶

Drain errors fall into three buckets:

1. Deadline-related. context.DeadlineExceeded. Common; not actionable.
2. Transient. "Could not flush Kafka buffer." Logged; investigated.
3. Programmer. "Drain called before Start." Should panic.

Use errors.Is to dispatch on bucket 1; log buckets 2 and 3 with full context.

Idle-Wait Drains¶

Sometimes there is no clean "drain me" signal — new work can keep arriving from internal sources even after external intake stops. An idle-wait drain handles this by declaring the system drained when it has been idle for a configurable quiet period.

type IdleWaiter struct {
    inFlight atomic.Int64
    closed   atomic.Bool
}

func (w *IdleWaiter) Enter() bool {
    if w.closed.Load() {
        return false
    }
    w.inFlight.Add(1)
    return true
}

func (w *IdleWaiter) Exit() { w.inFlight.Add(-1) }

func (w *IdleWaiter) WaitIdle(ctx context.Context, quietFor time.Duration) error {
    w.closed.Store(true)
    lastBusy := time.Now()
    t := time.NewTicker(20 * time.Millisecond)
    defer t.Stop()
    for {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-t.C:
            if w.inFlight.Load() == 0 {
                if time.Since(lastBusy) >= quietFor {
                    return nil
                }
            } else {
                lastBusy = time.Now()
            }
        }
    }
}

quietFor is the silence the system must exhibit before declaring drained. Use longer values (200ms–1s) when work can be re-enqueued internally; shorter values (50ms) when intake stops cleanly.

When to prefer idle-wait¶

• Job pipelines where stages can re-emit downstream.
• Saga-like flows where a transaction is many small steps.
• Systems with timer-driven internal events.

When to prefer wait-group¶

• Bounded job sets.
• Request/response flows where in-flight count is exact.

Deadline-Bounded Drain Designs¶

A deadline-bounded drain combines a wait with context.WithTimeout. Two structural choices:

Single deadline for the whole drain¶

drainCtx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
defer cancel()
_ = httpServer.Drain(drainCtx)
_ = workerPool.Drain(drainCtx)
_ = db.Close()

If httpServer.Drain takes 24 seconds, workerPool.Drain has only 1 second left. This is intentional — the deadline is the total budget. If you want more time, raise the budget.

Per-phase deadline¶

phaseCtx := func(d time.Duration) (context.Context, context.CancelFunc) {
    return context.WithTimeout(context.Background(), d)
}

httpCtx, c := phaseCtx(10 * time.Second); _ = httpServer.Drain(httpCtx); c()
poolCtx, c := phaseCtx(10 * time.Second); _ = workerPool.Drain(poolCtx); c()

Each phase gets its own budget. Easier to debug ("HTTP drain has its own 10 seconds; if it failed, only HTTP is to blame"). Worse for total wall-clock control.

Mix the two: a top-level deadline plus per-phase budgets derived from the same root.

total, cancel := context.WithTimeout(context.Background(), 25*time.Second)
defer cancel()
httpCtx, _ := context.WithTimeout(total, 10*time.Second)
poolCtx, _ := context.WithTimeout(total, 10*time.Second)

If httpCtx is bounded by both 10s and the parent's remaining time, the tighter bound wins.

Drain vs Hard Stop — Decision Frame¶

Not every program should drain. Use this frame to decide:

A mature service may switch modes at runtime: drain on normal SIGTERM, hard stop on a "panic" signal or detected corruption.

Hard stop done right¶

A "hard stop" is not unmanaged exit — it is a fast, intentional exit:

log.Println("hard stop")
os.Exit(1)

The exit code matters. Use 0 only for clean shutdowns; 1 or specific codes for forced exits. Orchestrators alert on non-zero exit codes.

HTTP Server Drain¶

The canonical Go HTTP drain uses http.Server.Shutdown:

ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
defer cancel()
if err := srv.Shutdown(ctx); err != nil {
    log.Printf("http shutdown: %v", err)
}

What Shutdown does:

1. Closes all listeners (no new connections).
2. Closes all idle keep-alive connections.
3. Waits for active requests to finish.
4. Returns nil on success, ctx.Err() on timeout.

What Shutdown does not do:

• Notify active clients that a shutdown is in progress (they continue normally).
• Cancel hijacked or WebSocket connections (you must do that yourself).
• Wait for Server.RegisterOnShutdown hooks to complete (they run concurrently).

Hijacked connections¶

WebSocket and SSE connections are typically hijacked. Shutdown does not close them. You must keep your own registry:

type WS struct {
    mu    sync.Mutex
    conns map[*websocket.Conn]struct{}
}

func (w *WS) Register(c *websocket.Conn) {
    w.mu.Lock()
    defer w.mu.Unlock()
    w.conns[c] = struct{}{}
}

func (w *WS) Unregister(c *websocket.Conn) {
    w.mu.Lock()
    defer w.mu.Unlock()
    delete(w.conns, c)
}

func (w *WS) DrainAll(ctx context.Context) {
    w.mu.Lock()
    for c := range w.conns {
        _ = c.WriteControl(websocket.CloseMessage,
            websocket.FormatCloseMessage(websocket.CloseGoingAway, "server shutdown"),
            time.Now().Add(time.Second),
        )
    }
    w.mu.Unlock()
    // optional: wait for them to close
}

Call DrainAll before srv.Shutdown so clients see the close frame and the server's drain is shorter.

Custom shutdown ordering¶

http.Server.RegisterOnShutdown lets you hook into the shutdown:

srv.RegisterOnShutdown(func() {
    log.Println("server is shutting down: stop background flushes")
})

The hook runs concurrently with the shutdown logic. Use it for side actions, not as the primary drain logic.

TLS, h2c, h3¶

The drain semantics are the same for TLS (ListenAndServeTLS + Shutdown) and h2c via the h2c package. For HTTP/3 via QUIC (e.g. quic-go), the library has its own CloseGracefully semantics — read its docs.

gRPC Drain¶

grpc-go exposes Server.GracefulStop:

done := make(chan struct{})
go func() {
    srv.GracefulStop()
    close(done)
}()

select {
case <-done:
    log.Println("grpc graceful stop complete")
case <-ctx.Done():
    srv.Stop() // hard stop
    <-done
}

GracefulStop waits for in-flight RPCs to finish. Stop (no graceful) cancels them. The pattern above is "graceful with deadline": try graceful, fall back to hard stop on timeout.

Streaming RPCs¶

A streaming RPC may run for hours. GracefulStop waits for it. To bound this, propagate the drain context into stream handlers:

func (s *Server) BigStream(stream pb.Service_BigStreamServer) error {
    ctx := stream.Context()
    for {
        select {
        case <-ctx.Done():
            return status.FromContextError(ctx.Err()).Err()
        case msg := <-source:
            if err := stream.Send(msg); err != nil {
                return err
            }
        }
    }
}

When the server starts to shut down, the stream context is cancelled, the handler returns, the stream closes.

Multi-server¶

For services exposing both HTTP and gRPC, drain them concurrently:

g, _ := errgroup.WithContext(ctx)
g.Go(func() error { return httpSrv.Shutdown(drainCtx) })
g.Go(func() error { grpcSrv.GracefulStop(); return nil })
_ = g.Wait()

The two drains run in parallel; the overall drain takes the maximum of the two, not the sum.

Worker Pool Drain¶

A worker pool with a job channel:

type Pool struct {
    jobs    chan Job
    workers int
    wg      sync.WaitGroup
    closed  atomic.Bool
    mu      sync.Mutex
}

func NewPool(workers, buf int) *Pool {
    return &Pool{jobs: make(chan Job, buf), workers: workers}
}

func (p *Pool) Start(ctx context.Context, handler func(context.Context, Job)) {
    for i := 0; i < p.workers; i++ {
        p.wg.Add(1)
        go func() {
            defer p.wg.Done()
            for {
                select {
                case <-ctx.Done():
                    return
                case j, ok := <-p.jobs:
                    if !ok {
                        return
                    }
                    handler(ctx, j)
                }
            }
        }()
    }
}

func (p *Pool) Submit(j Job) error {
    if p.closed.Load() {
        return errors.New("pool closed")
    }
    p.mu.Lock()
    defer p.mu.Unlock()
    if p.closed.Load() {
        return errors.New("pool closed")
    }
    p.jobs <- j
    return nil
}

func (p *Pool) Drain(ctx context.Context) error {
    p.mu.Lock()
    if p.closed.CompareAndSwap(false, true) {
        close(p.jobs)
    }
    p.mu.Unlock()
    done := make(chan struct{})
    go func() { p.wg.Wait(); close(done) }()
    select {
    case <-done:
        return nil
    case <-ctx.Done():
        return ctx.Err()
    }
}

Key points:

• closed is atomic for the fast path in Submit.
• mu guards the close; double-close is prevented with CompareAndSwap.
• Drain blocks on the wait group with a deadline.
• The handler receives the start context, not the drain context. The start context is the one cancelled when the service shuts down.

Dynamic pool sizing¶

A drain that shrinks a pool sends a sentinel to N workers:

type sentinelJob struct{}
func (p *Pool) Shrink(n int) {
    for i := 0; i < n; i++ {
        p.jobs <- sentinelJob{}
    }
}

Workers detect the sentinel and exit. The wait group reflects the new size.

Queue Consumer Drain¶

Consumers of message queues — Kafka, NATS, RabbitMQ, Redis lists, SQS — share a pattern:

1. Stop fetching new messages.
2. Process the messages already fetched.
3. Acknowledge / commit offsets.
4. Disconnect.

The complication is offset / ack semantics. Acknowledging too early loses messages on a crash; too late causes duplicates.

Kafka with segmentio/kafka-go¶

reader := kafka.NewReader(kafka.ReaderConfig{ /* ... */ })
defer reader.Close() // important: closes connections cleanly

for {
    m, err := reader.FetchMessage(ctx)
    if err != nil {
        break
    }
    if err := process(ctx, m); err != nil {
        log.Printf("process: %v", err)
        continue
    }
    if err := reader.CommitMessages(ctx, m); err != nil {
        log.Printf("commit: %v", err)
    }
}

On ctx cancellation, FetchMessage returns immediately. Any uncommitted message is left for the next consumer to re-read. This is at-least-once delivery — the application must be idempotent.

For the drain, we cancel the context, let FetchMessage return, and let any in-flight process finish (because we share the outer context with the loop, not the inner). To improve this, separate the contexts:

fetchCtx, fetchCancel := context.WithCancel(rootCtx)
processCtx := rootCtx // not cancelled during drain

for {
    m, err := reader.FetchMessage(fetchCtx)
    if err != nil {
        break
    }
    _ = process(processCtx, m)
    _ = reader.CommitMessages(processCtx, m)
}

On drain: fetchCancel() stops new fetches. The loop drains any prefetched messages (the library buffers internally), processes them, commits them, then exits cleanly.

NATS¶

nc, _ := nats.Connect(...)
sub, _ := nc.Subscribe(...)

// drain:
_ = sub.Drain() // stops receiving, processes pending
_ = nc.Drain()  // closes connection

NATS bakes the pattern into the library: Drain() is the right name.

RabbitMQ¶

ch.Cancel(consumerTag, false) // stop delivery
// process the channel of deliveries until empty
ch.Close()

The consumer tag identifies the subscription. Cancelling stops further deliveries; in-flight messages must still be ack'd or nack'd.

Redis lists¶

// poll loop:
for {
    select {
    case <-ctx.Done():
        return
    default:
        val, err := client.BLPop(ctx, time.Second, "queue").Result()
        // handle val
    }
}

Drain by cancelling ctx; the BLPop returns on context error. There are no offsets to commit; the message was already removed from the list.

Drain With errgroup¶

golang.org/x/sync/errgroup wraps a wait group with error propagation:

g, gCtx := errgroup.WithContext(rootCtx)
g.Go(func() error { return server.Run(gCtx) })
g.Go(func() error { return worker.Run(gCtx) })
g.Go(func() error { return consumer.Run(gCtx) })

if err := g.Wait(); err != nil {
    log.Printf("group: %v", err)
}

If any goroutine returns an error, gCtx is cancelled. The other goroutines see it and drain. Wait returns the first error.

This is the cleanest way to run multiple long-lived goroutines in a service. Pair it with signal.NotifyContext for the top-level:

ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer cancel()

g, gCtx := errgroup.WithContext(ctx)
g.Go(func() error { return server.Run(gCtx) })
g.Go(func() error { return worker.Run(gCtx) })

<-gCtx.Done()
drainCtx, dcancel := context.WithTimeout(context.Background(), 25*time.Second)
defer dcancel()

_ = server.Drain(drainCtx)
_ = worker.Drain(drainCtx)
_ = g.Wait()

The two contexts have different roles: gCtx signals "stop running"; drainCtx bounds the drain. Keep them separate.

Composing Drains¶

Composition often goes wrong. Three rules:

1. Drain in the opposite order of construction. If A depends on B, drain A before B. The reverse leaves A talking to a closed B.
2. Each drainable manages its own subtree. A worker pool drains its own workers; the parent service does not reach in.
3. Pass the drain context, not a fresh one. Components must share the deadline.

A useful pattern is a Lifecycle struct that owns the ordering:

type Lifecycle struct {
    starters []func(context.Context) error
    drains   []func(context.Context) error
}

func (l *Lifecycle) Add(start, drain func(context.Context) error) {
    l.starters = append(l.starters, start)
    l.drains = append(l.drains, drain)
}

func (l *Lifecycle) Run(ctx context.Context) error {
    for _, s := range l.starters {
        if err := s(ctx); err != nil {
            return err
        }
    }
    return nil
}

func (l *Lifecycle) Drain(ctx context.Context) error {
    for i := len(l.drains) - 1; i >= 0; i-- {
        if err := l.drains[i](ctx); err != nil {
            log.Printf("drain step %d: %v", i, err)
        }
    }
    return nil
}

Add registers a pair; Drain walks them in reverse. This is the simplest dependency-aware drainer.

For more complex graphs (where A and B are siblings but C depends on both), use a DAG drainer with topological sort.

Testing Drain¶

A drain you do not test is a drain that breaks silently. Test these properties:

1. Drain completes within budget on empty system¶

func TestDrainEmptyFast(t *testing.T) {
    svc := New()
    ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
    defer cancel()
    if err := svc.Drain(ctx); err != nil {
        t.Fatalf("expected nil, got %v", err)
    }
}

2. Drain waits for in-flight¶

func TestDrainWaitsForInFlight(t *testing.T) {
    svc := New()
    ctx := context.Background()
    svc.Start(ctx, 1)
    done := make(chan struct{})
    svc.Submit(func() {
        time.Sleep(50 * time.Millisecond)
        close(done)
    })
    dctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
    defer cancel()
    _ = svc.Drain(dctx)
    select {
    case <-done:
    default:
        t.Fatal("in-flight job did not complete")
    }
}

3. Drain honours deadline¶

func TestDrainHonoursDeadline(t *testing.T) {
    svc := New()
    ctx := context.Background()
    svc.Start(ctx, 1)
    svc.Submit(func() { time.Sleep(time.Hour) }) // hung
    dctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
    defer cancel()
    err := svc.Drain(dctx)
    if !errors.Is(err, context.DeadlineExceeded) {
        t.Fatalf("expected deadline exceeded, got %v", err)
    }
}

4. Drain rejects new submissions¶

func TestDrainRejectsSubmissions(t *testing.T) {
    svc := New()
    ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
    defer cancel()
    go svc.Drain(ctx)
    time.Sleep(10 * time.Millisecond)
    if err := svc.Submit(func(){}); err == nil {
        t.Fatal("expected error on submit after drain")
    }
}

5. Drain is idempotent¶

func TestDrainIdempotent(t *testing.T) {
    svc := New()
    ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
    defer cancel()
    _ = svc.Drain(ctx)
    if err := svc.Drain(ctx); err != nil {
        t.Fatalf("second drain: %v", err)
    }
}

6. No goroutine leaks¶

Use go.uber.org/goleak or count manually:

func TestNoLeaksAfterDrain(t *testing.T) {
    before := runtime.NumGoroutine()
    svc := New()
    ctx := context.Background()
    svc.Start(ctx, 4)
    dctx, cancel := context.WithTimeout(context.Background(), time.Second)
    defer cancel()
    _ = svc.Drain(dctx)
    time.Sleep(50 * time.Millisecond) // let goroutines fully exit
    if after := runtime.NumGoroutine(); after > before {
        t.Errorf("leak: before=%d after=%d", before, after)
    }
}

Always run drain tests with -race. Race detector catches more drain bugs than any other testing strategy.

Observability For Drain¶

Three categories of telemetry:

Metrics¶

• drain_duration_seconds{component} (histogram).
• drain_in_flight_at_start{component} (gauge / summary).
• drain_force_cancelled_total{component} (counter).
• drain_calls_total{result=success|deadline|error} (counter).

Logs¶

Structured logs with the component name, start time, duration, error:

log.With(
    "component", "worker_pool",
    "in_flight_start", n,
    "duration_ms", time.Since(start).Milliseconds(),
).Info("drain complete")

Traces¶

Each drain phase as a span. A drain trace looks like a Gantt chart of stages.

ctx, span := tracer.Start(ctx, "drain.http")
defer span.End()

Set span attributes for in_flight_start and the result. A failed drain shows up red in the trace UI.

Pitfalls At The Middle Level¶

Forgetting to close the inbound channel¶

A worker pool whose Drain does not close p.jobs blocks forever waiting on range p.jobs. Always close the channel.

Closing the channel from the wrong goroutine¶

If two goroutines both close the channel, the second panics. Use sync.Once or atomic.CompareAndSwap to guarantee single close.

Mixing intake context with drain context¶

If you derive the drain context from the cancelled root context, the drain has zero time.

// wrong
drainCtx := context.WithTimeout(rootCtx, 25*time.Second)

Once rootCtx is cancelled, drainCtx is also cancelled — the timeout is moot.

// right
drainCtx, cancel := context.WithTimeout(context.Background(), 25*time.Second)

Drain without health-check flip¶

The LB keeps sending traffic to a draining pod. Drain stays busy with new work. The deadline expires before in-flight reaches zero. Always flip readiness first.

Stopping the HTTP server before workers¶

If workers receive jobs via HTTP, stopping the HTTP server is fine first — but workers handling background jobs (cron-like) should drain first if HTTP responses depend on them.

Not bounding sub-drains¶

A top-level drain budget of 25s with a sub-drain of context.Background() (no timeout) means a single hung sub-component can block forever. Always derive sub-contexts from the drain context.

Async producers without Flush¶

A Kafka producer that buffers writes must Flush during drain. Otherwise messages in the buffer are lost.

defer producer.Close()
if err := producer.Flush(drainCtx); err != nil {
    log.Printf("flush: %v", err)
}

Closing the database before workers exit¶

Workers holding connections panic on use after close. Drain workers first.

Drain on a nil channel¶

Closing a nil channel panics. Defensive coding:

if p.jobs != nil {
    close(p.jobs)
}

But better: initialise it in the constructor.

Drain that does not propagate cancellation¶

A drain that closes a channel but does not cancel the context leaves workers stuck on long sleeps. Combine both signals.

Patterns That Often Help¶

Functional options¶

func WithDrainDeadline(d time.Duration) Option {
    return func(s *Server) { s.drainDeadline = d }
}

Per-environment drain budgets. Production gets 25s; tests get 100ms.

chan struct{} for signalling¶

A done channel is cheaper than a wait group when you only need "wait until x":

done := make(chan struct{})
go func() {
    work()
    close(done)
}()
<-done

sync.Once for single-shot drain¶

type Server struct {
    once sync.Once
}

func (s *Server) Drain(ctx context.Context) error {
    var err error
    s.once.Do(func() { err = s.drain(ctx) })
    return err
}

Idempotent drain: callable from any number of places.

Drain hooks¶

type Server struct {
    beforeDrain []func()
    afterDrain  []func()
}

func (s *Server) OnBeforeDrain(f func()) { s.beforeDrain = append(s.beforeDrain, f) }

Lets components register clean-up that runs at the right moment.

Patterns That Often Hurt¶

• Drain inside a handler. A handler calling srv.Shutdown deadlocks (the handler is holding the WaitGroup that Shutdown is waiting on).
• Goroutines that never exit. A for { ... } with no context check leaks past drain.
• Calls to os.Exit from goroutines. Skips drain entirely.
• A defer wg.Wait() in main. If Wait blocks forever, main never returns.
• Long time.Sleep instead of select with <-ctx.Done. Goroutine cannot drain on schedule.

Comparative Languages¶

The conceptual model is universal; only the API surface differs.

Mid-Level Cheat Sheet¶

1. Define Drainable interface; implement it on every long-lived component.
2. signal.NotifyContext at the top; errgroup for parallel run.
3. Per-component Drain(ctx) with deadline propagation.
4. Order drains by dependency: outer first, inner last.
5. Test drain with empty, in-flight, hung, and double-call cases.
6. Emit drain metrics; alert on force-cancellations.

type Drainable interface{ Drain(context.Context) error }

func Run(ctx context.Context, ds ...Drainable) error {
    <-ctx.Done()
    dctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
    defer cancel()
    for i := len(ds) - 1; i >= 0; i-- {
        if err := ds[i].Drain(dctx); err != nil {
            log.Printf("drain %T: %v", ds[i], err)
        }
    }
    return nil
}

Self-Assessment Checklist¶

• I can design a Drainable interface and apply it across components.
• I know when to use errgroup and how to wire it to a signal context.
• I can drain HTTP, gRPC, worker pool, and queue consumer.
• I can choose between idle-wait and wait-group based drains.
• I can express drain order via a Lifecycle registry.
• I can write five different kinds of drain test.
• I emit drain metrics and alert on force-cancellations.
• I avoid the dozen pitfalls listed above.

Summary¶

The middle level of drain is about composition. One drain is easy. Five drains in the right order with shared deadline and metrics — that takes real care. The patterns in this file (Drainable, Lifecycle, errgroup, per-phase context) are the building blocks of every well-shut-down Go service. Master them and you can write systems that deploy four times a day with zero downtime.

Further Reading¶

• The net/http package source for Server.Shutdown internals.
• The grpc-go source for GracefulStop.
• The golang.org/x/sync/errgroup documentation.
• Kubernetes documentation on terminationGracePeriodSeconds and preStop hooks.
• The segmentio/kafka-go Reader and Writer documentation, especially Close semantics.

Extended Example: A Full Mid-Sized Service¶

Below is a complete mid-sized service that exercises every concept in this file. It is roughly the shape of a real microservice: HTTP for synchronous requests, a worker pool for async jobs, a Kafka producer for events, a Postgres connection.

package main

import (
    "context"
    "database/sql"
    "errors"
    "fmt"
    "log"
    "net/http"
    "os"
    "os/signal"
    "sync"
    "sync/atomic"
    "syscall"
    "time"

    "golang.org/x/sync/errgroup"
    _ "github.com/lib/pq"
)

type Producer interface {
    Send(ctx context.Context, key string, value []byte) error
    Flush(ctx context.Context) error
    Close() error
}

type FakeProducer struct {
    buf   chan []byte
    wg    sync.WaitGroup
    closed atomic.Bool
}

func NewFakeProducer() *FakeProducer {
    p := &FakeProducer{buf: make(chan []byte, 1024)}
    p.wg.Add(1)
    go p.run()
    return p
}

func (p *FakeProducer) run() {
    defer p.wg.Done()
    for v := range p.buf {
        time.Sleep(5 * time.Millisecond)
        _ = v
    }
}

func (p *FakeProducer) Send(ctx context.Context, key string, value []byte) error {
    if p.closed.Load() {
        return errors.New("producer closed")
    }
    select {
    case p.buf <- value:
        return nil
    case <-ctx.Done():
        return ctx.Err()
    }
}

func (p *FakeProducer) Flush(ctx context.Context) error {
    for len(p.buf) > 0 {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-time.After(10 * time.Millisecond):
        }
    }
    return nil
}

func (p *FakeProducer) Close() error {
    if p.closed.CompareAndSwap(false, true) {
        close(p.buf)
    }
    p.wg.Wait()
    return nil
}

type Job struct{ ID int }

type Pool struct {
    jobs    chan Job
    wg      sync.WaitGroup
    closed  atomic.Bool
    mu      sync.Mutex
    prod    Producer
}

func NewPool(prod Producer, buf int) *Pool {
    return &Pool{jobs: make(chan Job, buf), prod: prod}
}

func (p *Pool) Start(ctx context.Context, n int) {
    for i := 0; i < n; i++ {
        p.wg.Add(1)
        go p.run(ctx, i)
    }
}

func (p *Pool) run(ctx context.Context, id int) {
    defer p.wg.Done()
    for {
        select {
        case <-ctx.Done():
            return
        case j, ok := <-p.jobs:
            if !ok {
                return
            }
            _ = p.prod.Send(ctx, fmt.Sprintf("job-%d", j.ID), []byte("done"))
        }
    }
}

func (p *Pool) Submit(j Job) error {
    if p.closed.Load() {
        return errors.New("pool closed")
    }
    p.mu.Lock()
    defer p.mu.Unlock()
    if p.closed.Load() {
        return errors.New("pool closed")
    }
    p.jobs <- j
    return nil
}

func (p *Pool) Drain(ctx context.Context) error {
    p.mu.Lock()
    if p.closed.CompareAndSwap(false, true) {
        close(p.jobs)
    }
    p.mu.Unlock()
    done := make(chan struct{})
    go func() { p.wg.Wait(); close(done) }()
    select {
    case <-done:
        return nil
    case <-ctx.Done():
        return ctx.Err()
    }
}

type App struct {
    pool     *Pool
    prod     Producer
    db       *sql.DB
    draining atomic.Bool
}

func (a *App) Handler() http.Handler {
    mux := http.NewServeMux()
    mux.HandleFunc("/ready", func(w http.ResponseWriter, r *http.Request) {
        if a.draining.Load() {
            w.WriteHeader(http.StatusServiceUnavailable)
            return
        }
        w.WriteHeader(http.StatusOK)
    })
    mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
        w.WriteHeader(http.StatusOK)
    })
    mux.HandleFunc("/enqueue", func(w http.ResponseWriter, r *http.Request) {
        if a.draining.Load() {
            http.Error(w, "draining", http.StatusServiceUnavailable)
            return
        }
        _ = a.pool.Submit(Job{ID: int(time.Now().Unix())})
        fmt.Fprintln(w, "enqueued")
    })
    return mux
}

func main() {
    ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
    defer cancel()

    prod := NewFakeProducer()
    pool := NewPool(prod, 64)
    app := &App{pool: pool, prod: prod}

    pool.Start(ctx, 4)

    srv := &http.Server{Addr: ":8080", Handler: app.Handler()}
    g, _ := errgroup.WithContext(ctx)
    g.Go(func() error {
        err := srv.ListenAndServe()
        if errors.Is(err, http.ErrServerClosed) {
            return nil
        }
        return err
    })

    <-ctx.Done()
    log.Println("drain: start")
    start := time.Now()

    app.draining.Store(true)
    time.Sleep(2 * time.Second)

    dctx, dcancel := context.WithTimeout(context.Background(), 25*time.Second)
    defer dcancel()

    if err := srv.Shutdown(dctx); err != nil {
        log.Printf("http: %v", err)
    }
    if err := pool.Drain(dctx); err != nil {
        log.Printf("pool: %v", err)
    }
    if err := prod.Flush(dctx); err != nil {
        log.Printf("flush: %v", err)
    }
    if err := prod.Close(); err != nil {
        log.Printf("prod close: %v", err)
    }
    _ = g.Wait()
    log.Printf("drain: complete in %s", time.Since(start))
}

Pay attention to the order in the drain block: HTTP first (no new requests), then pool (workers exit before producer is closed), then producer flush, then producer close. The wait at the end (g.Wait) catches any straggler error from the listener.

This template scales from "internal hackathon project" to "core production service" with only minor changes (real Kafka client, real Postgres, real metrics). The skeleton is correct.

Appendix A: Designing a Drain-Aware Library¶

If you publish a Go library that owns goroutines (a queue client, a metrics shipper, a websocket gateway), bake in drain from the start. A clean drain-aware library API looks like:

// Run starts the library's goroutines. Blocks until ctx is cancelled or
// an unrecoverable error occurs. Returns nil on graceful shutdown.
func (c *Client) Run(ctx context.Context) error

// Drain waits for in-flight operations to complete, bounded by ctx.
// Safe to call exactly once. After Drain returns, no further Run or
// Drain calls succeed.
func (c *Client) Drain(ctx context.Context) error

Document both functions thoroughly. State explicitly:

• The thread-safety of each method.
• Whether Drain is required after Run returns.
• What happens if Drain deadline is exceeded.
• Whether the library has its own internal deadlines beyond ctx.

Libraries that lie about drain ("call Close() and we'll figure it out") are the source of countless production incidents. Be precise.

Anti-example: a library that does not expose drain¶

type BadClient struct{ /* ... */ }
func NewBadClient() *BadClient { /* spawns goroutines internally */ }
func (b *BadClient) Stop() { /* unspecified what happens */ }

Stop is ambiguous: is it cooperative? Is it forced? Does it wait? No way to bound it. No way to detect a stuck shutdown. The library is unusable in production.

Better example¶

type GoodClient struct{ /* ... */ }
func NewGoodClient(opts ...Option) *GoodClient { /* no goroutines yet */ }
func (g *GoodClient) Start(ctx context.Context) error { /* spawns */ }
func (g *GoodClient) Drain(ctx context.Context) error { /* waits */ }

Two distinct lifecycle hooks. Both take context. Both have documented semantics. Caller is in control.

Appendix B: Drain Phases in Detail¶

Let's expand the phase model from earlier into a detailed reference.

Phase 0 — Pre-drain¶

This phase is optional but valuable. Run it before any blocking work.

• Log the start of drain with a timestamp and current in-flight counters.
• Emit a drain_started_total metric.
• Set a tracing span attribute so logs and traces correlate.
• Optionally, send an "I am shutting down" event to a control plane.

func (s *Server) preDrain() {
    s.logger.Info("drain started",
        "in_flight", s.inflight.Load(),
        "workers", s.workers,
    )
    metrics.DrainStarted.Inc()
}

Phase 1 — Stop intake¶

This phase flips the gate. After it returns, no new work enters the system.

• Set the draining flag.
• Wait for the readiness probe propagation window.
• Close the HTTP listener (or pause the queue consumer).

The wait for readiness is the most commonly skipped step. Without it, the LB keeps routing requests during drain and you cannot reach steady-state in-flight count of zero.

Phase 2 — Wait in-flight¶

This phase blocks until the system is idle.

• Wait on the in-flight WaitGroup or the in-flight counter.
• Bounded by the drain context's deadline.
• Return context.DeadlineExceeded if the deadline fires.

Most drain bugs hide here. Either the wait is unbounded, or the in-flight count is wrong (someone forgot a Done), or there is a goroutine that does not honour the context.

Phase 3 — Flush¶

Async producers and write-behind caches have internal buffers. This phase asks them to flush.

• Kafka producer flush.
• Metrics flush.
• Log shipper flush.

Each has its own deadline. Stack them inside the drain budget; if one runs over, the next gets less time. Track this with a "budget remaining" helper:

func remaining(ctx context.Context) time.Duration {
    d, ok := ctx.Deadline()
    if !ok {
        return time.Hour // pretend infinite
    }
    r := time.Until(d)
    if r < 0 {
        return 0
    }
    return r
}

Log the remaining budget at each phase boundary. It is gold during incident review.

Phase 4 — Close downstream¶

This is the last phase. Closing here is safe because nothing else is using the resources.

• Database connection pool close.
• Redis client close.
• File handles.
• OpenTelemetry exporters.

Order matters: drain consumers of the database before closing the database, drain Kafka producers before closing Kafka connections, etc.

Phase 5 — Post-drain¶

Optional cleanup before exit.

• Final flush of structured logs.
• Emit drain_complete_total metric.
• Return from main (which exits the process).

defer func() {
    log.Sync()
    log.Info("drain complete", "duration_ms", elapsed.Milliseconds())
}()

log.Sync flushes any buffered logs (some loggers buffer for performance).

Appendix C: A Tour of Mid-Level Anti-Patterns¶

Anti-pattern: done := make(chan bool)¶

A chan bool for signalling is a code smell. The value never carries information. Use chan struct{}:

done := make(chan struct{})
close(done) // signal
<-done       // wait

struct{} is zero-size; the compiler optimises better.

Anti-pattern: Blocking on a context outside a select¶

<-ctx.Done() // fine at the top of a function

This is OK. But inside a loop with other channels, always select:

select {
case <-ctx.Done():
    return
case j := <-jobs:
    process(j)
}

A bare <-ctx.Done() inside a loop is an unconditional exit point.

Anti-pattern: time.Sleep(d) to wait for drain¶

time.Sleep(5 * time.Second)

This does not wait for drain; it just adds latency. The system may have drained in 200ms, or still be working after 5s. Use WaitGroup.Wait with a deadline.

Anti-pattern: A drain that calls panic on error¶

func (s *Server) Drain(ctx context.Context) {
    if err := s.cleanup(); err != nil {
        panic(err)
    }
}