a73x

internal/agent/serialpump/serialpump.go

Ref:   Size: 11.9 KiB   History

// Package serialpump owns the durability of VM serial consoles. A backend's
// ConsoleSource opens the raw byte stream (cloud-hypervisor serves it on a
// unix socket, one client at a time — a new connection kicks the old one —
// and depending on CH version, output written while no client is connected is
// dropped (older) or buffered in a bounded replay ring (current)); the pump is
// that one client, always: a supervised goroutine per running VM opens the
// stream, drains continuously into a bounded in-memory ring (backlog for new
// viewers) and a capped on-disk serial.log (survives agent restart, unbounded
// history), fans live bytes out to attached console viewers, and forwards
// viewer input back to the stream. A slow viewer is dropped, never allowed to
// stall the drain. For cloud-hypervisor's socket shape specifically, do NOT
// connect other clients (socat etc.) — they would steal the line from the
// pump.
package serialpump

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"sync"
	"time"
)

// ConsoleSource opens the guest console for vmID. Implementations are
// per-backend: cloud-hypervisor serves the serial line on a unix socket, vfkit
// hands out a PTY. Open is called inside the pump's reconnect loop, so it must
// be cheap and safely retryable, and it may fail while the VM is down. A nil
// error means a usable, non-nil stream — the pump closes it unconditionally
// when the drain ends.
type ConsoleSource interface {
	Open(vmID string) (io.ReadWriteCloser, error)
}

const (
	defaultRingMax = 256 << 10 // 256 KiB — enough for a boot log
	defaultLogMax  = 4 << 20   // 4 MiB on disk, then rotate once to .old
	viewerDepth    = 64        // live-tail channel depth before a viewer is dropped
)

// Manager runs one Pump per VM. The console source and log path are injected so
// the package stays a leaf (no dependency on the agent's state store, and no
// knowledge of how any backend exposes its console).
type Manager struct {
	src     ConsoleSource
	logPath func(vmID string) string
	ringMax int
	logMax  int64

	mu    sync.Mutex
	pumps map[string]*pump
}

// NewManager returns a Manager opening each VM's console through src and
// resolving its on-disk log through logPath.
func NewManager(src ConsoleSource, logPath func(vmID string) string) *Manager {
	return &Manager{
		src:     src,
		logPath: logPath,
		ringMax: defaultRingMax,
		logMax:  defaultLogMax,
		pumps:   map[string]*pump{},
	}
}

// Ensure starts the VM's pump if it is not already running. Idempotent.
func (m *Manager) Ensure(vmID string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if p, ok := m.pumps[vmID]; ok {
		// Pump already running. A pump outlives its console only when the
		// hypervisor went away on its own — a crash, or a guest that powered
		// itself off — since a stop the fleet asked for stops the pump too. The
		// fleet restarts such a VM through the backend's Boot, which lands here
		// on a pump possibly parked deep in dial backoff (up to 30s); the fresh
		// socket must not wait that out, because unconsumed early boot output is
		// dropped or truncated by CH. Poke the dial loop to retry now. The ring
		// rides across that one restart, and should: its last words are what the
		// guest said as it died, which is what a viewer comes for.
		select {
		case p.poke <- struct{}{}:
		default: // a poke is already pending
		}
		return
	}
	p := &pump{
		vmID:    vmID,
		src:     m.src,
		logPath: m.logPath(vmID),
		ringMax: m.ringMax,
		logMax:  m.logMax,
		viewers: map[int]chan []byte{},
		done:    make(chan struct{}),
		poke:    make(chan struct{}, 1),
	}
	m.pumps[vmID] = p
	go p.run()
}

// Stop tears down the VM's pump: the guest was powered off or destroyed, so
// its console goes with it — the ring is released and every attached viewer is
// ended. The next Ensure (a power-on) builds a pump with an empty ring, which
// is what keeps a dead boot's output from answering for a live guest. No-op for
// unknown VMs.
func (m *Manager) Stop(vmID string) {
	m.mu.Lock()
	p := m.pumps[vmID]
	delete(m.pumps, vmID)
	m.mu.Unlock()
	if p != nil {
		p.stop()
	}
}

// StopAll tears down every pump (agent shutdown, tests).
func (m *Manager) StopAll() {
	m.mu.Lock()
	pumps := m.pumps
	m.pumps = map[string]*pump{}
	m.mu.Unlock()
	for _, p := range pumps {
		p.stop()
	}
}

// Attach bridges rw to the VM's console: onReady (may be nil) fires after
// validation and before any bytes — the caller uses it to send its protocol
// reply — then the ring backlog is replayed, then live output flows; bytes
// read from rw are forwarded to the guest as input. Blocks until ctx ends,
// rw errors, or the viewer is dropped as too slow. Returns an error
// immediately (before onReady) when the VM has no pump. The caller must close
// rw (unblocking its Read) once Attach returns, or the input-forwarding
// goroutine leaks parked in rw.Read.
func (m *Manager) Attach(ctx context.Context, vmID string, rw io.ReadWriter, onReady func() error) error {
	m.mu.Lock()
	p := m.pumps[vmID]
	m.mu.Unlock()
	if p == nil {
		return fmt.Errorf("no console for vm %q (powered off, or not on this host?)", vmID)
	}
	if onReady != nil {
		if err := onReady(); err != nil {
			return err
		}
	}
	return p.attach(ctx, rw)
}

// pump drains one VM's console. It reopens forever (the VMM restarts on VM
// stop/start; the console may not exist yet at boot) until stop() is called.
type pump struct {
	vmID    string
	src     ConsoleSource
	logPath string
	ringMax int
	logMax  int64

	mu        sync.Mutex
	ring      []byte
	conn      io.ReadWriteCloser // current console stream; input writes go here
	viewers   map[int]chan []byte
	nextID    int
	logF      *os.File
	logSize   int64
	logWarned bool // one breadcrumb per pump when the log path fails

	done     chan struct{}
	poke     chan struct{} // buffered(1); re-Ensure nudges a backed-off dial loop
	stopOnce sync.Once
}

func (p *pump) stop() {
	p.stopOnce.Do(func() {
		close(p.done)
		p.mu.Lock()
		if p.conn != nil {
			p.conn.Close() // unblock the drain read
		}
		for id, ch := range p.viewers {
			close(ch)
			delete(p.viewers, id)
		}
		if p.logF != nil {
			p.logF.Close()
			p.logF = nil
		}
		p.mu.Unlock()
	})
}

func (p *pump) run() {
	backoff := 250 * time.Millisecond
	for {
		select {
		case <-p.done:
			return
		default:
		}
		conn, err := p.src.Open(p.vmID)
		if err != nil {
			select {
			case <-p.done:
				return
			case <-time.After(backoff):
				if backoff < 30*time.Second {
					backoff *= 2
				}
			case <-p.poke:
				// Re-Ensure of a live pump (VM restarted, fresh socket):
				// retry immediately with a fresh backoff instead of waiting
				// out a stale one.
				backoff = 250 * time.Millisecond
			}
			continue
		}
		backoff = 250 * time.Millisecond
		p.mu.Lock()
		select {
		case <-p.done:
			// stop() ran between Open and here: it closed p.conn (nil at that
			// point) but not THIS conn — do not resurrect the pump.
			p.mu.Unlock()
			conn.Close()
			return
		default:
		}
		p.conn = conn
		p.mu.Unlock()
		p.drain(conn)
		p.mu.Lock()
		if p.conn == conn {
			p.conn = nil
		}
		p.mu.Unlock()
		conn.Close()
	}
}

func (p *pump) drain(conn io.Reader) {
	buf := make([]byte, 4096)
	for {
		n, err := conn.Read(buf)
		if n > 0 {
			p.publish(buf[:n])
		}
		if err != nil {
			return
		}
	}
}

// publish appends b to the ring (trimming from the front past ringMax), the
// on-disk log (rotating once at logMax), and every viewer. A viewer whose
// channel is full is dropped (closed + removed) — it must never stall the
// drain. b is copied: callers reuse their read buffer. The disk write happens
// under p.mu deliberately: the only contenders are keystrokes and attaches,
// and drain is the sole caller — a writer goroutine would buy complexity, not
// throughput.
func (p *pump) publish(b []byte) {
	cp := make([]byte, len(b))
	copy(cp, b)

	p.mu.Lock()
	defer p.mu.Unlock()

	select {
	case <-p.done:
		// stop() may have run while this publish waited on the lock. It already
		// closed the log file and viewer channels; touching them now would
		// reopen the log fd with nothing left to close it (run() is exiting).
		return
	default:
	}

	p.ring = append(p.ring, cp...)
	if over := len(p.ring) - p.ringMax; over > 0 {
		p.ring = p.ring[over:]
	}
	p.appendLogLocked(cp)
	for id, ch := range p.viewers {
		select {
		case ch <- cp:
		default:
			close(ch) // slow viewer: drop it, never block the pump
			delete(p.viewers, id)
		}
	}
}

// appendLogLocked writes to the capped on-disk log, rotating once (.old) at
// the cap. Log failures are swallowed — the console must keep working even if
// the disk is unhappy (the ring still serves backlog) — but breadcrumbed once
// per pump, so "history survives agent restart" going false is visible.
func (p *pump) appendLogLocked(b []byte) {
	if p.logF == nil {
		f, err := os.OpenFile(p.logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
		if err != nil {
			p.warnLogLocked("open", err)
			return
		}
		st, err := f.Stat()
		if err != nil {
			f.Close()
			p.warnLogLocked("stat", err)
			return
		}
		p.logF, p.logSize = f, st.Size()
	}
	if p.logSize+int64(len(b)) > p.logMax {
		p.logF.Close()
		p.logF = nil
		if err := os.Rename(p.logPath, p.logPath+".old"); err != nil {
			p.warnLogLocked("rotate", err)
			return
		}
		f, err := os.OpenFile(p.logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
		if err != nil {
			p.warnLogLocked("reopen", err)
			return
		}
		p.logF, p.logSize = f, 0
	}
	n, err := p.logF.Write(b)
	if err != nil {
		p.warnLogLocked("write", err)
	}
	p.logSize += int64(n)
}

// warnLogLocked emits one warning per pump lifetime: log failures are
// tolerated, but silently losing restart-surviving history is not.
func (p *pump) warnLogLocked(op string, err error) {
	if p.logWarned {
		return
	}
	p.logWarned = true
	slog.Warn("serial log persistence failing; console history will not survive agent restart",
		"op", op, "path", p.logPath, "err", err)
}

// errPumpStopped reports an attach racing a Stop (VM destroyed): the viewer
// must error out, not hang on a channel nothing will ever publish to or close.
var errPumpStopped = errors.New("console pump stopped")

// subscribe atomically snapshots the ring and registers a live channel — one
// lock, so no byte can fall between backlog and live, and stop() (which also
// takes the lock to close all viewer channels) cannot interleave: either we
// see done closed here, or stop sees our channel in the map and closes it.
func (p *pump) subscribe() (backlog []byte, ch chan []byte, cancel func(), err error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	select {
	case <-p.done:
		return nil, nil, nil, errPumpStopped
	default:
	}
	backlog = make([]byte, len(p.ring))
	copy(backlog, p.ring)
	ch = make(chan []byte, viewerDepth)
	id := p.nextID
	p.nextID++
	p.viewers[id] = ch
	return backlog, ch, func() {
		p.mu.Lock()
		defer p.mu.Unlock()
		if c, ok := p.viewers[id]; ok {
			close(c)
			delete(p.viewers, id)
		}
	}, nil
}

// writeInput forwards viewer keystrokes to the guest. Dropped silently when
// the console stream is not currently connected (VM stopped): the serial line
// simply isn't there, exactly like typing into an unplugged terminal.
func (p *pump) writeInput(b []byte) {
	p.mu.Lock()
	conn := p.conn
	p.mu.Unlock()
	if conn != nil {
		_, _ = conn.Write(b)
	}
}

func (p *pump) attach(ctx context.Context, rw io.ReadWriter) error {
	backlog, ch, cancel, err := p.subscribe()
	if err != nil {
		return err
	}
	defer cancel()
	if len(backlog) > 0 {
		if _, err := rw.Write(backlog); err != nil {
			return err
		}
	}
	// Input pump: rw → guest. Ends when rw read errors (stream closed);
	// cancel() then makes the output loop below observe the closed channel.
	go func() {
		buf := make([]byte, 1024)
		for {
			n, err := rw.Read(buf)
			if n > 0 {
				p.writeInput(buf[:n])
			}
			if err != nil {
				cancel()
				return
			}
		}
	}()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case b, ok := <-ch:
			if !ok {
				return fmt.Errorf("console viewer detached (slow reader, stream closed, or pump stopped)")
			}
			if _, err := rw.Write(b); err != nil {
				return err
			}
		}
	}
}