Stop saves pending trust changes and no longer waits on a hung RPC

handshake.go

@@ -13,6 +13,7 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 	"unicode"
 
@@ -83,6 +84,12 @@ const (
 	maxPendingHandshakes  = 256                    // cap pending (unapproved) handshake requests
 	maxPendingPerSource   = 16
 	maxJustificationLen   = 1024
+	// stopRPCWait bounds how long Stop waits for in-flight background
+	// RPCs before it saves and returns. A registry RPC can sit in its
+	// 30 s read deadline, and the daemon's whole shutdown budget is
+	// shorter than that, so Stop cannot wait them out. The deadline is
+	// set by the first Stop; later calls do not extend it.
+	stopRPCWait = 2 * time.Second
 	// pendingHandshakeTTL caps how long a pending (unapproved-by-
 	// operator) handshake request can sit in hm.pending before the
 	// reaper drops it. Without this, a request from a peer that the
@@ -115,9 +122,15 @@ type Manager struct {
 	stopping        bool                 // set under mu.Lock() before wg.Wait() in Stop
 	wg              sync.WaitGroup       // tracks background RPCs for clean shutdown
 	reapStop        chan struct{}        // signals replay reaper to stop
-	stopOnce        sync.Once            // ensures reapStop is closed only once
+	stopOnce        sync.Once            // first Stop: sets stopDeadline, starts rpcsDone, closes reapStop
+	stopDeadline    time.Time            // first Stop + stopRPCWait; no Stop waits for RPCs past it
+	rpcsDone        chan struct{}        // closed once wg drains after the first Stop
+	doneOnce        sync.Once            // ensures done is closed only once
 	dirty           chan struct{}        // buffered 1: non-blocking signal to drain goroutine
 	done            chan struct{}        // closed in Stop to signal drain goroutine exit
+	drainDone       chan struct{}        // closed when drainSaves returns
+	saveMu          sync.Mutex           // serializes saveTrust: concurrent saves would share trust.json.tmp
+	unsaved         atomic.Bool          // changed since the last successful save; set by markDirty
 
 	// Replay protection
 	replayMu   sync.Mutex
@@ -147,6 +160,7 @@ func NewManager(rt Runtime) *Manager {
 		trustWaiters:    make(map[uint32][]chan struct{}),
 		dirty:           make(chan struct{}, 1),
 		done:            make(chan struct{}),
+		drainDone:       make(chan struct{}),
 	}
 
 	if path := rt.IdentityPath(); path != "" {
@@ -160,7 +174,31 @@ func NewManager(rt Runtime) *Manager {
 	return hm
 }
 
-// Stop waits for all background RPCs to finish and stops the replay reaper.
+// Stop stops the replay reaper and the save goroutine, and saves trust
+// state one last time: every trust change made before it returns is
+// written, except one made by an RPC still running after the wait below,
+// or one whose write fails (logged).
+//
+// It first waits for in-flight background RPCs, because some of them
+// change trust (backfillPeerKey binds or drops a record), but only until
+// stopRPCWait after the first Stop: an RPC that is still running then
+// keeps running, and its change is saved by a later Stop, if any. Only
+// after that wait does it stop the drain goroutine, and it waits for that
+// goroutine to finish any save it has started before saving itself.
+// Closing done first, as Stop used to, lost changes: drainSaves picks
+// between a pending dirty signal and done at random, and nothing drained
+// the signals RPCs sent while Stop waited for them.
+//
+// The final save happens only if something changed since the last
+// successful save. A start followed by a stop leaves trust.json as it
+// was, including one that failed to parse at startup, which the next
+// real trust change overwrites but Stop on its own must not.
+//
+// Stop is safe to call more than once, concurrently included. The daemon
+// calls it twice (Daemon.Stop and the plugin runtime's Service.Stop).
+// Every call saves what is unsaved, so a second call picks up what RPCs
+// that outlived the first call's wait have changed since, but it does not
+// wait again.
 func (hm *Manager) Stop() {
 	// Mark stopping under mu so goRPC cannot race wg.Add vs wg.Wait.
 	// Any goRPC that holds mu.RLock() and sees stopping=false has already
@@ -169,16 +207,47 @@ func (hm *Manager) Stop() {
 	hm.stopping = true
 	hm.mu.Unlock()
 
-	// Both done and reapStop share stopOnce so a double-Stop (common in
-	// test cleanup paths that defer hm.Stop() while a parent also calls it)
-	// doesn't panic with "close of closed channel".
+	// The once makes a double-Stop (common in test cleanup paths that
+	// defer hm.Stop() while a parent also calls it) safe: reapStop is
+	// closed once, and every call shares one deadline and one waiter.
 	hm.stopOnce.Do(func() {
-		close(hm.done)
+		hm.stopDeadline = time.Now().Add(stopRPCWait)
+		hm.rpcsDone = make(chan struct{})
+		go func() {
+			hm.wg.Wait()
+			close(hm.rpcsDone)
+		}()
 		if hm.reapStop != nil {
 			close(hm.reapStop)
 		}
 	})
-	hm.wg.Wait()
+	if !hm.waitForRPCs() {
+		slog.Warn("handshake: stopping with background RPCs still running; only a later Stop saves what they change",
+			"waited", stopRPCWait)
+	}
+
+	hm.doneOnce.Do(func() { close(hm.done) })
+	<-hm.drainDone
+	hm.saveTrustIfUnsaved()
+}
+
+// waitForRPCs waits until the background RPCs have finished or the stop
+// deadline has passed, and reports whether they finished. Called by Stop
+// after stopOnce has set stopDeadline and rpcsDone.
+func (hm *Manager) waitForRPCs() bool {
+	select {
+	case <-hm.rpcsDone:
+		return true
+	default:
+	}
+	timer := time.NewTimer(time.Until(hm.stopDeadline))
+	defer timer.Stop()
+	select {
+	case <-hm.rpcsDone:
+		return true
+	case <-timer.C:
+		return false
+	}
 }
 
 // goRPC launches a tracked background goroutine.
@@ -333,14 +402,48 @@ type revokedSnapshotEntry struct {
 	Until  string `json:"until"` // RFC3339 timestamp of cooldown expiry
 }
 
+// saveTrust writes the current trust state to storePath.
 func (hm *Manager) saveTrust() {
 	if hm.storePath == "" {
 		return
 	}
 
+	// Saves run from drainSaves and from Stop, and Stop may run more than
+	// once at a time. AtomicWrite stages every write in the same
+	// trust.json.tmp, so two at once could rename each other's half-written
+	// file into place. Taken before mu: no caller holds mu here.
+	hm.saveMu.Lock()
+	defer hm.saveMu.Unlock()
+	hm.saveTrustLocked()
+}
+
+// saveTrustIfUnsaved is saveTrust, skipped when nothing has changed since
+// the last successful save. The check is made under saveMu, so a caller
+// that skips has waited out any save already running, and sees its
+// outcome: a save that failed leaves unsaved set, and this one retries.
+func (hm *Manager) saveTrustIfUnsaved() {
+	if hm.storePath == "" {
+		return
+	}
+	hm.saveMu.Lock()
+	defer hm.saveMu.Unlock()
+	if hm.unsaved.Load() {
+		hm.saveTrustLocked()
+	}
+}
+
+// saveTrustLocked does the work of saveTrust. Caller MUST hold saveMu.
+func (hm *Manager) saveTrustLocked() {
 	hm.mu.RLock()
 	defer hm.mu.RUnlock()
 
+	// Every change to persisted state (trusted, pending, revoked) sets
+	// unsaved under mu.Lock, so clearing it here, under mu.RLock and
+	// before the snapshot, cannot drop one: a change made before this
+	// point is in the snapshot, and one made after sets it again. A failed
+	// write sets it again too.
+	hm.unsaved.Store(false)
+
 	snap := trustSnapshot{}
 	for _, r := range hm.trusted {
 		snap.Trusted = append(snap.Trusted, trustSnapshotEntry{
@@ -377,11 +480,13 @@ func (hm *Manager) saveTrust() {
 
 	dir := filepath.Dir(hm.storePath)
 	if err := os.MkdirAll(dir, 0700); err != nil {
+		hm.unsaved.Store(true)
 		slog.Error("create trust state directory", "dir", dir, "err", err)
 		return
 	}
 
 	if err := fsutil.AtomicWrite(hm.storePath, data); err != nil {
+		hm.unsaved.Store(true)
 		slog.Error("write trust state", "err", err)
 		return
 	}
@@ -390,8 +495,11 @@ func (hm *Manager) saveTrust() {
 
 // drainSaves drains the dirty channel and persists trust state
 // asynchronously, so that AtomicWrite (including fsync) does not
-// block other handshake operations holding hm.mu.
+// block other handshake operations holding hm.mu. It returns once done
+// is closed, possibly with a dirty signal still pending; Stop waits on
+// drainDone and then saves whatever is still unsaved.
 func (hm *Manager) drainSaves() {
+	defer close(hm.drainDone)
 	for {
 		select {
 		case <-hm.dirty:
@@ -402,9 +510,11 @@ func (hm *Manager) drainSaves() {
 	}
 }
 
-// markDirty signals the drain goroutine that trust state needs
-// persistence. Caller MUST hold hm.mu (write-locked).
+// markDirty records that trust state has changed since the last
+// successful save and signals the drain goroutine to persist it. Caller
+// MUST hold hm.mu (write-locked).
 func (hm *Manager) markDirty() {
+	hm.unsaved.Store(true)
 	select {
 	case hm.dirty <- struct{}{}:
 	default: