panel/agent/internal/dispatch/heartbeat_exitwatch_test.go

package dispatch

// Tests for WI-05: heartbeats routed through the send queue (sendLoop as the
// sole owner of stream.Send) and the single-flight, ctx-bound exit watcher.

import (
	"context"
	"io"
	"log/slog"
	"sync"
	"testing"
	"time"

	"github.com/dbledeez/panel/agent/internal/runtime"
	panelv1 "github.com/dbledeez/panel/proto/panel/v1"
)

// recordingSender captures every envelope the sendLoop delivers.
type recordingSender struct {
	mu   sync.Mutex
	envs []*panelv1.AgentEnvelope
}

func (r *recordingSender) Send(e *panelv1.AgentEnvelope) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.envs = append(r.envs, e)
	return nil
}

func (r *recordingSender) snapshot() []*panelv1.AgentEnvelope {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]*panelv1.AgentEnvelope, len(r.envs))
	copy(out, r.envs)
	return out
}

// TestSendHeartbeatGoesThroughQueue verifies SendHeartbeat enqueues onto
// sendCh and the envelope reaches the sender via sendLoop — i.e. heartbeats
// no longer need a direct stream.Send from the session goroutine.
func TestSendHeartbeatGoesThroughQueue(t *testing.T) {
	d := newTestDispatcher()
	rs := &recordingSender{}
	d.SetSender(rs)

	now := time.Now()
	d.SendHeartbeat(now)

	deadline := time.After(2 * time.Second)
	for {
		for _, e := range rs.snapshot() {
			if hb := e.GetHeartbeat(); hb != nil {
				if got := hb.At.AsTime(); !got.Equal(now.UTC().Truncate(time.Nanosecond)) && got.Unix() != now.Unix() {
					t.Fatalf("heartbeat timestamp mismatch: got %v want %v", got, now)
				}
				return // delivered through the queue — pass
			}
		}
		select {
		case <-deadline:
			t.Fatal("heartbeat never delivered through sendLoop")
		case <-time.After(10 * time.Millisecond):
		}
	}
}

// TestHeartbeatIsDroppable pins the backpressure classification: a heartbeat
// must be shed (not block the producer) when the buffer is full.
func TestHeartbeatIsDroppable(t *testing.T) {
	env := &panelv1.AgentEnvelope{Payload: &panelv1.AgentEnvelope_Heartbeat{Heartbeat: &panelv1.Heartbeat{}}}
	if !isDroppableEnvelope(env) {
		t.Fatal("heartbeat envelope must be droppable telemetry")
	}
}

// fakeExitRuntime implements just enough of runtime.Runtime for watchExit.
// The embedded nil interface panics on any unexpected call — that's a test
// failure signal, not a hazard.
type fakeExitRuntime struct {
	runtime.Runtime
	mu          sync.Mutex
	activeWaits int
	maxWaits    int
	release     chan struct{} // close to make Wait return exit code 0
}

func (f *fakeExitRuntime) Wait(ctx context.Context, id string) (int, error) {
	f.mu.Lock()
	f.activeWaits++
	if f.activeWaits > f.maxWaits {
		f.maxWaits = f.activeWaits
	}
	f.mu.Unlock()
	defer func() {
		f.mu.Lock()
		f.activeWaits--
		f.mu.Unlock()
	}()
	select {
	case <-ctx.Done():
		return -1, ctx.Err()
	case <-f.release:
		return 0, nil
	}
}

func (f *fakeExitRuntime) InspectByName(ctx context.Context, name string) (runtime.RuntimeInstanceState, error) {
	return runtime.RuntimeInstanceState{Status: "exited"}, nil
}

func (f *fakeExitRuntime) counts() (active, max int) {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.activeWaits, f.maxWaits
}

func newExitWatchDispatcher(f *fakeExitRuntime) *Dispatcher {
	d := &Dispatcher{
		log:       slog.New(slog.NewTextHandler(io.Discard, nil)),
		runtime:   f,
		instances: map[string]*instanceRecord{},
		sendCh:    make(chan *panelv1.AgentEnvelope, 256),
	}
	go d.sendLoop()
	return d
}

// waitFor polls cond until true or the deadline elapses.
func waitFor(t *testing.T, d time.Duration, cond func() bool, msg string) {
	t.Helper()
	deadline := time.Now().Add(d)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal(msg)
}

// TestStartExitWatchSingleFlight re-arms the watcher many times (the shape of
// Announce firing on every controller reconnect) and asserts only ONE
// ContainerWait is ever in flight — the pre-fix behavior leaked one blocked
// Wait (and its docker connection) per reconnect.
func TestStartExitWatchSingleFlight(t *testing.T) {
	f := &fakeExitRuntime{release: make(chan struct{})}
	d := newExitWatchDispatcher(f)
	rec := &instanceRecord{InstanceID: "x", ContainerID: "cid"}

	for i := 0; i < 8; i++ {
		d.startExitWatch(rec)
	}
	// All superseded watchers must exit; exactly one survivor remains.
	waitFor(t, 2*time.Second, func() bool { a, _ := f.counts(); return a == 1 },
		"expected exactly 1 in-flight ContainerWait after re-arms")
	// The peak may briefly exceed 1 while a cancelled watcher unwinds, but it
	// must never approach the number of re-arms (pre-fix: 8 concurrent Waits).
	if _, max := f.counts(); max >= 8 {
		t.Fatalf("watchExit not single-flighted: peak concurrent waits = %d", max)
	}

	// Cancelling the record's watcher (handleStop / Shutdown path) must
	// release the last Wait.
	d.mu.Lock()
	cancel := rec.ExitWatchCancel
	rec.ExitWatchCancel = nil
	d.mu.Unlock()
	cancel()
	waitFor(t, 2*time.Second, func() bool { a, _ := f.counts(); return a == 0 },
		"cancelled exit watcher did not release ContainerWait")
}

// TestWatchExitCancelledEmitsNothing: a ctx-cancelled watcher must return
// silently — no spurious stopped/crashed InstanceState for a container that
// is still running.
func TestWatchExitCancelledEmitsNothing(t *testing.T) {
	f := &fakeExitRuntime{release: make(chan struct{})}
	d := newExitWatchDispatcher(f)
	rs := &recordingSender{}
	d.SetSender(rs)
	rec := &instanceRecord{InstanceID: "x", ContainerID: "cid"}

	d.startExitWatch(rec)
	waitFor(t, 2*time.Second, func() bool { a, _ := f.counts(); return a == 1 }, "watcher never started")
	d.mu.Lock()
	cancel := rec.ExitWatchCancel
	d.mu.Unlock()
	cancel()
	waitFor(t, 2*time.Second, func() bool { a, _ := f.counts(); return a == 0 }, "watcher never exited")
	time.Sleep(50 * time.Millisecond) // allow any (wrong) emission to flush
	for _, e := range rs.snapshot() {
		if e.GetInstanceState() != nil {
			t.Fatalf("cancelled watchExit emitted InstanceState: %v", e)
		}
	}
}

// TestWatchExitTerminalExit: when the container really exits (Wait returns
// 0 and inspect says "exited"), the watcher emits a STOPPED state.
func TestWatchExitTerminalExit(t *testing.T) {
	f := &fakeExitRuntime{release: make(chan struct{})}
	d := newExitWatchDispatcher(f)
	rs := &recordingSender{}
	d.SetSender(rs)
	rec := &instanceRecord{InstanceID: "x", ContainerID: "cid"}
	d.instances["x"] = rec

	d.startExitWatch(rec)
	waitFor(t, 2*time.Second, func() bool { a, _ := f.counts(); return a == 1 }, "watcher never started")
	close(f.release)
	waitFor(t, 2*time.Second, func() bool {
		for _, e := range rs.snapshot() {
			if st := e.GetInstanceState(); st != nil && st.Status == panelv1.InstanceStatus_INSTANCE_STATUS_STOPPED {
				return true
			}
		}
		return false
	}, "terminal exit never produced a STOPPED InstanceState")
}