a73x

internal/agent/serialpump/serialpump_test.go

Ref:   Size: 17.4 KiB   History

package serialpump

import (
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// Every wait in this file is bounded. A pump that stops forwarding bytes must
// report as a failed assertion in seconds; an unbounded read or write turns one
// broken claim into a ten-minute suite-wide panic that names no claim at all.
const ioTimeout = 5 * time.Second

// recvWithin receives one value from ch, failing with what if none arrives.
//
//nolint:ireturn // the type parameter IS the channel's element type (conn, error, signal)
func recvWithin[T any](t *testing.T, ch <-chan T, what string) T {
	t.Helper()
	select {
	case v := <-ch:
		return v
	case <-time.After(ioTimeout):
		t.Fatal(what)
		var zero T
		return zero
	}
}

// fakeCH is a unix-socket listener standing in for cloud-hypervisor's
// --serial socket=. It records input written by the pump and lets tests
// emit guest output.
type fakeCH struct {
	ln    net.Listener
	conns chan net.Conn
}

func newFakeCH(t *testing.T, sock string) *fakeCH {
	t.Helper()
	ln, err := net.Listen("unix", sock)
	require.NoError(t, err)
	f := &fakeCH{ln: ln, conns: make(chan net.Conn, 4)}
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			f.conns <- c
		}
	}()
	t.Cleanup(func() { ln.Close() })
	return f
}

func (f *fakeCH) conn(t *testing.T) net.Conn {
	t.Helper()
	return recvWithin(t, f.conns, "pump never dialed the serial socket")
}

// testSource dials the unix socket path returned by socketPath, standing in
// for a backend's ConsoleSource (e.g. cloudhv.ConsoleSource) in these tests.
type testSource func(vmID string) string

func (s testSource) Open(vmID string) (io.ReadWriteCloser, error) {
	return net.Dial("unix", s(vmID))
}

// waitConnected blocks until the pump has recorded its console stream. A
// listener accepting is NOT that moment: the pump stores p.conn only after
// Open returns, and viewer input arriving before then is dropped by design
// (typing into an unplugged terminal). Any test that forwards keystrokes must
// synchronise on the pump, not on the socket.
func waitConnected(t *testing.T, m *Manager, vmID string) {
	t.Helper()
	require.Eventually(t, func() bool {
		m.mu.Lock()
		p := m.pumps[vmID]
		m.mu.Unlock()
		if p == nil {
			return false
		}
		p.mu.Lock()
		defer p.mu.Unlock()
		return p.conn != nil
	}, ioTimeout, time.Millisecond, "pump never recorded a console stream")
}

func newTestManager(t *testing.T, dir string) *Manager {
	t.Helper()
	m := NewManager(
		testSource(func(vmID string) string { return filepath.Join(dir, vmID+".serial.sock") }),
		func(vmID string) string { return filepath.Join(dir, vmID+".serial.log") },
	)
	t.Cleanup(m.StopAll)
	return m
}

// pipeViewer returns an in-memory io.ReadWriter viewer plus the far ends the
// test uses to observe output and inject keystrokes.
func pipeViewer() (viewer io.ReadWriter, out io.Reader, in io.Writer) {
	or, ow := io.Pipe() // pump → viewer output
	ir, iw := io.Pipe() // test → viewer input
	type rw struct {
		io.Reader
		io.Writer
	}
	return rw{ir, ow}, or, iw
}

// readN reads exactly n bytes, failing the test if they never come. The read
// runs in its own goroutine because the viewer end is an io.Pipe, which has no
// deadline to set.
func readN(t *testing.T, r io.Reader, n int) []byte {
	t.Helper()
	buf := make([]byte, n)
	done := make(chan error, 1)
	go func() { _, err := io.ReadFull(r, buf); done <- err }()
	require.NoError(t, recvWithin(t, done, fmt.Sprintf("timed out reading %d bytes", n)))
	return buf
}

// writeAll writes b, failing rather than parking forever when nothing drains
// the far end — a stalled pump must not stall the test (the 1 MiB flood in
// TestSlowViewerIsDroppedNotBlocking would otherwise fill the socket buffer
// and block for good).
func writeAll(t *testing.T, w io.Writer, b []byte) {
	t.Helper()
	if c, ok := w.(net.Conn); ok {
		require.NoError(t, c.SetWriteDeadline(time.Now().Add(ioTimeout)))
		defer func() { _ = c.SetWriteDeadline(time.Time{}) }()
		_, err := c.Write(b)
		require.NoError(t, err)
		return
	}
	done := make(chan error, 1)
	go func() { _, err := w.Write(b); done <- err }()
	require.NoError(t, recvWithin(t, done, fmt.Sprintf("timed out writing %d bytes", len(b))))
}

func TestBacklogReplayThenLive(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	ch := newFakeCH(t, filepath.Join(dir, "vm1.serial.sock"))
	m.Ensure("vm1")
	guest := ch.conn(t)

	writeAll(t, guest, []byte("BOOT-LOG\n"))

	viewer, out, _ := pipeViewer()
	errc := make(chan error, 1)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() { errc <- m.Attach(ctx, "vm1", viewer, nil) }()

	// Backlog written before attach is replayed first...
	assert.Equal(t, "BOOT-LOG\n", string(readN(t, out, 9)))
	// ...then live bytes flow.
	writeAll(t, guest, []byte("LIVE\n"))
	assert.Equal(t, "LIVE\n", string(readN(t, out, 5)))
	cancel()
	assert.ErrorIs(t, recvWithin(t, errc, "Attach never returned after cancel"), context.Canceled)
}

func TestInputForwardedToSocket(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	ch := newFakeCH(t, filepath.Join(dir, "vm1.serial.sock"))
	m.Ensure("vm1")
	guest := ch.conn(t)
	waitConnected(t, m, "vm1") // keystrokes sent before this are dropped, not queued

	viewer, _, in := pipeViewer()
	ctx := t.Context()
	go m.Attach(ctx, "vm1", viewer, nil) //nolint:errcheck

	writeAll(t, in, []byte("ls\r"))
	assert.Equal(t, "ls\r", string(readN(t, guest, 3)))
}

func TestAttachUnknownVMErrorsBeforeOnReady(t *testing.T) {
	m := newTestManager(t, t.TempDir())
	called := false
	err := m.Attach(context.Background(), "nope", nil, func() error { called = true; return nil })
	assert.Error(t, err)
	assert.False(t, called, "onReady must not fire when the VM has no pump")
}

func TestOnReadyFiresBeforeBacklog(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	ch := newFakeCH(t, filepath.Join(dir, "vm1.serial.sock"))
	m.Ensure("vm1")
	guest := ch.conn(t)
	writeAll(t, guest, []byte("X"))

	viewer, out, _ := pipeViewer()
	ready := make(chan struct{})
	ctx := t.Context()
	go m.Attach(ctx, "vm1", viewer, func() error { close(ready); return nil }) //nolint:errcheck
	// onReady before any viewer write (protocol: reply frame precedes raw bytes).
	recvWithin(t, ready, "onReady never fired")
	assert.Equal(t, "X", string(readN(t, out, 1)))
}

func TestRingIsBounded(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	m.ringMax = 16 // shrink for the test
	ch := newFakeCH(t, filepath.Join(dir, "vm1.serial.sock"))
	m.Ensure("vm1")
	guest := ch.conn(t)

	writeAll(t, guest, []byte("0123456789ABCDEFGHIJ")) // 20 bytes > 16
	// Wait until the pump has drained all 20 bytes into the log.
	log := filepath.Join(dir, "vm1.serial.log")
	require.Eventually(t, func() bool {
		b, _ := os.ReadFile(log)
		return len(b) == 20
	}, 5*time.Second, 10*time.Millisecond)

	viewer, out, _ := pipeViewer()
	ctx := t.Context()
	go m.Attach(ctx, "vm1", viewer, nil) //nolint:errcheck
	// Backlog is only the LAST 16 bytes.
	assert.Equal(t, "456789ABCDEFGHIJ", string(readN(t, out, 16)))
}

func TestOnDiskLogRotatesAtCap(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	m.logMax = 32
	ch := newFakeCH(t, filepath.Join(dir, "vm1.serial.sock"))
	m.Ensure("vm1")
	guest := ch.conn(t)

	big := make([]byte, 40)
	for i := range big {
		big[i] = 'a'
	}
	writeAll(t, guest, big)
	require.Eventually(t, func() bool {
		_, err := os.Stat(filepath.Join(dir, "vm1.serial.log.old"))
		return err == nil
	}, 5*time.Second, 10*time.Millisecond, "log must rotate to .old at the cap")
}

// slowViewer consumes output correctly but slowly: Write sleeps then succeeds.
// (A never-reading viewer would park Attach inside rw.Write, where it could
// not observe the drop — the pump drops via the CHANNEL, so the viewer must
// keep returning from Write to come back to the channel receive.)
type slowViewer struct{ done chan struct{} }

func (v slowViewer) Read(p []byte) (int, error) { <-v.done; return 0, io.EOF }
func (v slowViewer) Write(p []byte) (int, error) {
	time.Sleep(3 * time.Millisecond)
	return len(p), nil
}

func TestSlowViewerIsDroppedNotBlocking(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	ch := newFakeCH(t, filepath.Join(dir, "vm1.serial.sock"))
	m.Ensure("vm1")
	guest := ch.conn(t)

	v := slowViewer{done: make(chan struct{})}
	defer close(v.done)
	errc := make(chan error, 1)
	go func() { errc <- m.Attach(context.Background(), "vm1", v, nil) }()

	// Flood: 256 chunks of 4096 bytes. The pump's read buffer is 4096, so
	// coalescing cannot reduce this below 256 distinct publishes — far past
	// viewerDepth (64) — while the viewer consumes only ~1 per 3ms. The
	// channel overflows, the pump drops the viewer, and Attach's next channel
	// receive observes the close and errors out.
	junk := make([]byte, 4096)
	for range 256 {
		writeAll(t, guest, junk)
	}
	select {
	case err := <-errc:
		assert.Error(t, err, "slow viewer is dropped with an error")
	case <-time.After(10 * time.Second):
		t.Fatal("slow viewer was never dropped")
	}
	// Pump still healthy after the drop: fresh output still lands in the log.
	log := filepath.Join(dir, "vm1.serial.log")
	prev, _ := os.Stat(log)
	writeAll(t, guest, []byte("still-draining"))
	require.Eventually(t, func() bool {
		st, err := os.Stat(log)
		return err == nil && (prev == nil || st.Size() > prev.Size())
	}, 5*time.Second, 10*time.Millisecond)
}

func TestFanOutToTwoViewers(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	ch := newFakeCH(t, filepath.Join(dir, "vm1.serial.sock"))
	m.Ensure("vm1")
	guest := ch.conn(t)

	v1, out1, _ := pipeViewer()
	v2, out2, _ := pipeViewer()
	ctx1, cancel1 := context.WithCancel(context.Background())
	defer cancel1()
	ctx2 := t.Context()
	go m.Attach(ctx1, "vm1", v1, nil) //nolint:errcheck
	go m.Attach(ctx2, "vm1", v2, nil) //nolint:errcheck

	// Both viewers see the same bytes (whichever attached later gets them via
	// the ring replay — same contract either way).
	writeAll(t, guest, []byte("BOTH"))
	assert.Equal(t, "BOTH", string(readN(t, out1, 4)))
	assert.Equal(t, "BOTH", string(readN(t, out2, 4)))

	// One viewer detaching must not disturb the other.
	cancel1()
	writeAll(t, guest, []byte("SOLO"))
	assert.Equal(t, "SOLO", string(readN(t, out2, 4)))
}

func TestSubscribeAfterStopErrors(t *testing.T) {
	// The Attach-vs-Stop race: Manager.Attach can look a pump up just before
	// Stop() runs. subscribe must then refuse (not register a viewer channel
	// nothing will ever close) so the console errors instead of hanging.
	p := &pump{viewers: map[int]chan []byte{}, done: make(chan struct{})}
	p.stop()
	_, _, _, err := p.subscribe()
	assert.ErrorIs(t, err, errPumpStopped)
}

func TestPumpReconnectsAfterSocketRestart(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	sock := filepath.Join(dir, "vm1.serial.sock")
	ch := newFakeCH(t, sock)
	m.Ensure("vm1")
	guest := ch.conn(t)
	writeAll(t, guest, []byte("A"))
	// Listener FIRST, then conn: closing the conn first lets the pump re-dial
	// into the still-live old listener, parking it on a conn ch2 never sees.
	ch.ln.Close()
	guest.Close()
	// Go's UnixListener unlinks its socket file on Close, so the path is
	// usually gone already — this remove only guards against a leftover file.
	if err := os.Remove(sock); err != nil && !os.IsNotExist(err) {
		t.Fatal(err)
	}

	// CH restarts (VM stop/start): a new listener appears; pump must re-dial.
	ch2 := newFakeCH(t, sock)
	guest2 := ch2.conn(t) // blocks until the pump reconnects
	writeAll(t, guest2, []byte("B"))
}

// TestEnsurePokesBackedOffPump pins the re-Ensure nudge: a pump whose
// hypervisor died on its own — a crash, a guest that powered itself off —
// outlives its console, so when the fleet restarts that VM, Ensure finds an
// existing pump possibly parked deep in dial backoff. The poke must make it
// re-dial immediately, or the fresh socket sits unconsumed for up to 30s and
// early boot output is lost.
//
// Timing schedule (every wait is a lower bound — time.After never fires
// early, so scheduling jitter only pushes dials LATER):
//
//	unpoked dials with no listener land at ~0 / 0.25 / 0.75 / 1.75 / 3.75s
//	(backoff 0.25 → 0.5 → 1 → 2 → 4s); after the ~3.75s failure the pump is
//	parked in a 4s wait, so the next UNPOKED dial cannot happen before ~7.75s.
//
// We sleep 4.2s — the ~3.75s dial has ~450ms of jitter margin to fail before
// the listener exists — then create the listener and re-Ensure (the poke).
// Requiring a connection within 2s (~6.2s total) leaves ~1.5s of margin below
// the ~7.75s unpoked dial: only the poke can connect that early. (If the
// ~3.75s dial were somehow delayed past the sleep, the poke is buffered and
// fires the moment the pump enters its next wait — still well within bound.)
func TestEnsurePokesBackedOffPump(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	sock := filepath.Join(dir, "vm1.serial.sock")

	m.Ensure("vm1") // no listener yet: pump enters its dial-backoff loop
	time.Sleep(4200 * time.Millisecond)

	ch := newFakeCH(t, sock)
	m.Ensure("vm1") // idempotent — but must poke the backed-off dial loop

	select {
	case <-ch.conns:
		// poked pump re-dialed immediately
	case <-time.After(2 * time.Second):
		t.Fatal("pump did not re-dial after Ensure poke; still waiting out a stale backoff")
	}
}

// TestPowerOffDropsTheRingAndKeepsTheLog is the pump-lifecycle invariant read
// from the console: a guest that has been powered off has nothing to replay.
// The backend stops the pump when it stops the guest, so the ring — the backlog
// every new viewer is handed on attach — goes with it, and the next boot's
// viewer reads that boot alone. Without this, the first boot's login prompt
// answers for the second, which is the whole basis of the smoke's power-cycle
// boot proof. The history is not lost with the ring: serial.log is one file per
// VM, appended across pump lifetimes.
func TestPowerOffDropsTheRingAndKeepsTheLog(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	sock := filepath.Join(dir, "vm1.serial.sock")
	logPath := filepath.Join(dir, "vm1.serial.log")

	ch := newFakeCH(t, sock)
	m.Ensure("vm1")
	guest := ch.conn(t)
	writeAll(t, guest, []byte("BOOT-1 login: "))
	// The first boot must be IN the ring at the moment of the stop, or the test
	// proves nothing about dropping it.
	require.Eventually(t, func() bool {
		b, _ := os.ReadFile(logPath)
		return string(b) == "BOOT-1 login: "
	}, ioTimeout, time.Millisecond, "pump never drained the first boot")

	// Power off. Listener before conn, as in TestPumpReconnectsAfterSocketRestart.
	m.Stop("vm1")
	ch.ln.Close()
	guest.Close()
	if err := os.Remove(sock); err != nil && !os.IsNotExist(err) {
		t.Fatal(err)
	}
	assert.Error(t, m.Attach(t.Context(), "vm1", nil, nil),
		"a powered-off guest has no console to attach to")

	// Power on: a fresh socket, and Ensure builds a fresh pump on it.
	ch2 := newFakeCH(t, sock)
	m.Ensure("vm1")
	guest2 := ch2.conn(t)
	writeAll(t, guest2, []byte("BOOT-2"))

	viewer, out, _ := pipeViewer()
	go m.Attach(t.Context(), "vm1", viewer, nil) //nolint:errcheck
	assert.Equal(t, "BOOT-2", string(readN(t, out, 6)),
		"the dead boot must not be replayed to the live one's viewer")

	require.Eventually(t, func() bool {
		b, _ := os.ReadFile(logPath)
		return string(b) == "BOOT-1 login: BOOT-2"
	}, ioTimeout, time.Millisecond, "serial.log must keep the history the ring dropped")
}

// TestPowerOffEndsAttachedViewers pins the other half of Stop's contract. A
// viewer watching a VM that powers off must be ended, not left parked on a
// channel nothing will ever publish to — its console is genuinely gone, and the
// error is what tells the far end (a browser, the smoke's tail) to re-dial.
func TestPowerOffEndsAttachedViewers(t *testing.T) {
	dir := t.TempDir()
	m := newTestManager(t, dir)
	ch := newFakeCH(t, filepath.Join(dir, "vm1.serial.sock"))
	m.Ensure("vm1")
	guest := ch.conn(t)

	viewer, out, _ := pipeViewer()
	errc := make(chan error, 1)
	go func() { errc <- m.Attach(context.Background(), "vm1", viewer, nil) }()
	writeAll(t, guest, []byte("LIVE\n"))
	assert.Equal(t, "LIVE\n", string(readN(t, out, 5))) // attached and flowing

	m.Stop("vm1")

	assert.Error(t, recvWithin(t, errc, "Attach never returned after the guest powered off"))
}

func TestStopUnknownVMIsNoop(t *testing.T) {
	m := newTestManager(t, t.TempDir())
	m.Stop("never-started") // must not panic
}

// fakeSource hands out an in-memory console stream and records how many
// times Open was called, pinning that the pump reconnects through the seam.
type fakeSource struct {
	mu    sync.Mutex
	opens int
	conn  io.ReadWriteCloser
	err   error
}

func (f *fakeSource) Open(vmID string) (io.ReadWriteCloser, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.opens++
	if f.err != nil {
		return nil, f.err
	}
	return f.conn, nil
}

func (f *fakeSource) count() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.opens
}

func TestManagerOpensConsoleThroughSource(t *testing.T) {
	guest, host := net.Pipe() // host end stands in for any ReadWriteCloser
	src := &fakeSource{conn: host}
	m := NewManager(src, func(vmID string) string {
		return filepath.Join(t.TempDir(), "serial.log")
	})
	defer m.StopAll()

	m.Ensure("vm-1")
	go func() { _, _ = guest.Write([]byte("hello console")) }()

	deadline := time.After(2 * time.Second)
	for {
		if src.count() > 0 {
			return
		}
		select {
		case <-deadline:
			t.Fatal("pump never called ConsoleSource.Open")
		case <-time.After(10 * time.Millisecond):
		}
	}
}