package web
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 := NewDrain()
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) {
lw, fl, ok := liveStream(w)
if !ok {
t.Error("FAIL: the test server cannot flush")
return
}
lw.Header().Set("Content-Type", "text/event-stream")
io.WriteString(lw, ": open\n\n")
fl.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 := NewDrain()
d.StopStreams()
ended := make(chan error, 1)
h := d.Wrap(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
liveStream(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 := NewDrain()
h := d.Wrap(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
liveStream(w)
}))
for i := 0; i < 5; i++ {
h.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest("GET", "/stream", nil))
}
d.mu.Lock()
n := len(d.live)
d.mu.Unlock()
if n != 0 {
t.Fatalf("FAIL: %d finished streams are still tracked", n)
}
t.Log("ok: finished streams are not tracked")
}