Graceful draining

Deep dive

How it works

In: a termination signal and a grace budget · Out: new work refused, in-flight work finished or cut, a going-away notice to clients, exit · State: a draining flag, an in-flight counter, a deadline; session state lives outside the process · Cost: O(1) per request; drain time ≤ grace.

1. Unroute. Flip readiness to false, or run a preStop sleep, so the load balancer stops choosing this instance. Propagation is not instant, and Kubernetes runs preStop to completion before it sends SIGTERM.
2. SIGTERM. Set the draining flag (one atomic) and close or hand over the listening socket.
3. Refuse new work. Increment the in-flight counter first, then test the flag; if draining, decrement and answer "going away". The drain loop cannot miss a request that slipped in.
4. Notify long-lived clients. Send a close frame with code 1001 ("going away", RFC 6455), ideally staggered.
5. Wait until the counter reaches 0 or the deadline passes.
6. Cut what is left; the client replays work that never committed.
7. Exit before the supervisor's SIGKILL: the deadline must be shorter than the grace period.
8. Resume (client). Reconnect after a random delay via the load balancer, present token plus last acknowledged sequence number; the new instance loads the session from the shared store.

 t=0 pod deleted ── preStop (readiness false, LB converges) ── SIGTERM ── drain ≤ deadline ── exit
     |<------------------------ terminationGracePeriodSeconds ----------------------------->| SIGKILL

Code

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst};
use std::sync::{Arc, Mutex};
use std::thread::sleep;
use std::time::{Duration, Instant};

pub type Store = Arc<Mutex<HashMap<u64, u32>>>; // shared store: resume token -> last committed seq

#[derive(Debug, PartialEq)]
pub enum Reply { Done(u32), GoingAway, Cut } // GoingAway ~ WebSocket close 1001

pub struct Server { draining: AtomicBool, force: AtomicBool, inflight: AtomicUsize, store: Store }

impl Server {
    pub fn new(store: Store) -> Arc<Self> {
        Arc::new(Server { draining: false.into(), force: false.into(), inflight: 0.into(), store })
    }
    /// Readiness probe: the load balancer stops routing here when false.
    pub fn ready(&self) -> bool { !self.draining.load(SeqCst) }

    /// Handle request `seq` of session `token`; `work` simulates processing time.
    pub fn handle(&self, token: u64, seq: u32, work: Duration) -> Reply {
        self.inflight.fetch_add(1, SeqCst);      // count first ...
        if self.draining.load(SeqCst) {          // ... then check, so drain() cannot miss us
            self.inflight.fetch_sub(1, SeqCst);
            return Reply::GoingAway;
        }
        let end = Instant::now() + work;
        let mut cut = false;
        while Instant::now() < end && !cut {
            cut = self.force.load(SeqCst);
            sleep(Duration::from_millis(2));
        }
        if !cut {
            let mut s = self.store.lock().unwrap();
            let last = s.entry(token).or_insert(0);
            if seq == *last + 1 { *last = seq; } // commit once; a replayed seq is a no-op
        }
        self.inflight.fetch_sub(1, SeqCst);
        if cut { Reply::Cut } else { Reply::Done(seq) }
    }

    /// SIGTERM handler: refuse new work, wait to the deadline, cut the rest (returns the count).
    pub fn drain(&self, grace: Duration) -> usize {
        self.draining.store(true, SeqCst);       // readiness is false from here on
        let deadline = Instant::now() + grace;
        while self.inflight.load(SeqCst) > 0 && Instant::now() < deadline {
            sleep(Duration::from_millis(1));
        }
        let stragglers = self.inflight.load(SeqCst);
        self.force.store(true, SeqCst);
        while self.inflight.load(SeqCst) > 0 { sleep(Duration::from_millis(1)); }
        stragglers
    }
}

/// Client: on GoingAway or Cut try the next instance with the same token and seq.
pub fn call(instances: &[Arc<Server>], token: u64, seq: u32, work: Duration) -> Reply {
    for s in instances {
        if let Reply::Done(n) = s.handle(token, seq, work) { return Reply::Done(n); }
        // real clients add jittered backoff before reconnecting
    }
    Reply::GoingAway
}

fn main() {
    let store: Store = Default::default();
    let (a, b) = (Server::new(store.clone()), Server::new(store.clone()));
    let ms = Duration::from_millis;
    assert!(a.ready());
    assert_eq!(a.handle(7, 1, ms(5)), Reply::Done(1));

    // Two quick requests and one slow one are in flight on A when SIGTERM arrives.
    let go = |tok, seq, w| { let s = a.clone(); std::thread::spawn(move || s.handle(tok, seq, w)) };
    let (q1, q2, slow) = (go(1, 1, ms(20)), go(2, 1, ms(30)), go(7, 2, ms(2000)));
    sleep(ms(10));
    let t = Instant::now();
    let cut = a.drain(ms(150));
    let took = t.elapsed();
    assert!(!a.ready() && cut == 1 && took >= ms(150) && took < ms(400));
    assert_eq!(q1.join().unwrap(), Reply::Done(1));
    assert_eq!(q2.join().unwrap(), Reply::Done(1));
    assert_eq!(slow.join().unwrap(), Reply::Cut);
    assert_eq!(a.handle(9, 1, ms(1)), Reply::GoingAway); // new work is refused

    // Client 7 resumes on B: seq 2 never committed, so it is replayed; a duplicate is a no-op.
    assert_eq!(store.lock().unwrap()[&7], 1);
    assert_eq!(call(&[a.clone(), b.clone()], 7, 2, ms(5)), Reply::Done(2));
    assert_eq!(call(&[a, b], 7, 2, ms(5)), Reply::Done(2));
    assert_eq!(store.lock().unwrap()[&7], 2);
    println!("busy drain took {took:?}");
}

Complexity and constants

Per request: three atomic operations on a shared counter, cheap when uncontended (not measured here); under many cores its cache line bounces, so shard it if profiles show it. Measured on this Mac (Apple M4 Pro), a 150 ms deadline with one 2 s straggler returned after 154.1-155.6 ms (5 runs). Memory: a few words per server, one store record per resumable client.

Efficiency

• Cost: Per request, three atomic operations. A drain is bounded by its deadline, not by CPU. The real cost is capacity: each closed connection costs a reconnect (TLS handshake, session load) elsewhere. The bound is the grace budget and the reconnect rate the rest of the fleet absorbs.
• Levers:
  • Deregister first (readiness false, preStop) — the balancer stops routing here before the listener closes (Kubernetes docs) — behind any load balancer.
  • Jittered reconnect, first attempt random in 0-5 s, then exponential backoff — no reconnect storm (RFC 6455 §7.2.3) — many clients closed together.
  • Gradual drain: the share of connections told to close grows as drain time elapses (Envoy --drain-strategy gradual, the default) — spreads reconnects — large fleets.
  • Resume with sequence numbers — replay only unacknowledged messages, no full resync — large session state.
  • Socket handover (nginx, Envoy) — no refused connects — single hosts without a balancer.
• Naive → standard → state of the art:
  • kill -9, blind client retries · every connection reset, all clients reconnect at once.
  • SIGTERM, readiness false, drain to a deadline, close 1001, jittered reconnect · only stragglers cut, reconnects spread over seconds.
  • Hot restart with handover, gradual drain, resume token in a shared store · no refused connects, no lost session state.

Implementation notes

• Budget: preStop + drain deadline + exit margin < terminationGracePeriodSeconds. The countdown starts before preStop runs and covers it and the stop after SIGTERM (docs' example: 60 s grace, 55 s hook, 10 s stop gets killed).
• Kubernetes sends TERM to process 1 of each container; a shell wrapper that does not forward signals never drains.
• systemd: TimeoutStopSec= bounds ExecStop= and the wait after SIGTERM, then SIGKILL; infinity disables it. See process supervision.
• Go http.Server.Shutdown(ctx) closes listeners and idle connections, then waits for active ones, but ignores hijacked connections such as WebSockets: notify and wait for those yourself.
• WebSocket close: send 1001, await the peer's close, then the server closes TCP first and holds TIME_WAIT (RFC 6455 §7.1.1).
• Handover: SO_REUSEPORT (Linux 3.9) lets processes of one UID bind the same address, each with its own accept queue. SCM_RIGHTS over a Unix socket passes an open file description, like dup(2) into the other process, so the successor inherits the same listener.
• nginx: USR2 starts a new master, WINCH gracefully stops the old workers, QUIT the old master; until QUIT you can roll back. Envoy's new process initializes, then asks the old one for its listen sockets over Unix domain sockets and tells it to drain.
• Keep session state outside the process (Redis, a database, a signed token).

Pitfalls

• Closing the listener before readiness converges: requests hit a dead port.
• Grace period shorter than the drain deadline: SIGKILL fires mid-drain.
• Everyone gets 1001 in the same millisecond and reconnects together: stagger closes, jitter reconnects.
• Non-idempotent replay: a cut request may have partly applied. Use sequence numbers or idempotency keys, commit once.
• Reconnecting to the same draining instance because the balancer still lists it. AWS: a deregistering target that drops a connection before the delay elapses gives the client a 5xx.
• (From memory) Closing one SO_REUSEPORT listener resets connections still in its accept queue; hand over the socket instead.

Variants and alternatives

In production

Kubernetes (kubelet termination flow, preStop), nginx, Envoy (hot restart, drain manager), Go net/http, axum, AWS Elastic Load Balancing, systemd, Cloudflare tableflip; in web and API fleets, multiplayer and chat gateways, queue consumers and database proxies.

Implementations

Sources

• Kubernetes, Pod Lifecycle. docs · Container lifecycle hooks. docs
• systemd, systemd.service(5), TimeoutStopSec=. man
• Go, net/http package, Server.Shutdown. docs
• nginx, Controlling nginx (upgrading the executable on the fly). docs
• Envoy, Hot restart. docs · Command line options. docs
• I. Fette, A. Melnikov, The WebSocket Protocol, RFC 6455, 2011 (close codes, §7.1.1, §7.2.3). RFC
• Linux man-pages, socket(7) (SO_REUSEPORT) and unix(7) (SCM_RIGHTS). socket · unix
• AWS, Edit target group attributes for your Application Load Balancer (deregistration delay). docs
• Cloudflare, Graceful upgrades in Go. blog