daemon: shutdown waits for a handshake poll in flight; follow-ups to #500

CHANGELOG.md

@@ -285,6 +285,21 @@ Detailed per-release notes are on the
     `~/.pilot/update-state.json`.
 
 ### Fixed
+- **A daemon stopping mid-poll no longer loses relayed handshakes.** The
+  registry empties a node's handshake inbox as it answers a poll, and
+  shutdown closed the registry client without waiting for a poll in flight,
+  so a restart during one dropped the requests and approvals it carried.
+  Shutdown now waits for it before stopping the handshake manager, so what it
+  brings back is still acted on and saved, and starts no new poll once it
+  has begun. The wait shares the 5s that shutdown already allows its
+  background goroutines.
+- **A full table of handshake windows lets go of one settled peer, not all
+  of them.** With 64 peers tracked, a new request evicted every peer whose
+  fast-poll window had closed, and those peers lost the hold-off that stops
+  an automatic handshake from reopening their window early. Now only the
+  one whose hold-off ends soonest makes room.
+- **Bytes after a known beacon notify kind are ignored** instead of the
+  whole notify being dropped, so a later beacon can extend the message.
 - **Stream segments fit one packet.** A full stream segment was 4096 bytes,
   about 4.2 KB on the wire and three IP fragments on a 1500-byte path. NATs,
   firewalls and some virtual networks drop fragments, so on those paths

pkg/daemon/daemon.go

@@ -1634,6 +1634,10 @@ func (d *Daemon) Stop() error {
 }
 
 func (d *Daemon) doStop() {
+	// The background goroutines and a relayed-handshake poll in flight
+	// share one deadline.
+	drainDeadline := time.Now().Add(5 * time.Second)
+
 	// Wait for all daemon-scoped background goroutines to notice
 	// stopCh and exit before tearing down shared infrastructure.
 	// Use a 5-second timeout to prevent a hung goroutine (e.g.
@@ -1646,7 +1650,7 @@ func (d *Daemon) doStop() {
 	}()
 	select {
 	case <-done:
-	case <-time.After(5 * time.Second):
+	case <-time.After(time.Until(drainDeadline)):
 		slog.Warn("timed out waiting for background goroutines to exit", "leaked", true)
 	}
 
@@ -1702,6 +1706,13 @@ func (d *Daemon) doStop() {
 		slog.Info("closed active connections", "count", len(conns))
 	}
 
+	// A relayed-handshake poll in flight finishes first, while the handshake
+	// manager can still act on and save what it brings back, and no new one
+	// starts: the registry empties the inbox as it answers.
+	if !d.hsPoll.shutdown(drainDeadline) {
+		slog.Warn("relayed-handshake poll still running at shutdown; stopping without its results")
+	}
+
 	// Wait for background handshake RPCs to drain
 	if d.handshakes != nil {
 		d.handshakes.Stop()
@@ -6565,10 +6576,8 @@ func (d *Daemon) lookupPeerPubKey(nodeID uint32) (ed25519.PublicKey, error) {
 }
 
 // pollRelayedHandshakes checks the registry for handshake requests and
-// responses relayed to this node and processes them.
-//
-// timeout > 0 bounds the registry call (a local client is waiting on it);
-// 0 leaves it unbounded, as the background loop always ran it.
+// responses relayed to this node and processes them. It runs on the poll
+// goroutine started by pollHandshakes.
 func (d *Daemon) pollRelayedHandshakes() {
 	rc, nodeID := d.reg(), d.NodeID()
 	if rc == nil {

pkg/daemon/handshakepoll.go

@@ -23,8 +23,9 @@ import (
 //     polls every handshakeFastPollInterval until the peer answers, becomes
 //     trusted, or handshakeFastPollWindow has passed.
 //   - On demand: a local client asking for pending requests or the trust
-//     list, approving, rejecting or waiting for trust triggers one poll
-//     first, at most one per handshakeOnDemandGap.
+//     list, approving or rejecting triggers one poll first, at most one per
+//     handshakeOnDemandGap. Waiting for trust does too, but only while a
+//     request this node sent that peer is unanswered.
 //   - Poke: the beacon can tell this node that something is waiting for it
 //     at the registry (beaconMsgNotify). A poke triggers one poll, from a
 //     small token bucket, and never sooner than handshakeOnDemandGap after
@@ -103,6 +104,9 @@ type handshakePollSched struct {
 	// meanwhile waits for that one instead of starting another.
 	running chan struct{}
 
+	// closed is set by shutdown: no poll starts after it.
+	closed bool
+
 	// run performs one poll. It is d.pollRelayedHandshakes; tests replace it.
 	run func()
 
@@ -150,14 +154,21 @@ func (s *handshakePollSched) requestSent(peer uint32, explicit bool) {
 	case !tracked && len(s.waiting) >= handshakeMaxWaiting:
 		s.pruneLocked(now)
 		if len(s.waiting) >= handshakeMaxWaiting {
-			// Still full of peers whose window has closed and that are
-			// only kept for their rearm time. Let those go rather than
-			// refuse a new request its fast polling.
+			// Still full, partly of peers whose window has closed and that
+			// are only kept for their rearm time. Let the one whose rearm
+			// ends soonest go rather than refuse a new request its fast
+			// polling; the others keep their hold-off.
+			var oldest uint32
+			var oldestRearm time.Time
+			found := false
 			for p, w := range s.waiting {
-				if !now.Before(w.until) {
-					delete(s.waiting, p)
+				if !now.Before(w.until) && (!found || w.rearm.Before(oldestRearm)) {
+					oldest, oldestRearm, found = p, w.rearm, true
 				}
 			}
+			if found {
+				delete(s.waiting, oldest)
+			}
 		}
 		if len(s.waiting) >= handshakeMaxWaiting {
 			s.mu.Unlock()
@@ -266,6 +277,9 @@ func (s *handshakePollSched) begin(minGap time.Duration) (done chan struct{}, st
 	if s.running != nil {
 		return s.running, false
 	}
+	if s.closed {
+		return nil, false
+	}
 	if !s.claimLocked(minGap) {
 		return nil, false
 	}
@@ -285,6 +299,31 @@ func (s *handshakePollSched) end() {
 	s.nudge()
 }
 
+// shutdown stops any further poll from starting and waits until deadline
+// for the one in flight to finish, reporting whether none is running when it
+// returns. The daemon calls it before stopping the handshake manager and
+// closing the registry client: the registry empties the node's handshake
+// inbox as it answers, so a poll whose reply is cut off, or processed by a
+// stopped manager, loses what it carried. Setting closed under the same lock
+// that begin takes means no poll can slip in after the check.
+func (s *handshakePollSched) shutdown(deadline time.Time) bool {
+	s.mu.Lock()
+	s.closed = true
+	running := s.running
+	s.mu.Unlock()
+	if running == nil {
+		return true
+	}
+	t := time.NewTimer(time.Until(deadline))
+	defer t.Stop()
+	select {
+	case <-running:
+		return true
+	case <-t.C:
+		return false
+	}
+}
+
 // sinceLastPoll is how long ago the most recent poll started.
 func (s *handshakePollSched) sinceLastPoll() time.Duration {
 	s.mu.Lock()
@@ -357,6 +396,11 @@ func (s *handshakePollSched) pokeWait() (due bool, wait time.Duration) {
 // would be handshakes lost.
 func (d *Daemon) pollHandshakes(minGap, wait time.Duration) {
 	s := d.hsPoll
+	if d.stopping() {
+		// No new poll once shutdown has begun: the registry client is about
+		// to be closed, and a poll cut off by that loses what it fetched.
+		return
+	}
 	if d.reg() == nil {
 		// Nothing to poll; do not leave a poke owed (the loop would keep
 		// re-arming its timer for a poll that cannot run).

pkg/daemon/tunnel.go

@@ -2332,8 +2332,12 @@ func (tm *TunnelManager) handleBeaconMessage(data []byte, from *net.UDPAddr) {
 			slog.Debug("dropping notify from non-beacon source", "from", from)
 			return
 		}
-		if len(data) != 2 || data[1] != beaconNotifyHandshake {
+		if len(data) < 2 || data[1] != beaconNotifyHandshake {
 			// A kind this daemon does not know is not a reason to poll.
+			// Bytes after a known kind are ignored, so a later beacon can
+			// add to the message without daemons from this one on dropping
+			// it. (v1.16.0 requires exactly two bytes: a beacon cannot
+			// extend the message while v1.16.0 nodes remain.)
 			slog.Debug("dropping beacon notify of unknown kind", "len", len(data))
 			return
 		}
@@ -2345,7 +2349,7 @@ func (tm *TunnelManager) handleBeaconMessage(data []byte, from *net.UDPAddr) {
 	}
 }
 
-// beaconMsgNotify is a beacon → node message, [0x0A][kind(1)], telling the
+// beaconMsgNotify is a beacon → node message, [0x0A][kind(1)][...], telling the
 // node that the registry is holding something for it. kind 0x01 is a relayed
 // trust-handshake request or answer; the node polls for it. It carries no
 // node ID, address or payload. Daemons that predate it log it as an unknown

pkg/daemon/zz_handshake_poll_sched_test.go

@@ -381,6 +381,32 @@ func TestSettledPeersDoNotCrowdOutANewRequest(t *testing.T) {
 	if !s.waitingOn(1000) {
 		t.Fatalf("with %d answered peers tracked, a new request got no fast polling", handshakeMaxWaiting)
 	}
+	// Only one settled peer made room; the rest keep their hold-off, so an
+	// automatic handshake to them still cannot reopen a window early.
+	s.mu.Lock()
+	kept := len(s.waiting)
+	s.mu.Unlock()
+	if kept != handshakeMaxWaiting {
+		t.Fatalf("%d peers tracked after making room for one, want %d", kept, handshakeMaxWaiting)
+	}
+
+	// The one let go is the peer whose hold-off ends soonest.
+	clock := time.Now()
+	s = newHandshakePollSched()
+	s.now = func() time.Time { return clock }
+	for peer := uint32(1); peer <= handshakeMaxWaiting; peer++ {
+		s.requestSent(peer, true)
+		s.answered(peer)
+		clock = clock.Add(time.Second)
+	}
+	s.requestSent(1000, true)
+	s.mu.Lock()
+	_, firstKept := s.waiting[1]
+	_, lastKept := s.waiting[handshakeMaxWaiting]
+	s.mu.Unlock()
+	if firstKept || !lastKept {
+		t.Fatalf("made room by letting go of the wrong peer (earliest kept: %v, latest kept: %v)", firstKept, lastKept)
+	}
 
 	// Still a cap on requests that are genuinely outstanding.
 	s = newHandshakePollSched()
@@ -425,12 +451,93 @@ func TestBeaconNotifyOfUnknownKindIsDropped(t *testing.T) {
 
 	tm.handleBeaconMessage([]byte{beaconMsgNotify}, beacon)
 	tm.handleBeaconMessage([]byte{beaconMsgNotify, 0x02}, beacon)
-	tm.handleBeaconMessage([]byte{beaconMsgNotify, beaconNotifyHandshake, 0x00}, beacon)
 	if calls != 0 {
 		t.Fatalf("a malformed or unknown notify ran the handler %d times", calls)
 	}
 	tm.handleBeaconMessage([]byte{beaconMsgNotify, beaconNotifyHandshake}, beacon)
 	if calls != 1 {
 		t.Fatalf("a handshake notify ran the handler %d times, want 1", calls)
 	}
+	// A later beacon may add bytes after the kind; this daemon still acts on
+	// the kind it knows.
+	tm.handleBeaconMessage([]byte{beaconMsgNotify, beaconNotifyHandshake, 0x00, 0x07}, beacon)
+	if calls != 2 {
+		t.Fatalf("a handshake notify with trailing bytes ran the handler %d times in total, want 2", calls)
+	}
+}
+
+// Shutdown waits for a poll in flight before stopping the handshake manager
+// and closing the registry client — the registry has already taken what the
+// poll carries out of the inbox, so the manager must still be there to act on
+// it and save it — and starts no new poll once it has begun. In review, the
+// wait came after the manager had stopped, so the poll's results were lost
+// anyway.
+func TestShutdownWaitsForThePollInFlight(t *testing.T) {
+	t.Parallel()
+	reg, rc := startTestRegistry(t)
+	t.Cleanup(func() { reg.Close() })
+	t.Cleanup(func() { rc.Close() })
+	d := New(Config{KeepaliveInterval: time.Hour})
+	d.regConn.Store(rc)
+	hs := &stopRecordingHandshakes{fakeHandshakeService: &fakeHandshakeService{}}
+	d.RegisterHandshakeService(hs)
+
+	release := make(chan struct{})
+	var runs atomic.Int32
+	var managerStoppedFirst atomic.Bool
+	d.hsPoll.run = func() {
+		runs.Add(1)
+		<-release
+		managerStoppedFirst.Store(hs.stopped.Load())
+	}
+	d.pollHandshakes(0, 0)
+
+	stopped := make(chan struct{})
+	go func() {
+		_ = d.Stop()
+		close(stopped)
+	}()
+	time.Sleep(100 * time.Millisecond)
+	select {
+	case <-stopped:
+		t.Fatal("Stop returned with a poll in flight")
+	default:
+	}
+	d.pollHandshakes(0, 0) // asked for after shutdown began
+	close(release)
+	select {
+	case <-stopped:
+	case <-time.After(10 * time.Second):
+		t.Fatal("Stop did not return after the poll finished")
+	}
+	if managerStoppedFirst.Load() {
+		t.Fatal("the handshake manager was stopped before the poll in flight finished")
+	}
+	if !hs.stopped.Load() {
+		t.Fatal("Stop did not stop the handshake manager")
+	}
+	if n := runs.Load(); n != 1 {
+		t.Fatalf("%d polls ran, want 1: none may start once shutdown has begun", n)
+	}
+}
+
+// stopRecordingHandshakes records, race-free, when the daemon stops it.
+type stopRecordingHandshakes struct {
+	*fakeHandshakeService
+	stopped atomic.Bool
+}
+
+func (s *stopRecordingHandshakes) Stop() { s.stopped.Store(true) }
+
+// Once shutdown has run, no poll starts, even for a trigger that got past the
+// daemon's stopping check before shutdown began.
+func TestNoPollStartsAfterShutdown(t *testing.T) {
+	t.Parallel()
+	s := newHandshakePollSched()
+	if !s.shutdown(time.Now()) {
+		t.Fatal("shutdown with no poll in flight reported one running")
+	}
+	if done, start := s.begin(0); start || done != nil {
+		t.Fatalf("begin after shutdown = (%v, %v), want no poll", done, start)
+	}
 }