package main
import (
"bytes"
"context"
"errors"
"io"
"log"
"net"
"net/http"
"os"
"path/filepath"
"regexp"
"strconv"
"strings"
"sync"
"sync/atomic"
"syscall"
"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
}
// hangGuard bounds a wait that should end by an event, not by time. It
// is not an assertion about speed: every wait below finishes as soon as
// the event it names happens, and the guard only turns a genuine hang
// into a failure instead of the test binary's ten-minute timeout.
const hangGuard = time.Minute
func await(t *testing.T, ch <-chan struct{}, what string) {
t.Helper()
select {
case <-ch:
case <-time.After(hangGuard):
t.Fatalf("FAIL: %s never happened", what)
}
}
// witness is an ordinary request held in flight until release. Its
// context is cancelled only when the drain bound runs out (streams are
// cancelled on their own), so a live witness context proves that
// whatever just happened happened before the bound, without timing it.
type witness struct {
in chan struct{}
ctx chan context.Context
release chan struct{}
answer chan string
}
func newWitness() *witness {
return &witness{in: make(chan struct{}), ctx: make(chan context.Context, 1),
release: make(chan struct{}), answer: make(chan string, 1)}
}
func (w *witness) handler(rw http.ResponseWriter, req *http.Request) {
w.ctx <- req.Context()
close(w.in)
<-w.release
io.WriteString(rw, "done")
}
// send issues the witness request and waits until the handler holds it.
func (w *witness) send(t *testing.T, url string) context.Context {
t.Helper()
go func() {
r, err := http.Get(url)
if err != nil {
w.answer <- err.Error()
return
}
b, _ := io.ReadAll(r.Body)
r.Body.Close()
w.answer <- string(b)
}()
await(t, w.in, "the witness request reaching its handler")
return <-w.ctx
}
// On a stop the forge refuses new connections, ends its live streams at
// once, lets a request already in flight finish, and returns once it has.
// "At once" is shown by order, not by a clock: the stream ends while the
// in-flight request is still running with a live context, which it could
// not be if the stream had waited for the drain bound.
func TestTheForgeStopsInOrder(t *testing.T) {
dr := drain.New()
streamOpen := make(chan struct{})
streamGone := make(chan struct{})
wit := newWitness()
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", wit.handler)
url, stop, returned := startServing(t, dr.Wrap(mux), dr, hangGuard)
resp, err := http.Get(url + "/stream")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
await(t, streamOpen, "the stream opening")
slowCtx := wit.send(t, url+"/slow")
stop()
await(t, streamGone, "the live stream ending")
if err := slowCtx.Err(); err != nil {
t.Fatalf("FAIL: the stream ended only when the drain cancelled everything: %v", err)
}
t.Log("ok: the live stream ended while the in-flight request was still live")
select {
case <-returned:
t.Fatal("FAIL: the server stopped with a request still in flight")
default:
}
t.Log("ok: the stop waits for the in-flight request")
close(wit.release)
if got := <-wit.answer; got != "done" {
t.Fatalf("FAIL: the in-flight request was cut: %q", got)
}
t.Log("ok: the in-flight request finished")
var serr error
await(t, chanOf(returned, &serr), "the server stopping")
if serr != nil {
t.Fatalf("FAIL: serveUntil: %v", serr)
}
t.Log("ok: the server stopped once the request finished")
if _, err := http.Get(url + "/slow"); err == nil {
t.Fatal("FAIL: the stopped server still accepts requests")
}
t.Log("ok: new requests are refused")
}
// chanOf turns serveUntil's result into an event, keeping its error.
func chanOf(returned <-chan error, err *error) <-chan struct{} {
ch := make(chan struct{})
go func() {
*err = <-returned
close(ch)
}()
return ch
}
// A request that never finishes and never declared itself a stream is
// cancelled when the bound runs out, and the stop ends then. The bound
// is checked from below only (the stop cannot end before it, with the
// request holding it); how soon after it the stop ends is the machine's
// business, so it is shown by order: the request was cancelled before
// serveUntil returned.
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 + "/")
await(t, in, "the request reaching its handler")
began := time.Now()
stop()
var serr error
await(t, chanOf(returned, &serr), "the server stopping")
if took := time.Since(began); took < bound {
t.Fatalf("FAIL: the stop ended after %v, before its %v bound, with a request holding it", took, bound)
}
t.Log("ok: the stop held the request for the whole bound")
select {
case <-cancelled:
t.Log("ok: the request that outlived the bound was cancelled before the stop ended")
default:
t.Fatal("FAIL: the stop ended without cancelling the request that outlived the bound")
}
}
// 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. As above,
// "at once" is shown by an in-flight request whose context is still live.
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)
}
wit := newWitness()
mux := http.NewServeMux()
mux.HandleFunc("/witness", wit.handler)
mux.Handle("/", h)
dr := drain.New()
url, stop, returned := startServing(t, dr.Wrap(mux), dr, hangGuard)
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)
}
witCtx := wit.send(t, url+"/witness")
stop()
await(t, gone, "the API stream ending")
if err := witCtx.Err(); err != nil {
t.Fatalf("FAIL: the API stream ended only when the drain cancelled everything: %v", err)
}
t.Log("ok: the API stream's context was cancelled while an ordinary request was still live")
close(wit.release)
var serr error
await(t, chanOf(returned, &serr), "the server stopping")
t.Log("ok: the server stopped")
}
// A handler that ignores its cancelled context still finishes before
// serveUntil returns, because the teardown that follows closes the
// index, the events bus and the database under it. The handler is held
// by a channel rather than a sleep, so the order is fixed by events.
func TestTeardownWaitsForHandlersThatIgnoredTheStop(t *testing.T) {
dr := drain.New()
in := make(chan struct{})
cancelled := make(chan struct{})
release := make(chan struct{})
var finished atomic.Bool
h := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
close(in)
<-req.Context().Done()
close(cancelled)
<-release
finished.Store(true)
})
const bound = 200 * time.Millisecond
url, stop, returned := startServing(t, dr.Wrap(h), dr, bound)
go http.Get(url + "/")
await(t, in, "the request reaching its handler")
stop()
await(t, cancelled, "the bound cancelling the handler")
select {
case <-returned:
t.Fatal("FAIL: serveUntil returned while a handler was still running")
default:
}
close(release)
var serr error
await(t, chanOf(returned, &serr), "the server stopping")
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,
// and wait on: until reports when a line containing a string arrives.
type lockedLog struct {
mu sync.Mutex
buf bytes.Buffer
changed chan struct{}
}
func (l *lockedLog) Write(p []byte) (int, error) {
l.mu.Lock()
defer l.mu.Unlock()
n, err := l.buf.Write(p)
if l.changed != nil {
close(l.changed)
l.changed = nil
}
return n, err
}
func (l *lockedLog) String() string {
l.mu.Lock()
defer l.mu.Unlock()
return l.buf.String()
}
func (l *lockedLog) until(substr string) <-chan struct{} {
ch := make(chan struct{})
go func() {
for {
l.mu.Lock()
if strings.Contains(l.buf.String(), substr) {
l.mu.Unlock()
close(ch)
return
}
if l.changed == nil {
l.changed = make(chan struct{})
}
changed := l.changed
l.mu.Unlock()
<-changed
}
}()
return ch
}
// The stub zit announces its own start on one named pipe and then blocks
// reading another until the test releases it, so the test knows the
// child is running and decides when it ends. No step waits on a clock.
const blockingStubZit = `#!/bin/sh
echo started > "$STUB_STARTED"
read release < "$STUB_RELEASE"
echo 0123456789abcdef
`
func mkfifo(t *testing.T, dir, name string) string {
t.Helper()
path := filepath.Join(dir, name)
if err := syscall.Mkfifo(path, 0o600); err != nil {
t.Fatal(err)
}
return path
}
// A stop with a store command still running holds until that zit child
// has ended, and its log names the child, its pid 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()
if err := os.WriteFile(filepath.Join(bin, "zit"), []byte(blockingStubZit), 0o755); err != nil {
t.Fatal(err)
}
pipes := t.TempDir()
started, releasePipe := mkfifo(t, pipes, "started"), mkfifo(t, pipes, "release")
t.Setenv("PATH", bin+string(os.PathListSeparator)+os.Getenv("PATH"))
t.Setenv("STUB_STARTED", started)
t.Setenv("STUB_RELEASE", releasePipe)
// Opened read-write, the release pipe never blocks the test, and the
// stub's read returns as soon as a line is written. Done once, and
// again at cleanup so a failing test never leaves the stub blocked.
var once sync.Once
release := func() {
once.Do(func() {
f, err := os.OpenFile(releasePipe, os.O_RDWR, 0)
if err != nil {
t.Errorf("open release pipe: %v", err)
return
}
io.WriteString(f, "go\n")
f.Close()
})
}
t.Cleanup(release)
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)
}()
// The stub's own signal: opening the pipe for reading returns once
// the stub has opened it to write, so the child is running.
startedSeen := make(chan struct{})
go func() {
if f, err := os.Open(started); err == nil {
io.ReadAll(f)
f.Close()
}
close(startedSeen)
}()
await(t, startedSeen, "the stub zit announcing its start")
t.Log("ok: the zit child is running")
dr := drain.New()
_, stop, returned := startServing(t, dr.Wrap(http.NotFoundHandler()), dr, hangGuard)
stop()
var serr error
await(t, chanOf(returned, &serr), "the server stopping")
// The stop closes the process's children for good; reopen them for
// the tests that run after this one.
t.Cleanup(children.Default.Reset)
stopped := make(chan struct{})
go func() {
stopChildren()
close(stopped)
}()
await(t, lg.until("stopping: waiting up to"), "the stop beginning to wait for children")
select {
case <-stopped:
t.Fatal("FAIL: the stop returned while a zit child was still running")
default:
}
t.Log("ok: the stop is waiting on the live zit child")
release()
await(t, stopped, "the stop finishing")
if n := children.Default.Live(); n != 0 {
t.Fatalf("FAIL: the stop finished with %d zit children still tracked", n)
}
await(t, childDone, "the store call returning")
t.Log("ok: the stop finished only after the child ended")
var ended string
for _, line := range strings.Split(lg.String(), "\n") {
if strings.Contains(line, "zit child ended") && strings.Contains(line, "repo=stubrepo") {
ended = line
}
}
if ended == "" {
t.Fatalf("FAIL: the stop logged no end for the child:\n%s", lg.String())
}
for _, want := range []string{"verb=info", "exit=exit status 0", "took="} {
if !strings.Contains(ended, want) {
t.Fatalf("FAIL: the child's end line lacks %q: %s", want, ended)
}
}
m := regexp.MustCompile(`pid=(\d+)`).FindStringSubmatch(ended)
if m == nil {
t.Fatalf("FAIL: the child's end line has no pid: %s", ended)
}
if pid, _ := strconv.Atoi(m[1]); pid <= 0 {
t.Fatalf("FAIL: the child's end line has pid %d", pid)
}
t.Logf("ok: the stop logged the child's end with its pid: %s", ended)
}
// A stop with no zit child running still closes the process's children,
// so nothing started between the stop and the exit escapes the wait.
func TestAQuietStopStillRefusesNewChildren(t *testing.T) {
t.Cleanup(children.Default.Reset)
if n := children.Default.Live(); n != 0 {
t.Skipf("%d children live from another test", n)
}
stopChildren()
_, _, err := children.Default.Command(context.Background(), "demo", "late", "true")
if !errors.Is(err, children.ErrStopping) {
t.Fatalf("FAIL: a child after a quiet stop was not refused: %v", err)
}
t.Log("ok: a quiet stop still refuses new children")
}