internal/agent/syncclient/leak_test.go

Ref: main v0.0.5   Size: 4.2 KiB   History

package syncclient

import (
	"context"
	"runtime"
	"testing"
	"time"

	"github.com/a73x/eitri/internal/pb"
	"github.com/a73x/eitri/internal/transport"
	"github.com/stretchr/testify/require"
)

// settleGoroutines forces GC and polls until the live goroutine count drops to
// target or the deadline passes, returning the last observed count. It lets any
// per-session goroutines (recv/worker/console-accept + QUIC internals) unwind
// before we sample, so the assertion isn't racing normal teardown.
func settleGoroutines(target int, d time.Duration) int {
	deadline := time.Now().Add(d)
	n := runtime.NumGoroutine()
	for time.Now().Before(deadline) {
		runtime.GC()
		n = runtime.NumGoroutine()
		if n <= target {
			return n
		}
		time.Sleep(20 * time.Millisecond)
	}
	return n
}

// newDropServer runs a minimal QUIC server that, for every connection, completes
// the sync handshake (reads Hello, opens the down-stream, pushes one snapshot so
// the client marks the session connected and exercises a step) and then drops
// the connection — exactly the production reconnect trigger: the session ends
// because the down-stream read errors, NOT because the parent ctx was cancelled.
// It returns the listen addr, the cert fingerprint to pin, and a stop func.
func newDropServer(t *testing.T) (addr, fp string, stop func()) {
	t.Helper()
	var certPEM, keyPEM []byte
	certPEM, keyPEM, fp = genTestCert(t)
	lis := listenTestQUIC(t, "127.0.0.1:0", certPEM, keyPEM)
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		for {
			conn, err := lis.Accept(ctx)
			if err != nil {
				return
			}
			go func() {
				up, err := conn.AcceptStream(ctx)
				if err != nil {
					return
				}
				var first pb.AgentMessage
				if err := transport.ReadMsg(up, &first, transport.DefaultMaxFrame); err != nil {
					return
				}
				down, err := conn.OpenStreamSync(ctx)
				if err != nil {
					return
				}
				// One snapshot makes the client mark the session connected and run a
				// step; then drop so the client's recv goroutine errors and the
				// session ends the way a real server drop / idle timeout ends it.
				_ = transport.WriteMsg(down, &pb.ServerMessage{
					Msg: &pb.ServerMessage_Snapshot{Snapshot: &pb.Snapshot{}}})
				time.Sleep(30 * time.Millisecond)
				_ = conn.CloseWithError(0, "drop")
			}()
		}
	}()
	return lis.Addr().String(), fp, func() { cancel(); lis.Close() }
}

// TestSessionNoWorkerLeakAcrossReconnects drives many session start/stop cycles
// (each ends by dropping the server, exactly like a production reconnect/deploy)
// and asserts the live goroutine count returns to its post-warmup baseline.
//
// The regression it guards: the up-stream worker goroutine selects only on the
// long-lived parent ctx, a dead stepSignal, and a ticker that session stops on
// return — so when a session ends because the RECV goroutine errored (server
// drop / idle timeout), the worker has no signal to exit and leaks one
// goroutine per reconnect. Over long uptime these accumulate and the server-side
// sync eventually wedges.
func TestSessionNoWorkerLeakAcrossReconnects(t *testing.T) {
	addr, fp, stop := newDropServer(t)
	defer stop()
	c := newClient(t, addr, fp, "host-x", "host-x.deadbeef")

	ctx := t.Context()

	// Each session connects, gets one snapshot, then the server drops it, so
	// session() returns promptly. The parent ctx is never cancelled — precisely
	// the condition under which the worker goroutine cannot observe session end.
	runOne := func() {
		done := make(chan error, 1)
		go func() { done <- c.session(ctx) }()
		select {
		case <-done:
		case <-time.After(5 * time.Second):
			t.Fatal("session did not return after server drop")
		}
	}

	// Warm up to a steady goroutine state before sampling the baseline.
	runOne()
	runOne()
	base := settleGoroutines(0, 3*time.Second) // target 0 => returns settled count

	const cycles = 10
	for range cycles {
		runOne()
	}

	// With the leak, each cycle strands one worker goroutine, so the count sits
	// near base+cycles and never settles back. A small slack absorbs QUIC jitter.
	got := settleGoroutines(base+2, 5*time.Second)
	require.LessOrEqualf(t, got, base+2,
		"goroutine count grew across %d reconnects (base=%d, got=%d): a worker goroutine leaks per session",
		cycles, base, got)
}