Build your own reverse proxy and load balancer

Every production system has one of these in front of it, and most developers treat it as a configuration file somebody else wrote. A proxy is a few hundred lines, and writing one explains a list of things that are otherwise folklore: why X-Forwarded-For can lie, why some headers must be stripped, why a retry can corrupt data, and why your streaming endpoint buffers when it goes through nginx.

What a reverse proxy does

client  →  proxy  →  one of several backends
                 ←  response

Four jobs, and only the first is obvious:

• Distribute
• Spread requests over healthy backends
• Protect
• Timeouts, limits, and TLS termination in one place
• Hide
• One public address for many internal services
• Survive
• Take a backend out when it fails, put it back when it recovers

The core, which is shorter than you expect

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    backend, release, ok := p.pool.Pick(r)
    if !ok {
        http.Error(w, "no healthy backend", 503)
        return
    }
    defer release()

    outreq := r.Clone(r.Context())
    outreq.URL.Scheme = backend.URL.Scheme
    outreq.URL.Host = backend.URL.Host
    outreq.RequestURI = ""               // must be empty for a client
    stripHopByHop(outreq.Header)
    addForwarded(outreq, r)

    resp, err := p.transport.RoundTrip(outreq)
    if err != nil {
        p.pool.Fail(backend)
        http.Error(w, "upstream unavailable", 502)
        return
    }
    defer resp.Body.Close()

    stripHopByHop(resp.Header)
    copyHeader(w.Header(), resp.Header)
    w.WriteHeader(resp.StatusCode)
    p.copyBody(w, resp.Body)
}

Everything after this is the detail, and the detail is where proxies go wrong.

Hop-by-hop headers, and why they must be removed

Some headers describe this connection, not the message. Forwarding them onward is wrong because the next connection is a different connection:

var hopByHop = []string{
    "Connection", "Keep-Alive", "Proxy-Authenticate",
    "Proxy-Authorization", "TE", "Trailer",
    "Transfer-Encoding", "Upgrade",
}

And the subtle part: Connection can name other headers that are also hop-by-hop for this connection, so those must be stripped too.

X-Forwarded-For, and the header that lies

func addForwarded(out, in *http.Request) {
    ip, _, _ := net.SplitHostPort(in.RemoteAddr)

    prior := in.Header.Get("X-Forwarded-For")
    if prior != "" {
        ip = prior + ", " + ip
    }
    out.Header.Set("X-Forwarded-For", ip)
    out.Header.Set("X-Forwarded-Proto", scheme(in))
    out.Header.Set("X-Forwarded-Host", in.Host)
}

The chain grows as the request passes each proxy, so the leftmost entry is the original client and each later one was added by a hop.

Choosing a backend

Round robin is two lines and is the right default:

func (p *Pool) roundRobin() *Backend {
    n := atomic.AddUint64(&p.next, 1)
    healthy := p.healthy()
    if len(healthy) == 0 {
        return nil
    }
    return healthy[n%uint64(len(healthy))]
}

Least connections is better when request durations vary a lot, which is most real traffic:

func (p *Pool) leastConns() *Backend {
    var best *Backend
    for _, b := range p.healthy() {
        if best == nil ||
            b.inflight.Load() < best.inflight.Load() {
            best = b
        }
    }
    return best
}

• Round robin
• Even request counts. Fine when requests cost the same
• Least connections
• Even load. Better with mixed fast and slow endpoints
• Weighted
• Backends of different sizes
• Consistent hash
• The same key goes to the same backend, for cache locality or sticky sessions

Round robin's failure mode is worth knowing: one slow backend keeps receiving its share, so a quarter of your users wait on the sick machine while three healthy ones idle. Least connections drains away from it automatically, because its in-flight count stays high.

Health, actively and passively

Active: ask each backend, on a timer.

func (p *Pool) probe(b *Backend) {
    ctx, cancel := context.WithTimeout(
        context.Background(), 2*time.Second)
    defer cancel()

    req, _ := http.NewRequestWithContext(ctx, "GET",
        b.URL.String()+"/health", nil)
    resp, err := p.probeClient.Do(req)

    switch {
    case err != nil || resp.StatusCode >= 500:
        b.markUnhealthy()
    default:
        resp.Body.Close()
        b.markHealthy()
    }
}

Passive: notice real traffic failing. A backend that refuses three connections in a row is out, regardless of what its health endpoint says.

Recovery must be careful. A backend that has been out comes back half open: it gets one request, and only if that succeeds does it rejoin.

Thresholds matter too: two consecutive failures to eject, three consecutive successes to restore. One failure is a blip, and treating it as an outage makes the proxy less reliable than no proxy.

Retries, and the one that corrupts data

Retrying a failed request sounds obviously good. It is conditional:

func retryable(r *http.Request, err error, sentBytes bool) bool {
    if sentBytes {
        return false        // the backend may have acted on it
    }
    switch r.Method {
    case "GET", "HEAD", "OPTIONS", "PUT", "DELETE":
        return isConnectionError(err)
    default:
        return false        // POST and PATCH are not idempotent
    }
}

Streaming, and the buffering that breaks it

func (p *Proxy) copyBody(w http.ResponseWriter, body io.Reader) {
    if flusher, ok := w.(http.Flusher); ok && p.flushInterval > 0 {
        // Server-sent events, long polling, log tails: the client
        // needs each write as it happens, not when the buffer is
        // full or the response ends.
        copyAndFlush(w, body, flusher, p.flushInterval)
        return
    }
    io.Copy(w, body)
}

This is why your streaming endpoint works locally and arrives all at once through a proxy: the default is to buffer, and a stream needs the opposite. Detect it by Content-Type: text/event-stream, or by the absence of aContent-Length, and flush.

WebSockets need the connection itself, not the body: on a 101 Switching Protocols, hijack the client connection and copy bytes in both directions until either side closes.

Draining a backend

Deployment is where a proxy earns its place. Taking a backend out gracefullymeans: stop sending it new requests, let the in-flight ones finish, then stop the process.

func (b *Backend) Drain(ctx context.Context) error {
    b.accepting.Store(false)
    for b.inflight.Load() > 0 {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-time.After(50 * time.Millisecond):
        }
    }
    return nil
}

That loop, plus chapter twenty one's graceful shutdown in the backend, is a zero-downtime deploy. There is nothing else to it.

./backend -name a -port 9001
./backend -name b -port 9002 -latency 500ms
./backend -name c -port 9003 -fail-rate 0.5

Assignment

Commit your work with the message Reverse Proxy - Assignment One, exactly as written, so the leaderboard can score it.