package main
import (
"context"
"io"
"net"
"net/http"
"testing"
"time"
"craftdesign.group/quarry/internal/web"
)
// startServing runs serveUntil on a free port and reports when it returns.
func startServing(t *testing.T, h http.Handler, drain *web.Drain, bound time.Duration) (url string, stop context.CancelFunc, returned <-chan error) {
t.Helper()
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
ctx, stop := context.WithCancel(context.Background())
t.Cleanup(stop)
done := make(chan error, 1)
go func() { done <- serveUntil(ctx, newServer(h), ln, drain, bound) }()
return "http://" + ln.Addr().String(), stop, done
}
// On a stop the forge refuses new connections, ends its live streams at
// once, lets a request already in flight finish, and returns well
// inside the drain bound.
func TestTheForgeStopsInOrder(t *testing.T) {
drain := web.NewDrain()
streamOpen := make(chan struct{})
streamGone := make(chan struct{})
slowIn := make(chan struct{})
mux := http.NewServeMux()
mux.HandleFunc("/stream", func(w http.ResponseWriter, req *http.Request) {
web.MarkLive(w)
w.Header().Set("Content-Type", "text/event-stream")
io.WriteString(w, ": open\n\n")
w.(http.Flusher).Flush()
close(streamOpen)
<-req.Context().Done()
close(streamGone)
})
mux.HandleFunc("/slow", func(w http.ResponseWriter, req *http.Request) {
close(slowIn)
time.Sleep(300 * time.Millisecond)
io.WriteString(w, "done")
})
const bound = 5 * time.Second
url, stop, returned := startServing(t, drain.Wrap(mux), drain, bound)
resp, err := http.Get(url + "/stream")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
<-streamOpen
slow := make(chan string, 1)
go func() {
r, err := http.Get(url + "/slow")
if err != nil {
slow <- err.Error()
return
}
b, _ := io.ReadAll(r.Body)
r.Body.Close()
slow <- string(b)
}()
<-slowIn
began := time.Now()
stop()
select {
case <-streamGone:
t.Log("ok: the live stream's request context was cancelled")
case <-time.After(bound / 2):
t.Fatal("FAIL: the live stream was still open halfway through the drain bound")
}
if got := <-slow; got != "done" {
t.Fatalf("FAIL: the in-flight request was cut: %q", got)
}
t.Log("ok: the in-flight request finished")
select {
case err := <-returned:
if err != nil {
t.Fatalf("FAIL: serveUntil: %v", err)
}
case <-time.After(bound + time.Second):
t.Fatal("FAIL: the server did not stop within its bound")
}
if took := time.Since(began); took >= bound {
t.Fatalf("FAIL: the stop took %v, the whole bound", took)
}
t.Log("ok: the server stopped inside its bound")
if _, err := http.Get(url + "/slow"); err == nil {
t.Fatal("FAIL: the stopped server still accepts requests")
}
t.Log("ok: new requests are refused")
}
// A request that never finishes and never declared itself a stream is
// cancelled when the bound runs out, and the stop still ends on time.
func TestTheDrainBoundIsKept(t *testing.T) {
drain := web.NewDrain()
in := make(chan struct{})
cancelled := make(chan struct{})
h := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
close(in)
<-req.Context().Done()
close(cancelled)
})
const bound = 300 * time.Millisecond
url, stop, returned := startServing(t, drain.Wrap(h), drain, bound)
go http.Get(url + "/")
<-in
began := time.Now()
stop()
select {
case <-returned:
case <-time.After(bound + 2*time.Second):
t.Fatal("FAIL: the server outlived its drain bound")
}
t.Logf("ok: stopped after %v with a bound of %v", time.Since(began).Round(time.Millisecond), bound)
select {
case <-cancelled:
t.Log("ok: the request that outlived the bound was cancelled")
case <-time.After(2 * time.Second):
t.Fatal("FAIL: the request that outlived the bound was never cancelled")
}
}