Day 47 — Connection Draining

Learning Objectives

• Implement graceful backend removal without dropping active connections
• Build the drain state machine: active → draining → removed
• Test with 50 concurrent active connections

Exercise: Backend with Drain Support

package main

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

type DrainState int32

const (
    StateActive   DrainState = iota
    StateDraining
    StateRemoved
)

type DrainableBackend struct {
    Addr     string
    state    atomic.Int32
    inFlight sync.WaitGroup
}

func (b *DrainableBackend) State() DrainState {
    return DrainState(b.state.Load())
}

// AcquireSlot returns true if a new connection can be sent to this backend.
// Returns false if draining or removed.
func (b *DrainableBackend) AcquireSlot() bool {
    if b.state.Load() != int32(StateActive) {
        return false
    }
    b.inFlight.Add(1)

    // Re-check state after Add (race: state may have changed between check and Add)
    if b.state.Load() != int32(StateActive) {
        b.inFlight.Done()
        return false
    }
    return true
}

func (b *DrainableBackend) ReleaseSlot() {
    b.inFlight.Done()
}

// Drain marks the backend as draining and waits for in-flight to complete
func (b *DrainableBackend) Drain(timeout time.Duration) bool {
    b.state.Store(int32(StateDraining))
    fmt.Printf("[DRAIN] %s: draining (waiting for in-flight connections)\n", b.Addr)

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

    select {
    case <-done:
        b.state.Store(int32(StateRemoved))
        fmt.Printf("[DRAIN] %s: drained cleanly\n", b.Addr)
        return true
    case <-time.After(timeout):
        b.state.Store(int32(StateRemoved))
        fmt.Printf("[DRAIN] %s: drain timeout, force-removed\n", b.Addr)
        return false
    }
}

func main() {
    backend := &DrainableBackend{Addr: "10.0.0.1:8080"}

    // Simulate 50 concurrent connections
    var wg sync.WaitGroup
    for i := 0; i < 50; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            if !backend.AcquireSlot() {
                fmt.Printf("[conn %d] rejected (draining)\n", id)
                return
            }
            defer backend.ReleaseSlot()
            // Simulate request processing (variable duration)
            time.Sleep(time.Duration(id%500) * time.Millisecond)
        }(i)
    }

    // Drain after 100ms (while connections are still active)
    time.Sleep(100 * time.Millisecond)
    fmt.Printf("Active connections at drain time: state=ACTIVE\n")

    cleaned := backend.Drain(30 * time.Second)
    fmt.Printf("Drain result: clean=%v\n", cleaned)

    wg.Wait()
    fmt.Printf("Final state: %v\n", backend.State())
}

Real-World Context

Kubernetes deployment rollovers use drain: when a new version deploys, old pods receive SIGTERM, stop accepting new requests, and wait for in-flight to complete. Without drain: users get connection reset errors during deploys. With drain: deploys are invisible to users. terminationGracePeriodSeconds is your drain timeout.