package drain
import (
"context"
"io"
"net/http"
"net/http/httptest"
"testing"
"time"
)
// A stopping server ends its live streams at once, by cancelling their
// request contexts, and leaves an ordinary request in flight to finish:
// a stream never ends on its own, so waiting for it only spends the
// drain bound, while an ordinary request may be halfway through a write.
func TestStoppingStreamsCancelsLiveStreamsOnly(t *testing.T) {
d := New()
streamGone := make(chan struct{})
ordinaryIn := make(chan struct{})
ordinaryCtx := make(chan context.Context, 1)
release := make(chan struct{})
mux := http.NewServeMux()
mux.HandleFunc("/stream", func(w http.ResponseWriter, req *http.Request) {
MarkLive(w)
w.Header().Set("Content-Type", "text/event-stream")
io.WriteString(w, ": open\n\n")
w.(http.Flusher).Flush()
<-req.Context().Done()
close(streamGone)
})
mux.HandleFunc("/ordinary", func(w http.ResponseWriter, req *http.Request) {
ordinaryCtx <- req.Context()
close(ordinaryIn)
<-release
io.WriteString(w, "done")
})
srv := httptest.NewServer(d.Wrap(mux))
defer srv.Close()
resp, err := http.Get(srv.URL + "/stream")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
buf := make([]byte, 8)
if _, err := resp.Body.Read(buf); err != nil {
t.Fatalf("FAIL: the stream sent nothing: %v", err)
}
ordinary := make(chan string, 1)
go func() {
r, err := http.Get(srv.URL + "/ordinary")
if err != nil {
ordinary <- err.Error()
return
}
b, _ := io.ReadAll(r.Body)
r.Body.Close()
ordinary <- string(b)
}()
<-ordinaryIn
d.StopStreams()
select {
case <-streamGone:
t.Log("ok: the live stream's context was cancelled")
case <-time.After(2 * time.Second):
t.Fatal("FAIL: stopping left the live stream running")
}
if err := (<-ordinaryCtx).Err(); err != nil {
t.Fatalf("FAIL: stopping cancelled an ordinary request: %v", err)
}
t.Log("ok: the ordinary request's context is still live")
close(release)
if got := <-ordinary; got != "done" {
t.Fatalf("FAIL: the ordinary request did not finish: %q", got)
}
t.Log("ok: the ordinary request finished")
}
// A stream that opens after the stop began is ended as it opens, rather
// than living until the process is killed under it.
func TestAStreamOpenedWhileStoppingEndsAtOnce(t *testing.T) {
d := New()
d.StopStreams()
ended := make(chan error, 1)
h := d.Wrap(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
MarkLive(w)
select {
case <-req.Context().Done():
ended <- nil
case <-time.After(2 * time.Second):
ended <- context.DeadlineExceeded
}
}))
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest("GET", "/stream", nil))
if err := <-ended; err != nil {
t.Fatal("FAIL: a stream opened during the stop kept running")
}
t.Log("ok: a late stream is cancelled as it opens")
}
// A finished request leaves nothing behind for the stop to walk.
func TestAFinishedStreamIsForgotten(t *testing.T) {
d := New()
h := d.Wrap(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
MarkLive(w)
}))
for i := 0; i < 5; i++ {
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest("GET", "/stream", nil))
}
d.mu.Lock()
n, running := len(d.live), d.running
d.mu.Unlock()
if n != 0 || running != 0 {
t.Fatalf("FAIL: %d finished streams and %d handlers are still tracked", n, running)
}
t.Log("ok: finished requests are not tracked")
}
// Wait reports whether every handler has returned, so teardown does not
// close stores under one that ignored its cancelled context.
func TestWaitCoversHandlersStillRunning(t *testing.T) {
d := New()
in := make(chan struct{})
release := make(chan struct{})
h := d.Wrap(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
close(in)
<-release
}))
go h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest("GET", "/", nil))
<-in
if d.Wait(50 * time.Millisecond) {
t.Fatal("FAIL: Wait reported idle with a handler still running")
}
t.Log("ok: a running handler keeps Wait from reporting idle")
close(release)
if !d.Wait(2 * time.Second) {
t.Fatal("FAIL: Wait did not see the handler return")
}
t.Log("ok: Wait returns once the handler has")
}