Drain Pattern — Junior Level¶

Table of Contents¶

1. Introduction
2. Prerequisites
3. Glossary
4. Core Concepts
5. Real-World Analogies
6. Mental Models
7. Pros & Cons
8. Use Cases
9. Code Examples
10. Coding Patterns
11. Clean Code
12. Product Use / Feature
13. Error Handling
14. Security Considerations
15. Performance Tips
16. Best Practices
17. Edge Cases & Pitfalls
18. Common Mistakes
19. Common Misconceptions
20. Tricky Points
21. Test
22. Tricky Questions
23. Cheat Sheet
24. Self-Assessment Checklist
25. Summary
26. What You Can Build
27. Further Reading
28. Related Topics
29. Diagrams & Visual Aids

Introduction¶

Focus: "What is draining? Why can't I just call os.Exit(0)? How do I make sure my service finishes its work before it stops?"

A drain is the controlled emptying of a system: water from a tub, fuel from a tank, requests from a server. In Go services, draining is the act of telling your program "do not accept new work" while letting the work already in progress finish. Then, and only then, do you actually exit.

The drain pattern is the opposite of a hard stop. A hard stop is what happens when you press Ctrl+C twice or call os.Exit(0): the operating system tears your process down immediately. Any goroutine that was mid-database-write gets cut off. Any HTTP response that was half-sent disappears. Any Kafka message your consumer was processing but had not committed gets re-delivered to a different consumer — possibly causing a duplicate side effect like a double charge on a credit card.

A graceful drain avoids all of that. It gives in-flight work a chance to land safely on disk, in the network, or in the next service. It is one of the most important production patterns in Go, and yet it is missed in almost every tutorial.

After reading this file you will:

• Know what draining means in plain English
• Understand the difference between hard stop and graceful drain
• Be able to drain a simple worker goroutine using a context and a sync.WaitGroup
• Know what a deadline-bounded drain is and why you need one
• Recognise the three pieces of every drain: stop intake, wait for in-flight, close downstream
• See your first integration with signal.Notify and os.Interrupt
• Understand why a drain is often the difference between a clean rolling deploy and a customer-facing incident

You do not need to know about Kafka rebalances, complex supervisor trees, or two-phase shutdown protocols yet. Those live in middle, senior, and professional. This file is about the smallest, most useful version of the drain pattern — the one you will reach for on day one of any real service.

Prerequisites¶

• Required: Go 1.21 or newer. Run go version to check.
• Required: Comfort with goroutines and channels at the basic level. You should know what go f() does and how to send and receive on a channel.
• Required: Awareness of context.Context. You do not have to know every method; knowing that a context can be cancelled is enough.
• Required: A main() function and the ability to run go run main.go.
• Helpful: Familiarity with sync.WaitGroup. We will use it in every example.
• Helpful: A vague memory of os.Signal from operating systems class. We will explain SIGTERM and SIGINT from scratch.

If you can write a program that spawns three goroutines, sends them work on a channel, and waits for them with a WaitGroup, you are ready.

Glossary¶

Core Concepts¶

Drain is three steps, always in this order¶

Every drain, no matter how complex, breaks down into the same three steps:

1. Stop accepting new work. Close the inbound channel, return 503 from the HTTP listener, pause the queue consumer. New requests are politely refused or queued upstream for the next pod.
2. Wait for in-flight work to finish. The workers that already have a task on their plate continue. You block until they are all done.
3. Close downstream resources. Database connections, Kafka producers, file handles. Only after every worker has reported "no more work in my hands."

Skip step 1 and you keep adding work even as you try to stop. Skip step 2 and you cut off live requests. Skip step 3 and you leak file descriptors and corrupt state. The order matters.

Drain needs a deadline¶

If you wait forever for in-flight work, one hung worker will keep your pod alive forever. Kubernetes will eventually run out of patience and send SIGKILL — and now you are back to a hard stop, except you also missed the chance to flush whatever you could have flushed.

The fix is a deadline-bounded drain: "give workers up to 25 seconds to finish; after that, cancel their contexts and force them to exit." The deadline is shorter than the platform's grace period (30 seconds in Kubernetes) so you have time to clean up resources before the kernel arrives.

Drain is not just for shutdown¶

The drain pattern shows up in more places than you might think:

• Rolling deploys — when a new version of your pod is starting, the old one drains before exiting.
• Kafka rebalances — when partitions are reassigned, a consumer drains the in-flight messages from the partitions it is about to lose.
• Worker pool resize — shrinking a pool drains the workers being removed.
• Database failover — connections to the old primary drain before being routed to the new primary.

The mechanics are similar in every case: stop intake, wait, close.

Hard stop is sometimes correct¶

The drain pattern is the default for stateful or side-effecting services. But there are times when a hard stop is the right call:

• Read-only services with idempotent retries. If a load balancer will retry on the next pod, dropping the in-flight request costs almost nothing.
• Tests. A test process exiting on os.Exit(1) after a failure is normal.
• Detected corruption. If you discover your own state is bad, stopping fast prevents the corruption from spreading.

Knowing when not to drain is part of using the pattern well.

Real-World Analogies¶

A coffee shop closing for the night¶

It is 9:55 PM. The shop closes at 10:00. The barista turns the "Open" sign to "Closed" — no new customers walk in. Three customers are already at the counter; the barista finishes their drinks. At 10:05 the last drink is handed over. The barista locks the door, turns off the espresso machine, and goes home.

That is a drain. The "Closed" sign is step 1. Finishing the three drinks is step 2. Locking the door and turning off the machine is step 3. The 10-minute deadline (call it the grace period) is the bound. If a customer's order was somehow taking 30 minutes, the barista would still leave at 10:05 — that is the deadline kicking in.

A hard stop would be: at 9:55 the barista drops three half-made drinks on the counter, walks out, locks the door. The customers are upset, the espresso is wasted, the cups have to be thrown out tomorrow.

A factory line ending a shift¶

The conveyor belt feeding raw parts is stopped first. The line workers keep assembling whatever is already on the belt. Once the belt is empty and every part has been assembled or set aside, the workers turn off the machines and leave.

The conveyor belt is your inbound queue. The workers are your goroutines. Turning off the belt is closing your channel. Letting them finish the parts on the belt is the drain wait. Turning off the machines is closing your database connections.

An elevator at end of service¶

A maintenance elevator going out of service at midnight does not just stop. It refuses new floor-button presses (no new work) and then finishes the trip it is on (in-flight work). Only at the lobby does it park itself and the doors stay open. If somebody had been frantically pressing buttons for a thirty-third floor, that request is cancelled at the deadline — the elevator parks and the technician walks them out.

Mental Models¶

The funnel model¶

Think of your service as a funnel. Requests pour in the top, flow through the goroutines in the middle, and exit at the bottom into a database or another service. A drain is closing the top of the funnel — no new pour — and waiting for the funnel to empty from the bottom up. The drain is over when the funnel is dry.

The bucket-brigade model¶

Imagine a chain of people passing buckets of water. When the brigade ends, the first person stops scooping — no new buckets enter the chain. Each person passes whatever bucket they are holding to the next. Once the buckets reach the end and there is nothing left in anyone's hands, the brigade is over.

This maps well onto pipelined services: HTTP server → worker pool → database. The first stage drains first; each subsequent stage drains as its upstream empties.

The two-clock model¶

Every drain has two clocks running.

• Worker clock — how long it actually takes the workers to finish.
• Deadline clock — how long you are willing to wait.

A drain ends when either clock reaches zero. If the worker clock reaches zero first, you did a "clean drain" — every task finished. If the deadline clock reaches zero first, you did a "forced drain" — some tasks were cancelled. Production drains aim to almost always be clean drains, but a forced drain is still better than an unbounded wait followed by SIGKILL.

Pros & Cons¶

Pros¶

• Zero data loss on rolling deploys. Every accepted request gets a response. Every accepted Kafka message gets committed.
• Predictable shutdown time. With a deadline, you know the maximum cost of a drain.
• Customer-friendly. Users do not see a 502 in the middle of their checkout because your old pod went away.
• Plays well with orchestrators. Kubernetes, Nomad, ECS all assume SIGTERM will be handled this way.
• Compose-friendly. Drains chain naturally: server drains, then pool drains, then DB closes.

Cons¶

• More code. A non-draining service is shorter. The drain adds a context, a wait group, a signal handler, a deadline.
• Race-prone if done casually. Closing a channel that someone is still sending to is a panic. Order of operations matters.
• Untested in most codebases. Most teams never simulate a drain in CI; the first real test is a production deploy.
• Hides slow shutdowns. If your drain takes 25 seconds every deploy and nobody notices, you have a leak you do not see.

Use Cases¶

• HTTP API servers. Drain before exit so in-flight responses are completed.
• Kafka consumers. Commit offsets for fully processed messages before releasing partitions.
• Worker pools. Let queued jobs in the channel be picked up and processed before the pool dies.
• Cron-like job runners. A long-running job that took the lock should finish before the next pod takes over.
• WebSocket gateways. Close existing sessions politely, refuse new ones.
• gRPC servers. GracefulStop is the gRPC drain primitive.
• Background batch flushers. A buffer that flushes every 5 seconds should flush one last time on shutdown.

Code Examples¶

Example 1: The smallest possible drain¶

A worker reads from a channel and processes items. We want it to stop accepting new items and finish what is in the channel before exiting.

package main

import (
    "fmt"
    "sync"
    "time"
)

func worker(id int, in <-chan int, wg *sync.WaitGroup) {
    defer wg.Done()
    for item := range in {
        time.Sleep(100 * time.Millisecond) // simulate work
        fmt.Printf("worker %d done with %d\n", id, item)
    }
    fmt.Printf("worker %d exiting\n", id)
}

func main() {
    in := make(chan int, 10)
    var wg sync.WaitGroup

    for i := 1; i <= 3; i++ {
        wg.Add(1)
        go worker(i, in, &wg)
    }

    for i := 1; i <= 5; i++ {
        in <- i
    }

    // Drain step 1: stop accepting new work.
    close(in)

    // Drain step 2: wait for in-flight work.
    wg.Wait()

    // Drain step 3: nothing to close in this toy example.
    fmt.Println("clean shutdown")
}

Closing in tells every worker "no more items, exit when the channel is drained." wg.Wait() blocks until all three workers call wg.Done. The output ends with clean shutdown and the process exits with code 0.

Example 2: Adding a context for cancellation¶

Closing a channel is fine for "no more work coming," but sometimes you want to also tell workers "stop what you are doing." For that we add a context.Context.

package main

import (
    "context"
    "fmt"
    "sync"
    "time"
)

func worker(ctx context.Context, id int, in <-chan int, wg *sync.WaitGroup) {
    defer wg.Done()
    for {
        select {
        case <-ctx.Done():
            fmt.Printf("worker %d cancelled\n", id)
            return
        case item, ok := <-in:
            if !ok {
                fmt.Printf("worker %d done\n", id)
                return
            }
            time.Sleep(50 * time.Millisecond)
            fmt.Printf("worker %d processed %d\n", id, item)
        }
    }
}

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    in := make(chan int, 5)
    var wg sync.WaitGroup

    for i := 1; i <= 3; i++ {
        wg.Add(1)
        go worker(ctx, i, in, &wg)
    }

    for i := 1; i <= 4; i++ {
        in <- i
    }
    close(in)

    // Wait for natural drain.
    wg.Wait()
    cancel() // tidy up the context

    fmt.Println("done")
}

Here the workers exit either because the channel closes (clean drain) or because the context is cancelled (forced drain). On a normal shutdown the channel close wins; on an emergency the cancel wins.

Example 3: Drain with a deadline¶

A real service caps the drain time. If workers do not finish in 5 seconds, we force-cancel and exit anyway.

package main

import (
    "context"
    "fmt"
    "sync"
    "time"
)

func slowWorker(ctx context.Context, id int, in <-chan int, wg *sync.WaitGroup) {
    defer wg.Done()
    for {
        select {
        case <-ctx.Done():
            fmt.Printf("worker %d cancelled at deadline\n", id)
            return
        case item, ok := <-in:
            if !ok {
                return
            }
            select {
            case <-time.After(2 * time.Second):
                fmt.Printf("worker %d finished %d\n", id, item)
            case <-ctx.Done():
                fmt.Printf("worker %d gave up on %d\n", id, item)
                return
            }
        }
    }
}

func main() {
    in := make(chan int, 10)
    for i := 1; i <= 8; i++ {
        in <- i
    }
    close(in)

    drainCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()

    var wg sync.WaitGroup
    for i := 1; i <= 2; i++ {
        wg.Add(1)
        go slowWorker(drainCtx, i, in, &wg)
    }

    done := make(chan struct{})
    go func() {
        wg.Wait()
        close(done)
    }()

    select {
    case <-done:
        fmt.Println("clean drain")
    case <-drainCtx.Done():
        fmt.Println("drain deadline exceeded, forcing exit")
        <-done // still wait for goroutines to actually return
    }
}

This is the canonical deadline-bounded drain. Notice that even after the deadline fires we still wait on done. We do not abandon the goroutines; we let them see the cancelled context and exit. Otherwise we would leak.

Example 4: Drain on SIGTERM¶

In a real service the drain is triggered not by main choosing to stop, but by an external signal.

package main

import (
    "context"
    "fmt"
    "os"
    "os/signal"
    "sync"
    "syscall"
    "time"
)

func runWorker(ctx context.Context, in <-chan int, wg *sync.WaitGroup) {
    defer wg.Done()
    for {
        select {
        case <-ctx.Done():
            return
        case item, ok := <-in:
            if !ok {
                return
            }
            time.Sleep(200 * time.Millisecond)
            fmt.Printf("processed %d\n", item)
        }
    }
}

func main() {
    in := make(chan int, 100)

    stop := make(chan os.Signal, 1)
    signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

    rootCtx, cancelRoot := context.WithCancel(context.Background())
    defer cancelRoot()

    var wg sync.WaitGroup
    for i := 0; i < 3; i++ {
        wg.Add(1)
        go runWorker(rootCtx, in, &wg)
    }

    // Producer goroutine.
    go func() {
        i := 0
        for {
            select {
            case <-rootCtx.Done():
                close(in)
                return
            case in <- i:
                i++
                time.Sleep(50 * time.Millisecond)
            }
        }
    }()

    <-stop
    fmt.Println("signal received, draining")

    cancelRoot()

    drained := make(chan struct{})
    go func() { wg.Wait(); close(drained) }()

    select {
    case <-drained:
        fmt.Println("clean drain")
    case <-time.After(10 * time.Second):
        fmt.Println("drain timed out")
    }
}

Press Ctrl+C and the producer stops, the channel closes, the workers finish what is in their hand, and the program exits.

Example 5: HTTP server drain¶

net/http has a built-in drain via Server.Shutdown. It stops accepting new connections and waits for in-flight requests to finish, up to the deadline of the context you pass.

package main

import (
    "context"
    "fmt"
    "net/http"
    "os"
    "os/signal"
    "syscall"
    "time"
)

func main() {
    mux := http.NewServeMux()
    mux.HandleFunc("/slow", func(w http.ResponseWriter, r *http.Request) {
        time.Sleep(3 * time.Second)
        fmt.Fprintln(w, "done")
    })

    srv := &http.Server{Addr: ":8080", Handler: mux}

    go func() {
        if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
            fmt.Println("server error:", err)
        }
    }()

    stop := make(chan os.Signal, 1)
    signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
    <-stop
    fmt.Println("draining HTTP server")

    ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
    defer cancel()
    if err := srv.Shutdown(ctx); err != nil {
        fmt.Println("shutdown error:", err)
    }
    fmt.Println("server stopped")
}

If /slow is mid-request when the signal arrives, the server lets it finish (3 seconds) before exiting. New requests get connection refused because the listener is closed. After 10 seconds, anything still hanging is force-closed.

Example 6: Drain with idle-wait¶

Sometimes there is no neat channel to close — work arrives unpredictably and you must wait until the system has been idle for a while.

package main

import (
    "fmt"
    "sync/atomic"
    "time"
)

type Service struct {
    inFlight atomic.Int64
}

func (s *Service) Handle() {
    s.inFlight.Add(1)
    defer s.inFlight.Add(-1)
    time.Sleep(100 * time.Millisecond)
}

func (s *Service) Drain(idle, deadline time.Duration) {
    end := time.Now().Add(deadline)
    for time.Now().Before(end) {
        if s.inFlight.Load() == 0 {
            time.Sleep(idle)
            if s.inFlight.Load() == 0 {
                return
            }
        }
        time.Sleep(50 * time.Millisecond)
    }
}

func main() {
    s := &Service{}
    for i := 0; i < 5; i++ {
        go s.Handle()
    }
    s.Drain(200*time.Millisecond, 5*time.Second)
    fmt.Println("drained")
}

The idle-wait pattern is "be idle for at least 200ms before declaring drained." It is useful for systems where new work can arrive even after intake is supposedly closed (think: in-process retries).

Example 7: Refusing new work during drain¶

A handler that returns 503 once the service is draining.

package main

import (
    "net/http"
    "sync/atomic"
)

type App struct {
    draining atomic.Bool
}

func (a *App) Middleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if a.draining.Load() {
            http.Error(w, "shutting down", http.StatusServiceUnavailable)
            return
        }
        next.ServeHTTP(w, r)
    })
}

func (a *App) StartDrain() {
    a.draining.Store(true)
}

StartDrain is called as the first action of your shutdown handler — before you even call Server.Shutdown. This is important if your load balancer keeps sending requests for a moment after SIGTERM (most do, briefly).

Example 8: Tracking in-flight count by hand¶

The pattern from example 6 wrapped into a reusable type.

package main

import (
    "context"
    "errors"
    "sync"
    "sync/atomic"
    "time"
)

type Drainer struct {
    inFlight atomic.Int64
    closed   atomic.Bool
    mu       sync.Mutex
}

func (d *Drainer) Enter() error {
    if d.closed.Load() {
        return errors.New("draining")
    }
    d.inFlight.Add(1)
    if d.closed.Load() {
        d.inFlight.Add(-1)
        return errors.New("draining")
    }
    return nil
}

func (d *Drainer) Exit() {
    d.inFlight.Add(-1)
}

func (d *Drainer) Drain(ctx context.Context) error {
    d.closed.Store(true)
    for d.inFlight.Load() > 0 {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-time.After(20 * time.Millisecond):
        }
    }
    return nil
}

The Enter/Exit pair is the bookend of every request. Drain closes intake and polls until the counter reaches zero.

Coding Patterns¶

Pattern: the shutdown function returns¶

In Go, the idiomatic way to express drain is "the function that started the workers also stops them." A Run method takes a context and returns when everything has drained:

func (s *Server) Run(ctx context.Context) error {
    // start workers ...
    <-ctx.Done()
    // drain ...
    return nil
}

The caller passes a cancellable context. Cancelling it triggers the drain.

Pattern: errgroup for drain¶

golang.org/x/sync/errgroup lets you wait on a set of goroutines and capture the first error:

import "golang.org/x/sync/errgroup"

g, ctx := errgroup.WithContext(ctx)
g.Go(func() error { return server.Run(ctx) })
g.Go(func() error { return consumer.Run(ctx) })
if err := g.Wait(); err != nil { /* one died, others see ctx cancellation */ }

When any goroutine returns an error, the group context is cancelled and the others see it as a drain signal.

Pattern: split shutdown into phases¶

Phase 1: refuse new work. Phase 2: wait for in-flight. Phase 3: close downstream. Implement each as a method:

func (s *Server) Drain(ctx context.Context) error {
    s.StopIntake()
    if err := s.WaitInFlight(ctx); err != nil { return err }
    return s.CloseDownstream()
}

This is straightforward to test, easy to reason about, and easy to extend.

Clean Code¶

• Name the context. drainCtx, shutdownCtx — never just ctx if the context's job is to bound a drain.
• Always defer the cancel. Even after a clean drain, cancel() should run. defer cancel() right after WithTimeout.
• Keep main shallow. main wires signal handling and starts components. The drain logic lives in the components.
• Pass context, do not store it. A struct that stores a context will keep that context alive longer than expected. Pass it through Run/Drain parameters.
• One drain per component. Do not let the HTTP server reach into the worker pool to drain it. Each component drains itself and exposes a Drain(ctx) error.
• Log timings. A log.Printf("drain took %s", time.Since(start)) after every drain is gold during incident reviews.

Product Use / Feature¶

Imagine a checkout service that emits a payment.completed event to Kafka. Without drain: a rolling deploy interrupts a checkout half-way; the payment is captured but the event is never sent. The downstream "send a receipt email" service never runs. The customer is charged and gets no email.

With drain: the pod receives SIGTERM, marks itself draining (so the LB stops routing new checkouts to it), finishes the three in-flight checkouts (including emitting their events), and exits. Zero stuck payments, zero missing emails, no Sunday-evening pager.

The drain pattern is what makes "we deploy four times a day" boring.

Error Handling¶

A drain has its own family of errors:

• context.DeadlineExceeded — drain ran out of time. Common, expected, worth logging with the number of in-flight items that did not finish.
• context.Canceled — drain context was cancelled before completion. Usually means a second signal arrived (operator pressed Ctrl+C twice) and you should now exit fast.
• net.ErrClosed — listener already closed. Often safe to ignore; surface only at debug level.
• Worker errors after drain start — drop into a side-buffer and log. They are post-mortem evidence, not actionable.

Wrap them with fmt.Errorf("drain server: %w", err) so a top-level log shows the chain.

Security Considerations¶

• Do not skip drain to recover from a security incident. If you suspect compromise, hard stop is correct — drain might exfiltrate.
• Do not log message bodies during drain. A panic log on a half-processed message can leak PII into logs.
• Health-check spoofing. A draining pod that still says "healthy" can be picked by a malicious caller to keep it alive. Always flip the health endpoint to "draining" before the work signal.

Performance Tips¶

• Drain time is on the deploy critical path. A 30-second drain across 50 pods adds up.
• Idle-wait with a long poll is cheaper than busy-wait. 20–50ms sleep between checks is fine.
• Use sync.WaitGroup over atomic counters where possible. Wait blocks; you do not have to poll.
• Avoid global locks during drain. A worker waiting for a mutex held by a drainer is the recipe for a missed deadline.
• Pre-size buffered channels. If a channel is the intake buffer, do not let it grow forever during drain — that just adds work.

Best Practices¶

• Always handle SIGTERM and SIGINT in production binaries.
• Always cap drain time with a context deadline.
• Always flip a "draining" flag before stopping the listener — so health checks turn unhealthy first.
• Always log start and end of drain with timings.
• Always test drain with kill -TERM in your dev environment, not just Ctrl+C.
• Never call os.Exit from a library — only main decides exit.
• Never use a goroutine pool without an explicit Stop/Drain method.

Edge Cases & Pitfalls¶

• Sending on a closed channel. Panics. The producer must check the drain flag before sending.
• Double close of a channel. Panics. Use sync.Once or a "done" channel to guarantee single close.
• A worker that re-enqueues to itself. During drain, the channel may already be closed. Either drain before resignaling or have a "stop accepting" gate.
• select with a default case in a worker. If every case is blocking and default fires immediately, your "drain wait" becomes a CPU spin.
• time.After in a hot loop. Creates a timer every iteration. Use time.NewTimer and reset it.
• A nil channel in a select. Blocks forever. Useful trick for disabling a case — but easy to do by accident during drain.

Common Mistakes¶

1. Forgetting to call cancel. Even if context auto-expires, calling cancel releases resources earlier.
2. Calling Shutdown from inside a handler. A handler running on the server cannot wait for the server to stop. Trigger drain from outside.
3. Logging on every poll iteration. A drain that polls every 20ms and logs every poll produces 50 log lines per second of drain.
4. Closing the inbound channel from a goroutine that is itself an inbound producer. Race against other producers.
5. Using select without a <-ctx.Done() case in workers. No way to force exit on deadline.
6. Drain timeout shorter than the slowest operation. A 1-second drain on a service whose median handler is 800ms drops 30% of requests every deploy.

Common Misconceptions¶

• "os.Exit is fine, the OS cleans up." It does not flush your application buffers, commit your Kafka offsets, or send your final HTTP response.
• "defer runs on SIGTERM." Only if you handle the signal yourself. Without signal.Notify, your program is killed and no defers run.
• "A graceful shutdown means no errors." A drain can still emit errors (deadline exceeded, partial flush). Graceful means "as good as we could do," not "perfect."
• "Buffered channels are enough for drain." Buffering helps absorb spikes; it does not by itself implement drain. You still need close + wait.
• "WaitGroup waits forever." It waits until the counter is zero. If you never call Done, it waits forever — that is a bug, not a feature.

Tricky Points¶

Wait on a WaitGroup whose count was never incremented¶

Returns immediately. That is sometimes desirable (no workers spawned) but can hide a bug where you forgot Add. Always pair wg.Add(1) and defer wg.Done() tightly.

Closing a channel from a producer that may still have buffered sends¶

close is safe even if there are buffered items; the receiver can still read them. What is not safe is sending after close. Producers must check a stop signal before sending.

context.Background versus a fresh context for the drain¶

For the root of a server, context.Background is correct. For the drain phase, context.WithTimeout(context.Background(), drainBudget) — note we do not derive from the cancelled service context, because that context is already cancelled and the drain would have zero time.

signal.Notify with an unbuffered channel¶

If the signal arrives before your code reaches <-stop, the runtime drops it. Always pass a buffered channel (make(chan os.Signal, 1)).

Health endpoints during drain¶

Most load balancers do not stop sending traffic the instant you call Shutdown. There is a propagation delay. Flip the health endpoint to "draining" first, sleep for the LB's check interval (often 5 seconds), then start the server drain.

Test¶

package drain_test

import (
    "context"
    "sync/atomic"
    "testing"
    "time"
)

type Counter struct{ n atomic.Int64 }

func (c *Counter) Add()  { c.n.Add(1) }
func (c *Counter) Done() { c.n.Add(-1) }
func (c *Counter) Load() int64 { return c.n.Load() }

func TestDrainWaitsForInFlight(t *testing.T) {
    c := &Counter{}
    c.Add()

    ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
    defer cancel()

    go func() {
        time.Sleep(50 * time.Millisecond)
        c.Done()
    }()

    for c.Load() > 0 {
        select {
        case <-ctx.Done():
            t.Fatal("drain deadline exceeded")
        case <-time.After(10 * time.Millisecond):
        }
    }
}

func TestDrainHonoursDeadline(t *testing.T) {
    c := &Counter{}
    c.Add()

    ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
    defer cancel()

    done := false
    for c.Load() > 0 && !done {
        select {
        case <-ctx.Done():
            done = true
        case <-time.After(5 * time.Millisecond):
        }
    }
    if !done {
        t.Fatal("expected deadline to fire")
    }
}

Run with go test -race ./.... The first test proves drain waits for in-flight to clear; the second proves the deadline actually bounds the wait.

Tricky Questions¶

Q1. What is the difference between close(ch) and cancelling a context? A. close(ch) signals "no more values," and receivers see ok=false. Cancelling a context signals "stop what you are doing," and goroutines see <-ctx.Done(). Drain often uses both: close to say "done with input," cancel to say "abort current work."

Q2. Why is a deadline necessary? A. Because some work never finishes. Without a bound, a hung handler can keep your pod alive past the grace period — and the kernel will SIGKILL it anyway, undoing any progress you had made on the drain.

Q3. What happens if the LB still sends traffic after Shutdown is called? A. Those connections may be accepted briefly (depending on socket state) but the server will not start new request goroutines after Shutdown is called. The LB will get connection refused on subsequent attempts. To smooth the transition, mark the service unhealthy first, wait a moment, then call Shutdown.

Q4. Should I drain before or after closing my database? A. Drain first, close database after. The drain is the consumer of the database; closing it first cancels the in-flight work.

Q5. What is the simplest way to detect a goroutine leak post-drain? A. Call runtime.NumGoroutine before and after. If the post-drain number is much higher than the baseline, you have a leak.

Cheat Sheet¶

DRAIN = STOP_INTAKE -> WAIT_INFLIGHT -> CLOSE_DOWNSTREAM
Always bound the wait with a deadline.
SIGTERM -> drain. SIGKILL -> too late.
Health endpoint flips first, listener stops second.
Test with kill -TERM in dev. Verify with /metrics in prod.

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

go server.Run(ctx)
<-ctx.Done()

drainCtx, drainCancel := context.WithTimeout(context.Background(), 25*time.Second)
defer drainCancel()
_ = server.Drain(drainCtx)

Self-Assessment Checklist¶

• I can explain the three steps of a drain in one sentence each.
• I can write a worker that exits cleanly on a closed channel.
• I can write a worker that exits on a cancelled context.
• I can combine both in a select block.
• I can wire signal.Notify to a context and use it to trigger drain.
• I can bound a drain with context.WithTimeout.
• I can use http.Server.Shutdown and explain what it does.
• I can describe the difference between SIGTERM, SIGINT, and SIGKILL.
• I can identify the most common pitfall: sending on a closed channel.

Summary¶

The drain pattern is the discipline of stopping a service in three steps: refuse new work, wait for in-flight work, close downstream resources. Every step is bounded by a deadline so a hung worker cannot block the whole shutdown. The pattern is built on three Go primitives: context for cancellation and deadlines, sync.WaitGroup for tracking in-flight goroutines, and signal.Notify for catching the OS shutdown signal. With those three pieces and a clear sense of order, you can shut down any Go service cleanly enough to roll out new versions four times a day without paging anyone.

What You Can Build¶

• A worker pool with Start, Submit, and Drain(ctx) methods.
• An HTTP API server that returns 200 from /healthz until it starts draining, then 503.
• A CLI tool that processes a list of files and supports Ctrl+C to "finish the current file, skip the rest."
• A batch job runner that flushes its buffer on exit.
• A simple message-queue consumer (Redis list, NATS, RabbitMQ) that drains on signal.

Further Reading¶

• The net/http package documentation for Server.Shutdown.
• The Go blog post "Go Concurrency Patterns: Context."
• The golang.org/x/sync/errgroup package documentation.
• Kubernetes documentation on pod lifecycle, especially the termination grace period.
• The signal package documentation, especially signal.NotifyContext.

Related Topics¶

• Graceful Shutdown
• Context Cancellation
• Worker Pools
• Signal Handling
• Kafka Consumer Patterns

Diagrams & Visual Aids¶

                   ┌─────────────┐
                   │   SIGTERM    │
                   └──────┬──────┘
                          │
                          ▼
                ┌───────────────────┐
                │ 1. Stop Intake     │  (close listener / channel / flag)
                └─────────┬─────────┘
                          │
                          ▼
                ┌───────────────────┐
                │ 2. Wait In-Flight  │  (WaitGroup.Wait, bounded by deadline)
                └─────────┬─────────┘
                          │
                          ▼
                ┌───────────────────┐
                │ 3. Close Downstream│  (DB, Kafka, files)
                └─────────┬─────────┘
                          │
                          ▼
                       exit 0

Time ──────────────────────────────────────────────►

intake  ████████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░  (stops at SIGTERM)
work-1  ░░░░████████░░░░░░░░░░░░░░░░░░░░░░░░░░░░
work-2  ░░░░░░██████████░░░░░░░░░░░░░░░░░░░░░░░░
work-3  ░░░░░░░░░░██████████████░░░░░░░░░░░░░░░░  (still running)

SIGTERM ▲                  ▲ deadline
        │                  │
        │←──── drain ──────│

The first diagram is the pipeline of phases. The second is a time-series view: intake stops the instant the signal arrives; individual workers finish at their own pace; the deadline puts a hard cap on how long we wait.

Extended Walkthrough: Building Drain From Scratch¶

To really internalise the drain pattern, let us build it from nothing. We start with a service that does not handle shutdown at all. Then we add each piece, one at a time, watching what happens.

Step 0 — A service that ignores shutdown¶

package main

import (
    "fmt"
    "time"
)

func main() {
    jobs := make(chan int, 5)

    go func() {
        for j := range jobs {
            time.Sleep(500 * time.Millisecond)
            fmt.Printf("processed %d\n", j)
        }
    }()

    for i := 1; ; i++ {
        jobs <- i
        time.Sleep(100 * time.Millisecond)
    }
}

Run this and press Ctrl+C. The output stops mid-sentence. The worker may have been mid-Sleep when the program died. No defers ran. The channel is leaked, the goroutine is leaked, the OS reclaims memory but no application-level cleanup happened.

This is the baseline. We will not ship this.

Step 1 — Catch the signal¶

The smallest change: tell the runtime "I want to know about SIGINT and SIGTERM."

package main

import (
    "fmt"
    "os"
    "os/signal"
    "syscall"
    "time"
)

func main() {
    jobs := make(chan int, 5)

    go func() {
        for j := range jobs {
            time.Sleep(500 * time.Millisecond)
            fmt.Printf("processed %d\n", j)
        }
    }()

    stop := make(chan os.Signal, 1)
    signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

    go func() {
        for i := 1; ; i++ {
            jobs <- i
            time.Sleep(100 * time.Millisecond)
        }
    }()

    <-stop
    fmt.Println("got signal")
}

Now Ctrl+C prints got signal and exits. But the worker did not finish what it was doing. The drain is still missing.

Step 2 — Tell the worker to stop¶

We need a way to say "no more work." The classic way is to close the input channel.

package main

import (
    "fmt"
    "os"
    "os/signal"
    "sync"
    "syscall"
    "time"
)

func main() {
    jobs := make(chan int, 5)
    var wg sync.WaitGroup

    wg.Add(1)
    go func() {
        defer wg.Done()
        for j := range jobs {
            time.Sleep(500 * time.Millisecond)
            fmt.Printf("processed %d\n", j)
        }
        fmt.Println("worker exit")
    }()

    stop := make(chan os.Signal, 1)
    signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

    done := make(chan struct{})
    go func() {
        i := 1
        for {
            select {
            case <-done:
                return
            case jobs <- i:
                i++
                time.Sleep(100 * time.Millisecond)
            }
        }
    }()

    <-stop
    fmt.Println("draining")
    close(done)
    close(jobs)
    wg.Wait()
    fmt.Println("clean exit")
}

Now the sequence is: signal arrives, we tell the producer to stop, we close the jobs channel, the worker drains what is left in the buffer, prints worker exit, the wait group unblocks, and we print clean exit.

This is already a working drain. Read it twice. The pieces are: a stop channel, a producer that checks the stop channel before each send, a closeable input channel, a wait group around the worker, and a final Wait in main.

Step 3 — Add a deadline¶

What if the worker is hung? We bound the wait.

done := make(chan struct{})
go func() { wg.Wait(); close(done) }()

select {
case <-done:
    fmt.Println("clean drain")
case <-time.After(5 * time.Second):
    fmt.Println("drain timed out; some work was lost")
}

We do not block on wg.Wait directly; we wrap it in a goroutine that closes a channel when done, then select on that against a timer. This is a Go idiom you will use many times.

Step 4 — Add a cancellation context¶

If the worker is doing a long operation (say, an HTTP call to a slow backend), we want a way to cancel that operation, not just signal "no more work." Enter context.

package main

import (
    "context"
    "fmt"
    "os"
    "os/signal"
    "sync"
    "syscall"
    "time"
)

func worker(ctx context.Context, jobs <-chan int, wg *sync.WaitGroup) {
    defer wg.Done()
    for {
        select {
        case <-ctx.Done():
            fmt.Println("worker cancelled")
            return
        case j, ok := <-jobs:
            if !ok {
                fmt.Println("jobs closed")
                return
            }
            select {
            case <-time.After(500 * time.Millisecond):
                fmt.Printf("processed %d\n", j)
            case <-ctx.Done():
                fmt.Printf("abandoned %d\n", j)
                return
            }
        }
    }
}

func main() {
    jobs := make(chan int, 10)
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    var wg sync.WaitGroup
    wg.Add(1)
    go worker(ctx, jobs, &wg)

    go func() {
        for i := 1; ; i++ {
            select {
            case <-ctx.Done():
                close(jobs)
                return
            case jobs <- i:
                time.Sleep(80 * time.Millisecond)
            }
        }
    }()

    stop := make(chan os.Signal, 1)
    signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
    <-stop
    fmt.Println("draining")

    drainCtx, drainCancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer drainCancel()

    done := make(chan struct{})
    go func() { wg.Wait(); close(done) }()

    select {
    case <-done:
        fmt.Println("clean drain")
    case <-drainCtx.Done():
        fmt.Println("force cancelling")
        cancel()
        <-done
    }
}

Look carefully at the worker. It has two select blocks. The outer one waits for either a job or cancellation. The inner one waits for the per-job timer or cancellation. Without the inner block, a long job would not be interruptible. This nested select is one of the most useful patterns in Go.

Step 5 — Refactor into a service¶

The pattern above is correct but messy in main. Let us extract it.

package main

import (
    "context"
    "fmt"
    "os"
    "os/signal"
    "sync"
    "syscall"
    "time"
)

type Service struct {
    jobs chan int
    wg   sync.WaitGroup
}

func New() *Service {
    return &Service{jobs: make(chan int, 16)}
}

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

func (s *Service) Submit(j int) bool {
    select {
    case s.jobs <- j:
        return true
    default:
        return false
    }
}

func (s *Service) Drain(ctx context.Context) error {
    close(s.jobs)
    done := make(chan struct{})
    go func() { s.wg.Wait(); close(done) }()
    select {
    case <-done:
        return nil
    case <-ctx.Done():
        return ctx.Err()
    }
}

func (s *Service) run(ctx context.Context, id int) {
    defer s.wg.Done()
    for {
        select {
        case <-ctx.Done():
            return
        case j, ok := <-s.jobs:
            if !ok {
                return
            }
            select {
            case <-time.After(200 * time.Millisecond):
                fmt.Printf("worker %d done %d\n", id, j)
            case <-ctx.Done():
                return
            }
        }
    }
}

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

    svc := New()
    svc.Start(ctx, 3)

    go func() {
        for i := 1; ; i++ {
            if !svc.Submit(i) {
                time.Sleep(10 * time.Millisecond)
                continue
            }
            select {
            case <-ctx.Done():
                return
            case <-time.After(50 * time.Millisecond):
            }
        }
    }()

    <-ctx.Done()

    drainCtx, drainCancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer drainCancel()
    if err := svc.Drain(drainCtx); err != nil {
        fmt.Println("drain error:", err)
    }
    fmt.Println("done")
}

Now main does only three things: derive a signal-bound context, start the service, and drain on cancellation. The service exposes Start, Submit, and Drain. This is roughly what production Go services look like.

Step 6 — Track in-flight count externally¶

Sometimes you cannot rely on WaitGroup alone — for example, in an HTTP server where requests are managed by net/http. You can track in-flight count yourself with an atomic counter.

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

func (t *Tracker) Enter() bool {
    if t.closed.Load() {
        return false
    }
    t.inFlight.Add(1)
    if t.closed.Load() {
        t.inFlight.Add(-1)
        return false
    }
    return true
}

func (t *Tracker) Exit() { t.inFlight.Add(-1) }

func (t *Tracker) Drain(ctx context.Context) error {
    t.closed.Store(true)
    tk := time.NewTicker(10 * time.Millisecond)
    defer tk.Stop()
    for t.inFlight.Load() > 0 {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-tk.C:
        }
    }
    return nil
}

The double-check inside Enter handles the race where closed.Store happens between the first Load and the Add.

Step 7 — Verify with a test¶

func TestServiceDrain(t *testing.T) {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    svc := New()
    svc.Start(ctx, 2)

    for i := 0; i < 10; i++ {
        svc.Submit(i)
    }

    dctx, dcancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer dcancel()
    if err := svc.Drain(dctx); err != nil {
        t.Fatalf("drain failed: %v", err)
    }
}

Run with go test -race -count=10. If anything is wrong — a missing Done, a missing close, a goroutine leak — -race will tell you. Running 10 times catches flakiness.

Worked Example: A File-Processing CLI With Drain¶

Let us walk through a complete small program. The program reads a list of file paths from stdin, processes them concurrently, and supports Ctrl+C with a graceful drain.

package main

import (
    "bufio"
    "context"
    "fmt"
    "os"
    "os/signal"
    "sync"
    "syscall"
    "time"
)

type Result struct {
    Path string
    N    int
    Err  error
}

func processFile(ctx context.Context, path string) Result {
    select {
    case <-ctx.Done():
        return Result{Path: path, Err: ctx.Err()}
    case <-time.After(300 * time.Millisecond):
    }
    f, err := os.Open(path)
    if err != nil {
        return Result{Path: path, Err: err}
    }
    defer f.Close()
    scanner := bufio.NewScanner(f)
    n := 0
    for scanner.Scan() {
        n++
        if ctx.Err() != nil {
            return Result{Path: path, N: n, Err: ctx.Err()}
        }
    }
    return Result{Path: path, N: n, Err: scanner.Err()}
}

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

    in := make(chan string, 8)
    out := make(chan Result, 8)
    var wg sync.WaitGroup

    for i := 0; i < 4; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()
            for p := range in {
                select {
                case <-ctx.Done():
                    out <- Result{Path: p, Err: ctx.Err()}
                    return
                case out <- processFile(ctx, p):
                }
            }
        }()
    }

    go func() {
        scanner := bufio.NewScanner(os.Stdin)
        for scanner.Scan() {
            path := scanner.Text()
            select {
            case <-ctx.Done():
                close(in)
                return
            case in <- path:
            }
        }
        close(in)
    }()

    go func() {
        wg.Wait()
        close(out)
    }()

    for r := range out {
        if r.Err != nil {
            fmt.Printf("FAIL %s: %v\n", r.Path, r.Err)
        } else {
            fmt.Printf("OK   %s: %d lines\n", r.Path, r.N)
        }
    }
    fmt.Println("done")
}