// Package drain lets a stopping server end its live streams without
// cutting the rest of its requests short, and tells it when every
// handler has returned.
//
// http.Server.Shutdown waits for every handler to return, and a live
// stream never returns on its own: it waits on its request's context,
// which Shutdown does not cancel. Left alone, every open tab would hold
// a deploy for the whole drain bound. Cancelling every request instead
// would end a push or a save halfway through. So each request carries
// its own cancel, and only those that declared themselves streams are
// cancelled when the stop begins.
//
// A leaf package because streams are opened in two layers, the web
// pages and the API router, and the router cannot import the web.
package drain
import (
"context"
"net/http"
"sync"
"time"
)
// Drain tracks the requests passing through one server.
type Drain struct {
mu sync.Mutex
live map[*writer]struct{}
stopped bool
running int
waiting chan struct{}
waitingOnce sync.Once
}
// New returns a Drain with nothing in flight.
func New() *Drain {
return &Drain{live: map[*writer]struct{}{}, waiting: make(chan struct{})}
}
// Waiting closes when Wait is first called: the stop has reached the
// point of holding for its handlers. It lets a caller order an event
// after that point rather than guess at it with a clock.
func (d *Drain) Waiting() <-chan struct{} { return d.waiting }
// Wrap gives every request through h a context the Drain can cancel,
// and counts it until its handler returns.
func (d *Drain) Wrap(h http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
ctx, cancel := context.WithCancel(req.Context())
dw := &writer{ResponseWriter: w, d: d, cancel: cancel}
d.mu.Lock()
d.running++
d.mu.Unlock()
defer func() {
d.mu.Lock()
delete(d.live, dw)
d.running--
d.mu.Unlock()
cancel()
}()
h.ServeHTTP(dw, req.WithContext(ctx))
})
}
// StopStreams cancels every live stream in flight, and every stream
// that opens from now on as it opens. Ordinary requests are untouched.
func (d *Drain) StopStreams() {
d.mu.Lock()
d.stopped = true
live := d.live
d.live = map[*writer]struct{}{}
d.mu.Unlock()
for dw := range live {
dw.cancel()
}
}
// Wait reports whether every handler has returned, waiting up to
// bound. http.Server.Close drops connections without waiting for their
// handlers, so a handler that ignored its cancelled context can still
// be running when the caller moves on to closing the stores it uses.
func (d *Drain) Wait(bound time.Duration) bool {
d.waitingOnce.Do(func() { close(d.waiting) })
deadline := time.Now().Add(bound)
for {
d.mu.Lock()
n := d.running
d.mu.Unlock()
if n == 0 {
return true
}
if !time.Now().Before(deadline) {
return false
}
time.Sleep(10 * time.Millisecond)
}
}
func (d *Drain) markLive(dw *writer) {
d.mu.Lock()
stopped := d.stopped
if !stopped {
d.live[dw] = struct{}{}
}
d.mu.Unlock()
if stopped {
dw.cancel()
}
}
type writer struct {
http.ResponseWriter
d *Drain
cancel context.CancelFunc
}
func (dw *writer) markLive() { dw.d.markLive(dw) }
// Flush is forwarded because streams find their flusher by type
// assertion, and a wrapper without one would silently disable them.
func (dw *writer) Flush() {
if f, ok := dw.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
func (dw *writer) Unwrap() http.ResponseWriter { return dw.ResponseWriter }
// MarkLive declares that the response being written through w is a
// long-lived stream, so a stopping server may end it at once. It walks
// the writer's wrappers to find the Drain's; with no Drain in the chain
// it does nothing.
func MarkLive(w http.ResponseWriter) {
for w != nil {
if m, ok := w.(interface{ markLive() }); ok {
m.markLive()
return
}
u, ok := w.(interface{ Unwrap() http.ResponseWriter })
if !ok {
return
}
w = u.Unwrap()
}
}