From a01274685360cf533d144896ca3991e5d586edf1 Mon Sep 17 00:00:00 2001 From: Marco Munizaga Date: Thu, 23 Jul 2026 21:10:13 -0700 Subject: [PATCH] partialmessages: reconcile enable and disable lifecycle --- extensions.go | 4 +-- extensions_test.go | 41 ++++++++++++++++------------- partialmessages/partialmsgs.go | 6 ++--- partialmessages/partialmsgs_test.go | 14 +++++----- 4 files changed, 34 insertions(+), 31 deletions(-) diff --git a/extensions.go b/extensions.go index f1834cd8..47242edc 100644 --- a/extensions.go +++ b/extensions.go @@ -90,7 +90,7 @@ func (pe *PeerExtensions) ExtendRPC(rpc *RPC) *RPC { // Purposely not trying to make a generic extension interface as there is only // one real consumer (partial messages). This may change in the future. type partialMessageInterface interface { - OnClosedOutboundStream(peer.ID) + OnPeerUnavailable(peer.ID) HandleRPC(from peer.ID, rpc *pubsub_pb.PartialMessagesExtension) error Heartbeat() EmitGossip(topic string, peers []peer.ID) @@ -224,7 +224,7 @@ func (es *extensionsState) reconcilePeerExtensions(id peer.ID, state *peerExtens es.disableTopicStreams(id) } if active.PartialMessages && es.partialMessagesExtension != nil { - es.partialMessagesExtension.OnClosedOutboundStream(id) + es.partialMessagesExtension.OnPeerUnavailable(id) } } if !state.active && shouldActivate { diff --git a/extensions_test.go b/extensions_test.go index 6e14ecbe..be499249 100644 --- a/extensions_test.go +++ b/extensions_test.go @@ -11,11 +11,11 @@ import ( ) type lifecyclePartialMessages struct { - closed []peer.ID + unavailable []peer.ID } -func (m *lifecyclePartialMessages) OnClosedOutboundStream(id peer.ID) { - m.closed = append(m.closed, id) +func (m *lifecyclePartialMessages) OnPeerUnavailable(id peer.ID) { + m.unavailable = append(m.unavailable, id) } func (*lifecyclePartialMessages) HandleRPC(peer.ID, *pubsub_pb.PartialMessagesExtension) error { @@ -66,18 +66,21 @@ func TestExtensionsDeactivateOnEitherHalfClosing(t *testing.T) { es := newPartialLifecycleState(cleanup) id := peer.ID("peer") activatePartialExtensions(t, es, id) + if len(cleanup.unavailable) != 0 { + t.Fatalf("activation unexpectedly performed partial-message lifecycle work: %v", cleanup.unavailable) + } test.close(es, id) if es.activePeerExtensions(id).PartialMessages { t.Fatal("extension remained active after stream closure") } - if len(cleanup.closed) != 1 || cleanup.closed[0] != id { - t.Fatalf("expected one cleanup for %q, got %v", id, cleanup.closed) + if len(cleanup.unavailable) != 1 || cleanup.unavailable[0] != id { + t.Fatalf("expected one cleanup for %q, got %v", id, cleanup.unavailable) } test.close(es, id) - if len(cleanup.closed) != 1 { - t.Fatalf("duplicate closure triggered %d cleanups", len(cleanup.closed)) + if len(cleanup.unavailable) != 1 { + t.Fatalf("duplicate closure triggered %d cleanups", len(cleanup.unavailable)) } }) } @@ -119,8 +122,8 @@ func TestExtensionsReplacementHalfReactivates(t *testing.T) { if !es.activePeerExtensions(id).PartialMessages { t.Fatal("replacement half did not reactivate extension") } - if len(cleanup.closed) != 1 { - t.Fatalf("expected one cleanup before reactivation, got %d", len(cleanup.closed)) + if len(cleanup.unavailable) != 1 || cleanup.unavailable[0] != id { + t.Fatalf("replacement lifecycle cleanup = %v; want one cleanup for %q", cleanup.unavailable, id) } }) } @@ -136,8 +139,8 @@ func TestExtensionsPartialCleanupUsesActiveSnapshot(t *testing.T) { es.myExtensions.PartialMessages = false es.OnClosedOutboundStream(id) - if len(cleanup.closed) != 1 || cleanup.closed[0] != id { - t.Fatalf("expected cleanup from negotiated snapshot, got %v", cleanup.closed) + if len(cleanup.unavailable) != 1 || cleanup.unavailable[0] != id { + t.Fatalf("expected cleanup from negotiated snapshot, got %v", cleanup.unavailable) } } @@ -325,8 +328,8 @@ func TestPeerExtensionsIncomingLifecycle(t *testing.T) { if es.peers[peerID] != nil && es.peers[peerID].active { t.Fatal("incoming close left extensions active") } - if disabled != 1 || partial.closed != 1 { - t.Fatalf("deactivation counts = topic %d, partial %d; want 1, 1", disabled, partial.closed) + if disabled != 1 || partial.unavailable != 1 { + t.Fatalf("deactivation counts = topic %d, partial %d; want 1, 1", disabled, partial.unavailable) } if es.Preprocess(&RPC{from: peerID, transport: peercomm.TransportTopic}) { t.Fatal("accepted topic RPC while extensions were inactive") @@ -343,11 +346,11 @@ func TestPeerExtensionsIncomingLifecycle(t *testing.T) { } es.OnClosedOutboundStream(peerID) - if disabled != 2 || partial.closed != 1 { - t.Fatalf("final deactivation counts = topic %d, partial %d; want 2, 1", disabled, partial.closed) + if disabled != 2 || partial.unavailable != 1 { + t.Fatalf("final deactivation counts = topic %d, partial %d; want 2, 1", disabled, partial.unavailable) } es.OnClosedOutboundStream(peerID) - if disabled != 2 || partial.closed != 1 { + if disabled != 2 || partial.unavailable != 1 { t.Fatal("repeated outbound close deactivated extensions twice") } } @@ -408,11 +411,11 @@ func extensionHello(from peer.ID, topicStreams, partialMessages bool) *RPC { } type recordingPartialMessageExtension struct { - closed int + unavailable int } -func (m *recordingPartialMessageExtension) OnClosedOutboundStream(peer.ID) { - m.closed++ +func (m *recordingPartialMessageExtension) OnPeerUnavailable(peer.ID) { + m.unavailable++ } func (*recordingPartialMessageExtension) HandleRPC(peer.ID, *pubsub_pb.PartialMessagesExtension) error { diff --git a/partialmessages/partialmsgs.go b/partialmessages/partialmsgs.go index 12584b74..566d269b 100644 --- a/partialmessages/partialmsgs.go +++ b/partialmessages/partialmsgs.go @@ -241,13 +241,13 @@ func (e *PartialMessagesExtension[PeerState]) initPeerState(topic string, gState } } -func (e *PartialMessagesExtension[PeerState]) OnClosedOutboundStream(id peer.ID) { +func (e *PartialMessagesExtension[PeerState]) OnPeerUnavailable(id peer.ID) { for topic, tState := range e.statePerTopicPerGroup { for _, gState := range tState { delete(gState.peerState, id) } if ctr, ok := e.peerInitiatedGroupCounter[topic]; ok { - ctr.OnClosedOutboundStream(id) + ctr.OnPeerUnavailable(id) } } } @@ -355,7 +355,7 @@ func (ctr *peerInitiatedGroupCounterState) Dec(id peer.ID) { } } -func (ctr *peerInitiatedGroupCounterState) OnClosedOutboundStream(id peer.ID) { +func (ctr *peerInitiatedGroupCounterState) OnPeerUnavailable(id peer.ID) { if n, ok := ctr.perPeer[id]; ok { ctr.total -= n delete(ctr.perPeer, id) diff --git a/partialmessages/partialmsgs_test.go b/partialmessages/partialmsgs_test.go index 080d15f5..93f0a997 100644 --- a/partialmessages/partialmsgs_test.go +++ b/partialmessages/partialmsgs_test.go @@ -71,8 +71,8 @@ func (m *mockNetworkPartialMessages) removePeers() { if a == b { continue } - m.handlers[a].OnClosedOutboundStream(b) - m.handlers[b].OnClosedOutboundStream(a) + m.handlers[a].OnPeerUnavailable(b) + m.handlers[b].OnPeerUnavailable(a) } } @@ -1169,8 +1169,8 @@ func TestGossipDelivery(t *testing.T) { } // Cleanup - h1Handler.OnClosedOutboundStream(h2ID) - h2Handler.OnClosedOutboundStream(h1ID) + h1Handler.OnPeerUnavailable(h2ID) + h2Handler.OnPeerUnavailable(h1ID) for range 10 { h1Handler.Heartbeat() h2Handler.Heartbeat() @@ -1274,10 +1274,10 @@ func TestPeerInitiatedCounter(t *testing.T) { } // All peers go away, and the counts should be back to 0 - handler.OnClosedOutboundStream("1") + handler.OnPeerUnavailable("1") for id := range 2 { otherPeer := fmt.Sprintf("peer%d", id) - handler.OnClosedOutboundStream(peer.ID(otherPeer)) + handler.OnPeerUnavailable(peer.ID(otherPeer)) } assertCounts(0, map[peer.ID]int{}) @@ -1377,7 +1377,7 @@ func FuzzPeerInitiatedCounter(f *testing.F) { delete(expectedPeercounts, peer.ID(otherPeer)) } - handler.OnClosedOutboundStream(peer.ID(otherPeer)) + handler.OnPeerUnavailable(peer.ID(otherPeer)) case 2: // heartbeat until everything is cleared script = script[1:] for range handler.GroupTTLByHeatbeat + 1 {