package web
import (
"context"
"net/http"
"sync"
)
// Drain lets a stopping server end its live streams without cutting
// the rest of its requests short.
//
// 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.
type Drain struct {
mu sync.Mutex
live map[*drainWriter]struct{}
stopped bool
}
// NewDrain returns a Drain with nothing in flight.
func NewDrain() *Drain {
return &Drain{live: map[*drainWriter]struct{}{}}
}
// Wrap gives every request through h a context the Drain can cancel.
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 := &drainWriter{ResponseWriter: w, d: d, cancel: cancel}
defer func() {
d.mu.Lock()
delete(d.live, dw)
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[*drainWriter]struct{}{}
d.mu.Unlock()
for dw := range live {
dw.cancel()
}
}
func (d *Drain) markLive(dw *drainWriter) {
d.mu.Lock()
stopped := d.stopped
if !stopped {
d.live[dw] = struct{}{}
}
d.mu.Unlock()
if stopped {
dw.cancel()
}
}
type drainWriter struct {
http.ResponseWriter
d *Drain
cancel context.CancelFunc
}
func (dw *drainWriter) 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 *drainWriter) Flush() {
if f, ok := dw.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
func (dw *drainWriter) 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()
}
}