From 02182f16848346445a40654548329a479c6848c2 Mon Sep 17 00:00:00 2001 From: Teo Calin Date: Wed, 7 Oct 2026 16:12:58 +0300 Subject: [PATCH 1/2] daemon: shutdown waits for a handshake poll in flight; smaller follow-ups to #500 From the second review of #500. - Shutdown closed the registry client without waiting for a poll in flight. The registry empties the node's handshake inbox as it answers, so a restart during a poll lost what it carried. doStop now waits up to 5s for the poll (waitIdle), and pollHandshakes starts none once the daemon is stopping. - At the 64-peer cap, only the closed-window peer whose rearm ends soonest is evicted; the others keep their hold-off against automatic handshakes reopening a window. - A beacon notify with bytes after the handshake kind is acted on, so a later beacon can extend the message without older daemons ignoring it. - Comments: the package doc describes the wait-for-trust condition, and pollRelayedHandshakes no longer documents the removed timeout. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 6 +++ pkg/daemon/daemon.go | 11 +++-- pkg/daemon/handshakepoll.go | 48 ++++++++++++++++--- pkg/daemon/tunnel.go | 4 +- pkg/daemon/zz_handshake_poll_sched_test.go | 54 +++++++++++++++++++++- 5 files changed, 110 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 88f8f37c..0cc042e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -285,6 +285,12 @@ 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 up to 5s for it, and starts no new poll once it has + begun. - **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 diff --git a/pkg/daemon/daemon.go b/pkg/daemon/daemon.go index 0128113e..90357fe1 100644 --- a/pkg/daemon/daemon.go +++ b/pkg/daemon/daemon.go @@ -1713,6 +1713,11 @@ func (d *Daemon) doStop() { // (including persisted deny/grudge lists). The 5-minute heartbeat TTL // reaps truly-dead nodes; users who explicitly want to leave call // `pilotctl deregister` via IPC (CmdDeregister) which is unaffected. + // A relayed-handshake poll in flight finishes first: closing the client + // would cut it off after the registry had already emptied the inbox. + if !d.hsPoll.waitIdle(5 * time.Second) { + slog.Warn("relayed-handshake poll still running at shutdown; closing the registry client anyway") + } if d.reg() != nil { d.reg().Close() } @@ -6565,10 +6570,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 { diff --git a/pkg/daemon/handshakepoll.go b/pkg/daemon/handshakepoll.go index ba1b02dd..895ec0cb 100644 --- a/pkg/daemon/handshakepoll.go +++ b/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 @@ -150,14 +151,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() @@ -285,6 +293,27 @@ func (s *handshakePollSched) end() { s.nudge() } +// waitIdle waits up to timeout for a poll in flight to finish, and reports +// whether none is running when it returns. Shutdown calls it before closing +// the registry client: the registry empties the node's handshake inbox as it +// answers, so a poll cut off mid-reply would lose what it carried. +func (s *handshakePollSched) waitIdle(timeout time.Duration) bool { + s.mu.Lock() + running := s.running + s.mu.Unlock() + if running == nil { + return true + } + t := time.NewTimer(timeout) + 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 +386,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). diff --git a/pkg/daemon/tunnel.go b/pkg/daemon/tunnel.go index 34aeb1e8..33c8c295 100644 --- a/pkg/daemon/tunnel.go +++ b/pkg/daemon/tunnel.go @@ -2332,8 +2332,10 @@ 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 older daemons dropping it. slog.Debug("dropping beacon notify of unknown kind", "len", len(data)) return } diff --git a/pkg/daemon/zz_handshake_poll_sched_test.go b/pkg/daemon/zz_handshake_poll_sched_test.go index 6d0b3494..ac98b807 100644 --- a/pkg/daemon/zz_handshake_poll_sched_test.go +++ b/pkg/daemon/zz_handshake_poll_sched_test.go @@ -381,6 +381,14 @@ 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) + } // Still a cap on requests that are genuinely outstanding. s = newHandshakePollSched() @@ -425,7 +433,6 @@ 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) } @@ -433,4 +440,49 @@ func TestBeaconNotifyOfUnknownKindIsDropped(t *testing.T) { 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 closing the registry client +// (a poll cut off mid-reply loses what the registry already took out of the +// inbox), and starts no new poll once it has begun. +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) + + release := make(chan struct{}) + var finished atomic.Int32 + d.hsPoll.run = func() { + <-release + finished.Add(1) + } + d.pollHandshakes(0, 0) + if d.hsPoll.waitIdle(50 * time.Millisecond) { + t.Fatal("waitIdle reported idle with a poll in flight") + } + go func() { + time.Sleep(100 * time.Millisecond) + close(release) + }() + if !d.hsPoll.waitIdle(2 * time.Second) { + t.Fatal("waitIdle gave up although the poll finished") + } + if finished.Load() != 1 { + t.Fatal("the poll did not run to completion") + } + + close(d.stopCh) + d.pollHandshakes(0, 0) + if got := d.RelayedHandshakePolls(); got != 1 { + t.Fatalf("registry polls = %d after shutdown began, want 1 (none new)", got) + } } From eb67eba52686cab1c03ae94db818100742ad34c0 Mon Sep 17 00:00:00 2001 From: Teo Calin Date: Wed, 7 Oct 2026 17:08:49 +0300 Subject: [PATCH 2/2] daemon: shutdown waits for the handshake poll before the manager stops Review fixes. - Shutdown waited for a poll in flight only after the handshake manager had stopped, so the poll's reply was processed by a stopped manager: requests and approvals were neither saved nor answered, and the registry had already emptied the inbox. The wait now comes before the manager stops. - No poll can start once shutdown has begun: the scheduler is closed under the same lock that starting a poll takes. Before, a caller past the stopping check could start one after the wait had returned. - The wait shares the 5s deadline shutdown already gives its background goroutines instead of adding 5s of its own. - The shutdown test now calls Stop and checks the order; tests cover the closed scheduler and which settled peer is let go. Each fails against the matching revert. - Changelog entries for the eviction and notify changes; the notify comments say v1.16.0 still requires exactly two bytes. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 13 +++- pkg/daemon/daemon.go | 18 +++-- pkg/daemon/handshakepoll.go | 22 ++++-- pkg/daemon/tunnel.go | 6 +- pkg/daemon/zz_handshake_poll_sched_test.go | 91 +++++++++++++++++----- 5 files changed, 116 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0cc042e1..aa17ec3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -289,8 +289,17 @@ Detailed per-release notes are on 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 up to 5s for it, and starts no new poll once it has - begun. + 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 diff --git a/pkg/daemon/daemon.go b/pkg/daemon/daemon.go index 90357fe1..80ea78f9 100644 --- a/pkg/daemon/daemon.go +++ b/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() @@ -1713,11 +1724,6 @@ func (d *Daemon) doStop() { // (including persisted deny/grudge lists). The 5-minute heartbeat TTL // reaps truly-dead nodes; users who explicitly want to leave call // `pilotctl deregister` via IPC (CmdDeregister) which is unaffected. - // A relayed-handshake poll in flight finishes first: closing the client - // would cut it off after the registry had already emptied the inbox. - if !d.hsPoll.waitIdle(5 * time.Second) { - slog.Warn("relayed-handshake poll still running at shutdown; closing the registry client anyway") - } if d.reg() != nil { d.reg().Close() } diff --git a/pkg/daemon/handshakepoll.go b/pkg/daemon/handshakepoll.go index 895ec0cb..e0264107 100644 --- a/pkg/daemon/handshakepoll.go +++ b/pkg/daemon/handshakepoll.go @@ -104,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() @@ -274,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 } @@ -293,18 +299,22 @@ func (s *handshakePollSched) end() { s.nudge() } -// waitIdle waits up to timeout for a poll in flight to finish, and reports -// whether none is running when it returns. Shutdown calls it before closing -// the registry client: the registry empties the node's handshake inbox as it -// answers, so a poll cut off mid-reply would lose what it carried. -func (s *handshakePollSched) waitIdle(timeout time.Duration) bool { +// 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(timeout) + t := time.NewTimer(time.Until(deadline)) defer t.Stop() select { case <-running: diff --git a/pkg/daemon/tunnel.go b/pkg/daemon/tunnel.go index 33c8c295..02926431 100644 --- a/pkg/daemon/tunnel.go +++ b/pkg/daemon/tunnel.go @@ -2335,7 +2335,9 @@ func (tm *TunnelManager) handleBeaconMessage(data []byte, from *net.UDPAddr) { 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 older daemons dropping it. + // 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 } @@ -2347,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 diff --git a/pkg/daemon/zz_handshake_poll_sched_test.go b/pkg/daemon/zz_handshake_poll_sched_test.go index ac98b807..d796acea 100644 --- a/pkg/daemon/zz_handshake_poll_sched_test.go +++ b/pkg/daemon/zz_handshake_poll_sched_test.go @@ -390,6 +390,24 @@ func TestSettledPeersDoNotCrowdOutANewRequest(t *testing.T) { 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() for peer := uint32(1); peer <= handshakeMaxWaiting; peer++ { @@ -448,9 +466,12 @@ func TestBeaconNotifyOfUnknownKindIsDropped(t *testing.T) { } } -// Shutdown waits for a poll in flight before closing the registry client -// (a poll cut off mid-reply loses what the registry already took out of the -// inbox), and starts no new poll once it has begun. +// 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) @@ -458,31 +479,65 @@ func TestShutdownWaitsForThePollInFlight(t *testing.T) { 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 finished atomic.Int32 + var runs atomic.Int32 + var managerStoppedFirst atomic.Bool d.hsPoll.run = func() { + runs.Add(1) <-release - finished.Add(1) + managerStoppedFirst.Store(hs.stopped.Load()) } d.pollHandshakes(0, 0) - if d.hsPoll.waitIdle(50 * time.Millisecond) { - t.Fatal("waitIdle reported idle with a poll in flight") - } + + stopped := make(chan struct{}) go func() { - time.Sleep(100 * time.Millisecond) - close(release) + _ = d.Stop() + close(stopped) }() - if !d.hsPoll.waitIdle(2 * time.Second) { - t.Fatal("waitIdle gave up although the poll finished") + time.Sleep(100 * time.Millisecond) + select { + case <-stopped: + t.Fatal("Stop returned with a poll in flight") + default: } - if finished.Load() != 1 { - t.Fatal("the poll did not run to completion") + 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) + } +} - close(d.stopCh) - d.pollHandshakes(0, 0) - if got := d.RelayedHandshakePolls(); got != 1 { - t.Fatalf("registry polls = %d after shutdown began, want 1 (none new)", got) +// 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) } }