package main
import (
"bytes"
"context"
"io"
"log"
"net"
"net/http"
"os"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"craftdesign.group/quarry/internal/api"
"craftdesign.group/quarry/internal/children"
"craftdesign.group/quarry/internal/drain"
"craftdesign.group/quarry/internal/ops"
"craftdesign.group/quarry/internal/zitcli"
)
// startServing runs serveUntil on a free port and reports when it returns.
func startServing(t *testing.T, h http.Handler, d *drain.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, d, 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) {
dr := drain.New()
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) {
drain.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, dr.Wrap(mux), dr, 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) {
dr := drain.New()
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, dr.Wrap(h), dr, 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")
}
}
// The API's streams (chat, attention) are served by the router, not by
// the web layer, and must end at once on a stop the same way.
func TestAnAPIStreamEndsWhenTheForgeStops(t *testing.T) {
type in struct{}
gone := make(chan struct{})
r := ops.NewRegistry()
if err := ops.Register(r, ops.Op[in, struct{}]{
Name: "probe.stream", Doc: "a stream that runs until its request ends",
Kind: ops.Stream, Auth: ops.Anonymous, Route: "GET /api/v1/probe/stream",
StreamRun: func(ctx context.Context, _ in, _ ops.Streamer) error {
<-ctx.Done()
close(gone)
return ctx.Err()
},
}); err != nil {
t.Fatal(err)
}
h, err := api.Router(r, "v")
if err != nil {
t.Fatal(err)
}
dr := drain.New()
const bound = 5 * time.Second
url, stop, returned := startServing(t, dr.Wrap(h), dr, bound)
resp, err := http.Get(url + "/api/v1/probe/stream")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("FAIL: the probe stream answered %d", resp.StatusCode)
}
buf := make([]byte, 8)
if _, err := resp.Body.Read(buf); err != nil {
t.Fatalf("FAIL: the stream sent nothing: %v", err)
}
began := time.Now()
stop()
select {
case <-gone:
t.Log("ok: the API stream's context was cancelled")
case <-time.After(bound / 2):
t.Fatal("FAIL: the API stream was still open halfway through the drain bound")
}
select {
case <-returned:
case <-time.After(bound + time.Second):
t.Fatal("FAIL: the server did not stop within its bound")
}
if took := time.Since(began); took >= bound/2 {
t.Fatalf("FAIL: the stop took %v", took)
}
t.Log("ok: the server stopped at once")
}
// A handler that ignores its cancelled context still finishes before
// serveUntil returns (within a short wait), because the teardown that
// follows closes the index, the events bus and the database under it.
func TestTeardownWaitsForHandlersThatIgnoredTheStop(t *testing.T) {
dr := drain.New()
in := make(chan struct{})
var finished atomic.Bool
h := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
close(in)
<-req.Context().Done()
time.Sleep(400 * time.Millisecond)
finished.Store(true)
})
const bound = 200 * time.Millisecond
url, stop, returned := startServing(t, dr.Wrap(h), dr, bound)
go http.Get(url + "/")
<-in
stop()
select {
case <-returned:
case <-time.After(bound + serveHandlerWait + time.Second):
t.Fatal("FAIL: serveUntil outlived its bound and its handler wait")
}
if !finished.Load() {
t.Fatal("FAIL: serveUntil returned while a handler was still running")
}
t.Log("ok: serveUntil waited for the handler that ignored the stop")
}
// lockedLog is a log destination a test can read while loggers write.
type lockedLog struct {
mu sync.Mutex
buf bytes.Buffer
}
func (l *lockedLog) Write(p []byte) (int, error) {
l.mu.Lock()
defer l.mu.Unlock()
return l.buf.Write(p)
}
func (l *lockedLog) String() string {
l.mu.Lock()
defer l.mu.Unlock()
return l.buf.String()
}
// A stop with a store command still running holds until that zit child
// has ended, and its log names the child and how it ended. In production
// the process returned in 0.2 s with two zit children alive, which
// systemd then SIGKILLed.
func TestAStopWaitsForALiveZitChildAndLogsIt(t *testing.T) {
bin := t.TempDir()
stub := "#!/bin/sh\nsleep 0.6\necho 0123456789abcdef\n"
if err := os.WriteFile(filepath.Join(bin, "zit"), []byte(stub), 0o755); err != nil {
t.Fatal(err)
}
t.Setenv("PATH", bin+string(os.PathListSeparator)+os.Getenv("PATH"))
var lg lockedLog
log.SetOutput(&lg)
t.Cleanup(func() { log.SetOutput(os.Stderr) })
store := filepath.Join(t.TempDir(), "stubrepo")
childDone := make(chan struct{})
go func() {
zitcli.New(store).Info(context.Background(), "trunk")
close(childDone)
}()
for seen := time.Now(); children.Default.Live() == 0; time.Sleep(5 * time.Millisecond) {
if time.Since(seen) > 3*time.Second {
t.Fatal("FAIL: the store client never registered its zit child")
}
}
dr := drain.New()
_, stop, returned := startServing(t, dr.Wrap(http.NotFoundHandler()), dr, time.Second)
began := time.Now()
stop()
<-returned
stopChildren()
took := time.Since(began)
select {
case <-childDone:
default:
t.Fatal("FAIL: the stop finished while the zit child was still running")
}
if took < 300*time.Millisecond {
t.Fatalf("FAIL: the stop returned in %v with a live zit child", took)
}
t.Logf("ok: the stop held %v for the zit child", took.Round(time.Millisecond))
out := lg.String()
for _, want := range []string{"repo=stubrepo", "verb=info", "pid=", "exit=exit status 0", "took="} {
if !strings.Contains(out, want) {
t.Fatalf("FAIL: the stop log lacks %q:\n%s", want, out)
}
}
t.Log("ok: the stop log names the child and how it ended")
}