diff --git a/extensions.go b/extensions.go index 6b921391..eae1adbd 100644 --- a/extensions.go +++ b/extensions.go @@ -6,6 +6,7 @@ import ( "github.com/libp2p/go-libp2p-pubsub/partialmessages" pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb" + "github.com/libp2p/go-libp2p-pubsub/vivaldi" "github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/core/protocol" ) @@ -13,6 +14,7 @@ import ( type PeerExtensions struct { TestExtension bool PartialMessages bool + Spread bool } type TestExtensionConfig struct { @@ -44,6 +46,7 @@ func peerExtensionsFromRPC(rpc *RPC) PeerExtensions { if hasPeerExtensions(rpc) { out.TestExtension = rpc.Control.Extensions.GetTestExtension() out.PartialMessages = rpc.Control.Extensions.GetPartialMessages() + out.Spread = rpc.Control.Extensions.GetSpread() } return out } @@ -67,6 +70,15 @@ func (pe *PeerExtensions) ExtendRPC(rpc *RPC) *RPC { } rpc.Control.Extensions.PartialMessages = &pe.PartialMessages } + if pe.Spread { + if rpc.Control == nil { + rpc.Control = &pubsub_pb.ControlMessage{} + } + if rpc.Control.Extensions == nil { + rpc.Control.Extensions = &pubsub_pb.ControlExtensions{} + } + rpc.Control.Extensions.Spread = &pe.Spread + } return rpc } @@ -91,6 +103,7 @@ type extensionsState struct { testExtension *testExtension partialMessagesExtension partialMessageInterface + spreadState *SpreadState } func newExtensionsState(myExtensions PeerExtensions, reportMisbehavior func(peer.ID), sendRPC func(peer.ID, *RPC, bool)) *extensionsState { @@ -101,6 +114,7 @@ func newExtensionsState(myExtensions PeerExtensions, reportMisbehavior func(peer reportMisbehavior: reportMisbehavior, sendRPC: sendRPC, testExtension: nil, + spreadState: NewSpreadState(), } } @@ -168,6 +182,10 @@ func (es *extensionsState) extensionsOnNewOutboundStream(id peer.ID) { if es.myExtensions.TestExtension && es.peerExtensions[id].TestExtension { es.testExtension.OnNewOutboundStream(id) } + + if es.myExtensions.Spread && es.peerExtensions[id].Spread { + es.spreadState.AddPeer(id) + } } // extensionsOnClosedOutboundStream is always called after extensionsOnNewOutboundStream. @@ -175,6 +193,9 @@ func (es *extensionsState) extensionsOnClosedOutboundStream(id peer.ID) { if es.myExtensions.PartialMessages && es.peerExtensions[id].PartialMessages { es.partialMessagesExtension.OnClosedOutboundStream(id) } + if es.myExtensions.Spread && es.peerExtensions[id].Spread { + es.spreadState.RemovePeer(id) + } } func (es *extensionsState) extensionsHandleRPC(rpc *RPC) error { @@ -189,6 +210,18 @@ func (es *extensionsState) extensionsHandleRPC(rpc *RPC) error { } } + // Update SPREAD topic membership when peers announce subscriptions + if es.myExtensions.Spread && es.peerExtensions[rpc.from].Spread && rpc.Subscriptions != nil { + for _, sub := range rpc.GetSubscriptions() { + topic := sub.GetTopicid() + if sub.GetSubscribe() { + es.spreadState.AddPeerTopic(topic, rpc.from) + } else { + es.spreadState.RemovePeerTopic(topic, rpc.from) + } + } + } + return nil } @@ -254,6 +287,41 @@ func PublishPartial[PeerState any](ps *PubSub, topic string, groupID []byte, pub } } +// WithVivaldi wires a Vivaldi service into the spread extension and starts +// the runner. Passing a nil service disables Vivaldi. +func WithVivaldi(vsvc *vivaldi.Service, cfg *VivaldiConfig) Option { + return func(ps *PubSub) error { + gs, ok := ps.rt.(*GossipSubRouter) + if !ok { + return errors.New("pubsub router is not gossipsub") + } + if gs.extensions == nil || gs.extensions.spreadState == nil { + return errors.New("spread extension state not initialized") + } + gs.extensions.spreadState.ConfigureVivaldi(vsvc, cfg) + // Start runner if service is provided + if vsvc != nil { + gs.extensions.spreadState.StartVivaldiRunner() + } + return nil + } +} + +// WithSpreadClusteringConfig configures spread clustering behaviour. +func WithSpreadClusteringConfig(cfg *SpreadClusteringConfig) Option { + return func(ps *PubSub) error { + gs, ok := ps.rt.(*GossipSubRouter) + if !ok { + return errors.New("pubsub router is not gossipsub") + } + if gs.extensions == nil || gs.extensions.spreadState == nil { + return errors.New("spread extension state not initialized") + } + gs.extensions.spreadState.ConfigureClustering(cfg) + return nil + } +} + type partialMessageRouter struct { gs *GossipSubRouter } diff --git a/gossip_protocols.go b/gossip_protocols.go new file mode 100644 index 00000000..d21cdbeb --- /dev/null +++ b/gossip_protocols.go @@ -0,0 +1,13 @@ +package pubsub + +// Gossip protocol choices selectable via WithProtocolChoice. GOSSIPSUB is the +// default behaviour; SPREAD enables the SPREAD anonymous-gossip selection for +// messages marked with msg.Spread. +const ( + GOSSIPSUB = 1 + SPREAD = 2 +) + +// GossipProtocolChoice selects which forwarding strategy the GossipSub router +// uses. See the GOSSIPSUB and SPREAD constants. +type GossipProtocolChoice int diff --git a/gossipsub.go b/gossipsub.go index 37bbdb14..2e3259f2 100644 --- a/gossipsub.go +++ b/gossipsub.go @@ -294,26 +294,28 @@ func NewGossipSubWithRouter(ctx context.Context, h host.Host, rt PubSubRouter, o func DefaultGossipSubRouter(h host.Host) *GossipSubRouter { params := DefaultGossipSubParams() rt := &GossipSubRouter{ - peers: make(map[peer.ID]protocol.ID), - mesh: make(map[string]map[peer.ID]struct{}), - fanout: make(map[string]map[peer.ID]struct{}), - lastpub: make(map[string]int64), - gossip: make(map[peer.ID][]*pb.ControlIHave), - control: make(map[peer.ID]*pb.ControlMessage), - backoff: make(map[string]map[peer.ID]time.Time), - peerhave: make(map[peer.ID]int), - peerdontwant: make(map[peer.ID]int), - unwanted: make(map[peer.ID]map[checksum]int), - iasked: make(map[peer.ID]int), - outbound: make(map[peer.ID]bool), - connect: make(chan connectInfo, params.MaxPendingConnections), - cab: pstoremem.NewAddrBook(), - mcache: NewMessageCache(params.HistoryGossip, params.HistoryLength), - protos: GossipSubDefaultProtocols, - feature: GossipSubDefaultFeatures, - tagTracer: newTagTracer(h.ConnManager()), - params: params, - reducePXRecords: defaultPXRecordReducer, + peers: make(map[peer.ID]protocol.ID), + mesh: make(map[string]map[peer.ID]struct{}), + fanout: make(map[string]map[peer.ID]struct{}), + lastpub: make(map[string]int64), + gossip: make(map[peer.ID][]*pb.ControlIHave), + control: make(map[peer.ID]*pb.ControlMessage), + backoff: make(map[string]map[peer.ID]time.Time), + peerhave: make(map[peer.ID]int), + peerdontwant: make(map[peer.ID]int), + unwanted: make(map[peer.ID]map[checksum]int), + iasked: make(map[peer.ID]int), + outbound: make(map[peer.ID]bool), + connect: make(chan connectInfo, params.MaxPendingConnections), + cab: pstoremem.NewAddrBook(), + mcache: NewMessageCache(params.HistoryGossip, params.HistoryLength), + protos: GossipSubDefaultProtocols, + feature: GossipSubDefaultFeatures, + tagTracer: newTagTracer(h.ConnManager()), + params: params, + reducePXRecords: defaultPXRecordReducer, + spreadPropagation: NewSpreadPropagation(nil), + spreadRelayDuplicateCounter: make(map[spreadDuplicateKey]int), } rt.extensions = newExtensionsState(PeerExtensions{}, func(p peer.ID) { @@ -364,6 +366,35 @@ func DefaultGossipSubParams() GossipSubParams { } } +// WithProtocolChoice is a gossipsub router option that selects the forwarding +// strategy. The default is GOSSIPSUB; pass SPREAD to enable SPREAD selection +// for messages marked with msg.Spread. +func WithProtocolChoice(choice GossipProtocolChoice) Option { + return func(p *PubSub) error { + // If p can be casted to GossipSub, set the protocol choice on the router. + gs, ok := p.rt.(*GossipSubRouter) + if !ok { + return fmt.Errorf("cannot set protocol choice since pubsub router is not gossipsub") + } + gs.gossipProtocolChoice = choice + return nil + } +} + +// WithSpreadPropagationConfig is a gossipsub router option that sets the SPREAD +// peer-selection parameters (intra/inter-cluster fanout and probabilities). +// Unset or invalid fields fall back to the package defaults. +func WithSpreadPropagationConfig(cfg *SpreadConfig) Option { + return func(ps *PubSub) error { + gs, ok := ps.rt.(*GossipSubRouter) + if !ok { + return fmt.Errorf("pubsub router is not gossipsub") + } + gs.spreadPropagation = NewSpreadPropagation(cfg) + return nil + } +} + // WithPeerScore is a gossipsub router option that enables peer scoring. func WithPeerScore(params *PeerScoreParams, thresholds *PeerScoreThresholds) Option { return func(ps *PubSub) error { @@ -650,6 +681,14 @@ type GossipSubRouter struct { // number of heartbeats since the beginning of time; this allows us to amortize some resource // clean up -- eg backoff clean up. heartbeatTicks uint64 + + // GossipProtocolChoice + gossipProtocolChoice GossipProtocolChoice + + // SPREAD propagation selector. + spreadPropagation *SpreadPropagation + // Per-(topic, message ID) duplicate SPREAD re-propagation counters. + spreadRelayDuplicateCounter map[spreadDuplicateKey]int } var _ BatchPublisher = &GossipSubRouter{} @@ -659,6 +698,11 @@ type connectInfo struct { spr *record.Envelope } +type spreadDuplicateKey struct { + topic string + msgID string +} + func (gs *GossipSubRouter) Protocols() []protocol.ID { return gs.protos } @@ -1349,6 +1393,54 @@ func (gs *GossipSubRouter) Publish(msg *Message) { } } +func (gs *GossipSubRouter) canRelaySeenSpreadDuplicate(msg *Message) bool { + if msg == nil || !msg.Spread { + return false + } + if gs.spreadPropagation == nil || gs.extensions == nil || !gs.extensions.myExtensions.Spread { + return false + } + return gs.spreadPropagation.config.DuplicateRepropagation > 0 +} + +func (gs *GossipSubRouter) hasSpreadDuplicateRelayBudget(topic, msgID string) bool { + maxRelays := gs.spreadPropagation.config.DuplicateRepropagation + if maxRelays <= 0 { + return false + } + + key := spreadDuplicateKey{topic: topic, msgID: msgID} + return gs.spreadRelayDuplicateCounter[key] < maxRelays +} + +func (gs *GossipSubRouter) consumeSpreadDuplicateRelayBudget(topic, msgID string) bool { + maxRelays := gs.spreadPropagation.config.DuplicateRepropagation + if maxRelays <= 0 { + return false + } + + key := spreadDuplicateKey{topic: topic, msgID: msgID} + used := gs.spreadRelayDuplicateCounter[key] + if used >= maxRelays { + return false + } + + gs.spreadRelayDuplicateCounter[key] = used + 1 + return true +} + +func (gs *GossipSubRouter) cleanupSpreadDuplicateRelayBudget() { + if len(gs.spreadRelayDuplicateCounter) == 0 { + return + } + for key := range gs.spreadRelayDuplicateCounter { + if gs.p.seenMessage(key.msgID) { + continue + } + delete(gs.spreadRelayDuplicateCounter, key) + } +} + func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] { return func(yield func(peer.ID, *RPC) bool) { gs.mcache.Put(msg) @@ -1405,7 +1497,39 @@ func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] { } } + // If this message was marked as SPREAD and SPREAD mode is enabled on this router, + // override the default selection with the SPREAD-selected peers. + if gs.gossipProtocolChoice == SPREAD && msg.Spread { + useAngular := false + if gs.spreadPropagation != nil && gs.spreadPropagation.config != nil { + useAngular = gs.spreadPropagation.config.UseAngularInterPeers + } + clusterPeers, interClusterPeers := gs.extensions.spreadState.GetPropagationPeers(topic, gs.p.host.ID(), useAngular) + if gs.spreadPropagation == nil { + gs.spreadPropagation = NewSpreadPropagation(nil) + } + selected := gs.spreadPropagation.GetPeersForPropagation(from, clusterPeers, interClusterPeers) + + spreadTosend := make(map[peer.ID]struct{}, len(selected)) + for _, p := range selected { + if p == from || p == peer.ID(msg.GetFrom()) { + continue + } + spreadTosend[p] = struct{}{} + } + + if len(spreadTosend) > 0 { + tosend = spreadTosend + } + } + out := rpcWithMessages(msg.Message) + // Set RPC-level spread extension if both the router is configured to use + // SPREAD by default and the message explicitly requests spread. + if gs.gossipProtocolChoice == SPREAD && msg.Spread { + v := true + out.Spread = &pb.SpreadExtension{SourceIsSpreadNode: &v} + } for pid := range tosend { if pid == from || pid == peer.ID(msg.GetFrom()) { continue @@ -1909,6 +2033,7 @@ func (gs *GossipSubRouter) heartbeat() { // advance the message history window gs.mcache.Shift() + gs.cleanupSpreadDuplicateRelayBudget() gs.extensions.Heartbeat() } diff --git a/gossipsub_test.go b/gossipsub_test.go index a10751fd..67a18bf8 100644 --- a/gossipsub_test.go +++ b/gossipsub_test.go @@ -3338,6 +3338,7 @@ func TestGossipsubIdontwantReceive(t *testing.T) { type mockRawTracer struct { onRecvRPC func(*RPC) + onSendRPC func(*RPC, peer.ID) } func (m *mockRawTracer) RecvRPC(rpc *RPC) { @@ -3356,13 +3357,207 @@ func (m *mockRawTracer) Leave(topic string) {} func (m *mockRawTracer) Prune(p peer.ID, topic string) {} func (m *mockRawTracer) RejectMessage(msg *Message, reason string) {} func (m *mockRawTracer) OnClosedOutboundStream(p peer.ID) {} -func (m *mockRawTracer) SendRPC(rpc *RPC, p peer.ID) {} -func (m *mockRawTracer) ThrottlePeer(p peer.ID) {} -func (m *mockRawTracer) UndeliverableMessage(msg *Message) {} -func (m *mockRawTracer) ValidateMessage(msg *Message) {} +func (m *mockRawTracer) SendRPC(rpc *RPC, p peer.ID) { + if m.onSendRPC != nil { + m.onSendRPC(rpc, p) + } +} +func (m *mockRawTracer) ThrottlePeer(p peer.ID) {} +func (m *mockRawTracer) UndeliverableMessage(msg *Message) {} +func (m *mockRawTracer) ValidateMessage(msg *Message) {} var _ RawTracer = &mockRawTracer{} +func withSpreadExtensionAdvertiseForTests() Option { + return func(ps *PubSub) error { + gs, ok := ps.rt.(*GossipSubRouter) + if !ok { + return fmt.Errorf("pubsub router is not gossipsub") + } + gs.extensions.myExtensions.Spread = true + return nil + } +} + +func publishWithSpreadForTests(ctx context.Context, topic *Topic, data []byte) error { + msg, err := topic.validate(ctx, data) + if err != nil { + return err + } + msg.Spread = true + return topic.p.val.sendMsgBlocking(msg) +} + +func TestGossipsubSpreadDuplicateRepropagationBudget(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + hosts := getDefaultHosts(t, 6) + sourceIdx := []int{0, 1, 2, 3} + middleIdx := 4 + sinkIdx := 5 + + topicName := "spread-duplicate-repropagation" + payload := []byte("same-payload") + hashMsgID := func(pmsg *pb.Message) string { + sum := sha256.Sum256(pmsg.GetData()) + return string(sum[:]) + } + + var sendsFromMiddleToLast atomic.Int32 + middleTracer := &mockRawTracer{ + onSendRPC: func(rpc *RPC, p peer.ID) { + if p != hosts[sinkIdx].ID() || len(rpc.GetPublish()) == 0 { + return + } + for _, msg := range rpc.GetPublish() { + if msg.GetTopic() == topicName && bytes.Equal(msg.GetData(), payload) { + sendsFromMiddleToLast.Add(1) + } + } + }, + } + + commonOpts := []Option{ + WithMessageIdFn(hashMsgID), + WithProtocolChoice(SPREAD), + WithSpreadPropagationConfig(&SpreadConfig{ + IntraFanout: 1, + InterFanout: 0, + IntraRho: 1, + InterProb: 0, + DuplicateRepropagation: 2, + }), + withSpreadExtensionAdvertiseForTests(), + } + + psubs := make([]*PubSub, len(hosts)) + for i := range hosts { + opts := append([]Option{}, commonOpts...) + if i == middleIdx { + opts = append(opts, WithRawTracer(middleTracer)) + } + psubs[i] = getGossipsub(ctx, hosts[i], opts...) + } + + for _, idx := range sourceIdx { + connect(t, hosts[idx], hosts[middleIdx]) + } + connect(t, hosts[middleIdx], hosts[sinkIdx]) + + topics := make([]*Topic, 0, len(psubs)) + for _, ps := range psubs { + topic, err := ps.Join(topicName) + if err != nil { + t.Fatal(err) + } + topics = append(topics, topic) + } + + for i, topic := range topics { + if i == 0 || i == 1 || i == 2 || i == 3 { + continue + } + sub, err := topic.Subscribe() + if err != nil { + t.Fatal(err) + } + go func(sub *Subscription) { + for { + _, err := sub.Next(ctx) + if err != nil { + return + } + } + }(sub) + } + + time.Sleep(2 * time.Second) + + for _, idx := range sourceIdx { + if err := publishWithSpreadForTests(ctx, topics[idx], payload); err != nil { + t.Fatal(err) + } + time.Sleep(200 * time.Millisecond) + } + + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if sendsFromMiddleToLast.Load() >= 3 { + break + } + time.Sleep(100 * time.Millisecond) + } + + if got := sendsFromMiddleToLast.Load(); got != 3 { + t.Fatalf("expected middle peer to relay 3 times (1 initial + 2 duplicate repropagations), got %d", got) + } +} + +// TestGossipsubSpreadSelectionDelivers checks that when SPREAD selection is +// active (InterProb=1, InterFanout=1) a published spread message is forwarded +// using the SPREAD-selected peer set and reaches the subscriber. With a single +// eligible inter peer the selection is deterministic, so this exercises the +// SPREAD override path end to end (publish -> SPREAD selection -> wire -> +// receive) rather than the gossipsub fallback path. +func TestGossipsubSpreadSelectionDelivers(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + hosts := getDefaultHosts(t, 2) + connect(t, hosts[0], hosts[1]) + + topicName := "spread-selection-delivers" + payload := []byte("spread-hello") + + opts := []Option{ + WithProtocolChoice(SPREAD), + WithSpreadPropagationConfig(&SpreadConfig{ + IntraFanout: 1, + InterFanout: 1, + IntraRho: 1, + InterProb: 1, + }), + withSpreadExtensionAdvertiseForTests(), + } + + psubs := make([]*PubSub, len(hosts)) + for i := range hosts { + psubs[i] = getGossipsub(ctx, hosts[i], opts...) + } + + topics := make([]*Topic, len(psubs)) + for i, ps := range psubs { + topic, err := ps.Join(topicName) + if err != nil { + t.Fatal(err) + } + topics[i] = topic + } + + sub, err := topics[1].Subscribe() + if err != nil { + t.Fatal(err) + } + + // Give subscriptions and SPREAD extension advertisements time to propagate. + time.Sleep(2 * time.Second) + + if err := publishWithSpreadForTests(ctx, topics[0], payload); err != nil { + t.Fatal(err) + } + + recvCtx, recvCancel := context.WithTimeout(ctx, 5*time.Second) + defer recvCancel() + got, err := sub.Next(recvCtx) + if err != nil { + t.Fatalf("subscriber did not receive the spread message: %v", err) + } + if !bytes.Equal(got.GetData(), payload) { + t.Fatalf("unexpected payload: got %q want %q", got.GetData(), payload) + } +} + func TestGossipsubNoIDONTWANTToMessageSender(t *testing.T) { synctestTest(t, func(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) diff --git a/pb/rpc.pb.go b/pb/rpc.pb.go index d3925f21..73b81973 100644 --- a/pb/rpc.pb.go +++ b/pb/rpc.pb.go @@ -1,7 +1,7 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.36.6 -// protoc v7.34.1 +// protoc-gen-go v1.36.10 +// protoc v6.33.4 // source: rpc.proto package pubsub_pb @@ -31,6 +31,8 @@ type RPC struct { // must use field numbers larger than 0x200000 to be encoded with at least 4 // bytes TestExtension *TestExtension `protobuf:"bytes,6492434,opt,name=testExtension" json:"testExtension,omitempty"` + // SPREAD extension: advertises that a peer supports SPREAD propagation + Spread *SpreadExtension `protobuf:"bytes,6492435,opt,name=spread" json:"spread,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -100,6 +102,13 @@ func (x *RPC) GetTestExtension() *TestExtension { return nil } +func (x *RPC) GetSpread() *SpreadExtension { + if x != nil { + return x.Spread + } + return nil +} + type Message struct { state protoimpl.MessageState `protogen:"open.v1"` From []byte `protobuf:"bytes,1,opt,name=from" json:"from,omitempty"` @@ -524,6 +533,7 @@ type ControlExtensions struct { // Experimental extensions must use field numbers larger than 0x200000 to be // encoded with 4 bytes TestExtension *bool `protobuf:"varint,6492434,opt,name=testExtension" json:"testExtension,omitempty"` + Spread *bool `protobuf:"varint,6492435,opt,name=spread" json:"spread,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -572,6 +582,13 @@ func (x *ControlExtensions) GetTestExtension() bool { return false } +func (x *ControlExtensions) GetSpread() bool { + if x != nil && x.Spread != nil { + return *x.Spread + } + return false +} + type PeerInfo struct { state protoimpl.MessageState `protogen:"open.v1"` PeerID []byte `protobuf:"bytes,1,opt,name=peerID" json:"peerID,omitempty"` @@ -730,6 +747,53 @@ func (x *PartialMessagesExtension) GetPartsMetadata() []byte { return nil } +type SpreadExtension struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Set when a SPREAD node forwards a SPREAD-marked message, so the receiving + // peer knows the message is being disseminated via SPREAD and should keep + // propagating it using SPREAD selection. + SourceIsSpreadNode *bool `protobuf:"varint,1,opt,name=sourceIsSpreadNode" json:"sourceIsSpreadNode,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SpreadExtension) Reset() { + *x = SpreadExtension{} + mi := &file_rpc_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SpreadExtension) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SpreadExtension) ProtoMessage() {} + +func (x *SpreadExtension) ProtoReflect() protoreflect.Message { + mi := &file_rpc_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SpreadExtension.ProtoReflect.Descriptor instead. +func (*SpreadExtension) Descriptor() ([]byte, []int) { + return file_rpc_proto_rawDescGZIP(), []int{12} +} + +func (x *SpreadExtension) GetSourceIsSpreadNode() bool { + if x != nil && x.SourceIsSpreadNode != nil { + return *x.SourceIsSpreadNode + } + return false +} + type RPC_SubOpts struct { state protoimpl.MessageState `protogen:"open.v1"` Subscribe *bool `protobuf:"varint,1,opt,name=subscribe" json:"subscribe,omitempty"` // subscribe or unsubcribe @@ -747,7 +811,7 @@ type RPC_SubOpts struct { func (x *RPC_SubOpts) Reset() { *x = RPC_SubOpts{} - mi := &file_rpc_proto_msgTypes[12] + mi := &file_rpc_proto_msgTypes[13] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -759,7 +823,7 @@ func (x *RPC_SubOpts) String() string { func (*RPC_SubOpts) ProtoMessage() {} func (x *RPC_SubOpts) ProtoReflect() protoreflect.Message { - mi := &file_rpc_proto_msgTypes[12] + mi := &file_rpc_proto_msgTypes[13] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -807,14 +871,15 @@ var File_rpc_proto protoreflect.FileDescriptor const file_rpc_proto_rawDesc = "" + "\n" + - "\trpc.proto\x12\tpubsub.pb\"\xce\x03\n" + + "\trpc.proto\x12\tpubsub.pb\"\x85\x04\n" + "\x03RPC\x12<\n" + "\rsubscriptions\x18\x01 \x03(\v2\x16.pubsub.pb.RPC.SubOptsR\rsubscriptions\x12,\n" + "\apublish\x18\x02 \x03(\v2\x12.pubsub.pb.MessageR\apublish\x123\n" + "\acontrol\x18\x03 \x01(\v2\x19.pubsub.pb.ControlMessageR\acontrol\x12=\n" + "\apartial\x18\n" + " \x01(\v2#.pubsub.pb.PartialMessagesExtensionR\apartial\x12A\n" + - "\rtestExtension\x18\x92\xa2\x8c\x03 \x01(\v2\x18.pubsub.pb.TestExtensionR\rtestExtension\x1a\xa3\x01\n" + + "\rtestExtension\x18\x92\xa2\x8c\x03 \x01(\v2\x18.pubsub.pb.TestExtensionR\rtestExtension\x125\n" + + "\x06spread\x18\x93\xa2\x8c\x03 \x01(\v2\x1a.pubsub.pb.SpreadExtensionR\x06spread\x1a\xa3\x01\n" + "\aSubOpts\x12\x1c\n" + "\tsubscribe\x18\x01 \x01(\bR\tsubscribe\x12\x18\n" + "\atopicid\x18\x02 \x01(\tR\atopicid\x12(\n" + @@ -854,11 +919,12 @@ const file_rpc_proto_rawDesc = "" + "\x10ControlIDontWant\x12\x1e\n" + "\n" + "messageIDs\x18\x01 \x03(\tR\n" + - "messageIDs\"f\n" + + "messageIDs\"\x81\x01\n" + "\x11ControlExtensions\x12(\n" + "\x0fpartialMessages\x18\n" + " \x01(\bR\x0fpartialMessages\x12'\n" + - "\rtestExtension\x18\x92\xa2\x8c\x03 \x01(\bR\rtestExtension\"N\n" + + "\rtestExtension\x18\x92\xa2\x8c\x03 \x01(\bR\rtestExtension\x12\x19\n" + + "\x06spread\x18\x93\xa2\x8c\x03 \x01(\bR\x06spread\"N\n" + "\bPeerInfo\x12\x16\n" + "\x06peerID\x18\x01 \x01(\fR\x06peerID\x12*\n" + "\x10signedPeerRecord\x18\x02 \x01(\fR\x10signedPeerRecord\"\x0f\n" + @@ -867,7 +933,9 @@ const file_rpc_proto_rawDesc = "" + "\atopicID\x18\x01 \x01(\tR\atopicID\x12\x18\n" + "\agroupID\x18\x02 \x01(\fR\agroupID\x12&\n" + "\x0epartialMessage\x18\x03 \x01(\fR\x0epartialMessage\x12$\n" + - "\rpartsMetadata\x18\x04 \x01(\fR\rpartsMetadataB1Z/github.com/libp2p/go-libp2p-pubsub/pb;pubsub_pb" + "\rpartsMetadata\x18\x04 \x01(\fR\rpartsMetadata\"A\n" + + "\x0fSpreadExtension\x12.\n" + + "\x12sourceIsSpreadNode\x18\x01 \x01(\bR\x12sourceIsSpreadNodeB1Z/github.com/libp2p/go-libp2p-pubsub/pb;pubsub_pb" var ( file_rpc_proto_rawDescOnce sync.Once @@ -881,7 +949,7 @@ func file_rpc_proto_rawDescGZIP() []byte { return file_rpc_proto_rawDescData } -var file_rpc_proto_msgTypes = make([]protoimpl.MessageInfo, 13) +var file_rpc_proto_msgTypes = make([]protoimpl.MessageInfo, 14) var file_rpc_proto_goTypes = []any{ (*RPC)(nil), // 0: pubsub.pb.RPC (*Message)(nil), // 1: pubsub.pb.Message @@ -895,26 +963,28 @@ var file_rpc_proto_goTypes = []any{ (*PeerInfo)(nil), // 9: pubsub.pb.PeerInfo (*TestExtension)(nil), // 10: pubsub.pb.TestExtension (*PartialMessagesExtension)(nil), // 11: pubsub.pb.PartialMessagesExtension - (*RPC_SubOpts)(nil), // 12: pubsub.pb.RPC.SubOpts + (*SpreadExtension)(nil), // 12: pubsub.pb.SpreadExtension + (*RPC_SubOpts)(nil), // 13: pubsub.pb.RPC.SubOpts } var file_rpc_proto_depIdxs = []int32{ - 12, // 0: pubsub.pb.RPC.subscriptions:type_name -> pubsub.pb.RPC.SubOpts + 13, // 0: pubsub.pb.RPC.subscriptions:type_name -> pubsub.pb.RPC.SubOpts 1, // 1: pubsub.pb.RPC.publish:type_name -> pubsub.pb.Message 2, // 2: pubsub.pb.RPC.control:type_name -> pubsub.pb.ControlMessage 11, // 3: pubsub.pb.RPC.partial:type_name -> pubsub.pb.PartialMessagesExtension 10, // 4: pubsub.pb.RPC.testExtension:type_name -> pubsub.pb.TestExtension - 3, // 5: pubsub.pb.ControlMessage.ihave:type_name -> pubsub.pb.ControlIHave - 4, // 6: pubsub.pb.ControlMessage.iwant:type_name -> pubsub.pb.ControlIWant - 5, // 7: pubsub.pb.ControlMessage.graft:type_name -> pubsub.pb.ControlGraft - 6, // 8: pubsub.pb.ControlMessage.prune:type_name -> pubsub.pb.ControlPrune - 7, // 9: pubsub.pb.ControlMessage.idontwant:type_name -> pubsub.pb.ControlIDontWant - 8, // 10: pubsub.pb.ControlMessage.extensions:type_name -> pubsub.pb.ControlExtensions - 9, // 11: pubsub.pb.ControlPrune.peers:type_name -> pubsub.pb.PeerInfo - 12, // [12:12] is the sub-list for method output_type - 12, // [12:12] is the sub-list for method input_type - 12, // [12:12] is the sub-list for extension type_name - 12, // [12:12] is the sub-list for extension extendee - 0, // [0:12] is the sub-list for field type_name + 12, // 5: pubsub.pb.RPC.spread:type_name -> pubsub.pb.SpreadExtension + 3, // 6: pubsub.pb.ControlMessage.ihave:type_name -> pubsub.pb.ControlIHave + 4, // 7: pubsub.pb.ControlMessage.iwant:type_name -> pubsub.pb.ControlIWant + 5, // 8: pubsub.pb.ControlMessage.graft:type_name -> pubsub.pb.ControlGraft + 6, // 9: pubsub.pb.ControlMessage.prune:type_name -> pubsub.pb.ControlPrune + 7, // 10: pubsub.pb.ControlMessage.idontwant:type_name -> pubsub.pb.ControlIDontWant + 8, // 11: pubsub.pb.ControlMessage.extensions:type_name -> pubsub.pb.ControlExtensions + 9, // 12: pubsub.pb.ControlPrune.peers:type_name -> pubsub.pb.PeerInfo + 13, // [13:13] is the sub-list for method output_type + 13, // [13:13] is the sub-list for method input_type + 13, // [13:13] is the sub-list for extension type_name + 13, // [13:13] is the sub-list for extension extendee + 0, // [0:13] is the sub-list for field type_name } func init() { file_rpc_proto_init() } @@ -928,7 +998,7 @@ func file_rpc_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_rpc_proto_rawDesc), len(file_rpc_proto_rawDesc)), NumEnums: 0, - NumMessages: 13, + NumMessages: 14, NumExtensions: 0, NumServices: 0, }, diff --git a/pb/rpc.proto b/pb/rpc.proto index c21b9360..ae7b1968 100644 --- a/pb/rpc.proto +++ b/pb/rpc.proto @@ -32,6 +32,9 @@ message RPC { // bytes optional TestExtension testExtension = 6492434; + // SPREAD extension: advertises that a peer supports SPREAD propagation + optional SpreadExtension spread = 6492435; + } message Message { @@ -85,6 +88,8 @@ message ControlExtensions { // Experimental extensions must use field numbers larger than 0x200000 to be // encoded with 4 bytes optional bool testExtension = 6492434; + + optional bool spread = 6492435; } message PeerInfo { @@ -104,3 +109,10 @@ message PartialMessagesExtension { // An encoded representation of the parts a peer has and wants. optional bytes partsMetadata = 4; } + +message SpreadExtension { + // Set when a SPREAD node forwards a SPREAD-marked message, so the receiving + // peer knows the message is being disseminated via SPREAD and should keep + // propagating it using SPREAD selection. + optional bool sourceIsSpreadNode = 1; +} diff --git a/pubsub.go b/pubsub.go index fce5c544..99b51adb 100644 --- a/pubsub.go +++ b/pubsub.go @@ -275,6 +275,7 @@ type Message struct { ReceivedFrom peer.ID ValidatorData any Local bool + Spread bool } func (m *Message) GetFrom() peer.ID { @@ -961,6 +962,9 @@ func WithAppSpecificRpcInspector(inspector func(peer.ID, *RPC) error) Option { // processLoop handles all inputs arriving on the channels func (p *PubSub) processLoop(ctx context.Context) { defer func() { + if gs, ok := p.rt.(*GossipSubRouter); ok && gs.extensions != nil && gs.extensions.spreadState != nil { + gs.extensions.spreadState.ShutdownVivaldi() + } // Clean up go routines. for _, queue := range p.peers { queue.Close() @@ -1571,7 +1575,11 @@ func (p *PubSub) handleIncomingRPC(rpc *RPC) { continue } - msg := &Message{Message: pmsg, ID: "", ReceivedFrom: rpc.from, ValidatorData: nil, Local: false} + msg := &Message{Message: pmsg, ID: "", ReceivedFrom: rpc.from, ValidatorData: nil, Local: false, Spread: false} + // If the enclosing RPC marks the publish as SPREAD, forward that marker to the internal message + if rpc.GetSpread() != nil { + msg.Spread = rpc.GetSpread().GetSourceIsSpreadNode() + } if p.shouldPush(msg) { toPush = append(toPush, msg) } @@ -1631,6 +1639,9 @@ func (p *PubSub) shouldPush(msg *Message) bool { id := p.idGen.ID(msg) if p.seenMessage(id) { p.tracer.DuplicateMessage(msg) + if gs, ok := p.rt.(*GossipSubRouter); ok && gs.canRelaySeenSpreadDuplicate(msg) && gs.hasSpreadDuplicateRelayBudget(msg.GetTopic(), id) { + return true + } return false } @@ -1642,6 +1653,13 @@ func (p *PubSub) pushMsg(msg *Message) { src := msg.ReceivedFrom id := p.idGen.ID(msg) + if p.seenMessage(id) { + if gs, ok := p.rt.(*GossipSubRouter); ok && gs.canRelaySeenSpreadDuplicate(msg) && gs.consumeSpreadDuplicateRelayBudget(msg.GetTopic(), id) { + gs.Publish(msg) + } + return + } + if !p.val.Push(src, msg) { return } diff --git a/spread.md b/spread.md new file mode 100644 index 00000000..99a74440 --- /dev/null +++ b/spread.md @@ -0,0 +1,148 @@ +# SPREAD Clustering And Propagation + +This document describes how the current SPREAD implementation derives a peer's "local cluster" and the set of "inter" peers used by the propagation selector. + +## Files + +- Clustering state and computation: `spread_state.go` (`SpreadState.GetPropagationPeers`) +- Propagation peer selection (fanout/probabilities): `spread_propagation.go` (`SpreadPropagation.GetPeersForPropagation`) +- Publish hook (switches from gossipsub -> SPREAD when enabled): `gossipsub.go` (`GossipSubRouter.rpcs`) +- Config options: + - `WithSpreadClusteringConfig` in `extensions.go` + - `WithSpreadPropagationConfig` in `gossipsub.go` + - Vivaldi runner wiring via `WithVivaldi` in `extensions.go` + +## High-Level Flow + +SPREAD is only used when both conditions hold: + +- Router is configured for SPREAD mode: `WithProtocolChoice(SPREAD)` +- The message is marked as spread at the `Message` level: `msg.Spread == true` + +When SPREAD is active for a message, its selection fully replaces the default +gossipsub peer set (provided it is non-empty). If SPREAD selects no peers, the +router falls back to the default gossipsub `tosend` set. + +```mermaid +flowchart TD + A[Publish msg] --> B{SPREAD mode AND msg.Spread?} + B -- no --> C[Use default gossipsub tosend set] + B -- yes --> D[Compute clusterPeers + interPeers] + D --> E[Select forwarding peers via SpreadPropagation] + E --> F{Selected > 0?} + F -- yes --> G[Override tosend with SPREAD selection] + F -- no --> C +``` + +## Basic Concepts + +### Topic-local peer universe + +Clustering is topic-aware and ignores non-subscribed peers: + +- Start from the set of SPREAD peers subscribed to the topic. +- Exclude `self` for selection. + +This peer set comes from SPREAD extension topic tracking in `SpreadState` (updated when peers announce subscriptions). + +### Vivaldi coordinates + +The implementation uses the Vivaldi "height-vector" distance to estimate RTT-like proximity: + +``` +dist(a, b) = sqrt((ax - bx)^2 + (ay - by)^2) + ah + bh +``` + +Only peers with a known peer state (`vsvc.GetPeerState(peer) != nil`) are considered "known" for coordinate-based sorting. + +### Known vs unknown peers + +- Known peers: have a Vivaldi state recorded in the local `vivaldi.Service`. +- Unknown peers: are spread peers in the topic, but do not yet have a recorded Vivaldi state. + +Unknown peers are still eligible for inter-peer selection (see below). + +## Clustering Algorithm + +### Configuration + +`SpreadClusteringConfig` controls cluster sizing and ring partitioning: + +- `ClusterPct` (float in `(0, 1]`): percentage of topic peers (excluding self) considered to be in the local cluster. +- `NumRings` (int > 0): number of equal-sized rings to partition the non-cluster known peers. + +Defaults (if unset/invalid): + +- `ClusterPct = 0.25` +- `NumRings = 3` + +### Steps + +Given a topic `t` and `self`: + +1. Build `topicSet` = all SPREAD peers subscribed to `t`, excluding `self`. +2. Compute `clusterSize = ceil(ClusterPct * |topicSet|)`, clamped to a minimum of `1`. +3. Fetch a constantly updated global list of known SPREAD peers sorted by `dist(self, peer)` ascending. +4. Filter that sorted list by topic membership, producing `knownInTopic`. +5. The local cluster is the first `clusterSize` peers of `knownInTopic` (or fewer if not enough known peers). +6. The remaining *known* peers in-topic are partitioned into `NumRings` equal-sized rings, then flattened back into a single list. +7. The inter-peer list is: + - all remaining known peers (post-cluster, post-ring partitioning; currently flattened), plus + - all unknown peers in the topic. + +**Notes:** + +- Cluster membership is *local* (there is no global agreement). +- When Vivaldi has not converged or peer states are sparse, the cluster may be small (or empty) because only known peers can appear in the cluster. +- Unknown peers are not discarded; they are treated as inter peers until they system learns their coordinates. + +## Inter-Peer Selection + +Propagation selection is controlled by `SpreadConfig` (defaults shown): + +- Intra: + - Always pick 1 peer from `clusterPeers` if non-empty. + - With probability `IntraRho` (default `0.6`), pick up to `IntraFanout - 1` additional random cluster peers. +- Inter: + - With probability `InterProb` (default `0.8`), pick `InterFanout` random peers from `interPeers`. + +All selections are random sampling with replacement; the final forwarding set is deduplicated by using a map. + +## Rings + +### Current behavior + +Rings are computed but not used as distinct tiers during selection: + +- Known non-cluster peers are split into `NumRings` equal-sized buckets. +- Those rings are immediately flattened into a single `inter` list. +- Inter selection samples randomly across the whole flattened set (plus unknown peers). + +### How rings could be used in the future + +The existing ring partitioning allows introducing explicit multi-scale propagation without changing the clustering interface: + +**Example**: Per-ring quotas, such as picking 1 peer per ring, or exponentially less per farther rings. + +## Caching and Updates (Performance) + +To stay lightweight, clustering maintains a local cache of: + +- a list of known spread peers sorted by Vivaldi distance to `self` +- a set form of the same list for O(1) "known vs unknown" checks + +The cache is invalidated ("dirty") when: + +- SPREAD peer membership changes +- topic membership changes +- Vivaldi updates run (runner ticks or explicit updates) + +The expensive sort is recomputed when a caller needs propagation peers and the cache is dirty. + +## Selection Semantics + +When SPREAD is enabled for a message: + +1. Build SPREAD-selected peers (`cluster + inter` via `SpreadPropagation`). +2. If the selection is non-empty, it fully replaces the default gossipsub `tosend` set. +3. If the selection is empty (e.g. no eligible SPREAD peers in the topic), the router keeps the default gossipsub `tosend` set so the message is not dropped. diff --git a/spread_clustering_test.go b/spread_clustering_test.go new file mode 100644 index 00000000..fe0b1ffd --- /dev/null +++ b/spread_clustering_test.go @@ -0,0 +1,84 @@ +package pubsub + +import ( + "testing" + + "github.com/libp2p/go-libp2p/core/peer" +) + +func TestSplitIntoEqualRings(t *testing.T) { + peers := []peer.ID{"p1", "p2", "p3", "p4", "p5"} + rings := splitIntoEqualRings(peers, 3) + if len(rings) != 3 { + t.Fatalf("expected 3 rings, got %d", len(rings)) + } + if got := len(rings[0]); got != 2 { + t.Fatalf("expected first ring size 2, got %d", got) + } + if got := len(rings[1]); got != 2 { + t.Fatalf("expected second ring size 2, got %d", got) + } + if got := len(rings[2]); got != 1 { + t.Fatalf("expected third ring size 1, got %d", got) + } + flat := flattenRings(rings) + if len(flat) != len(peers) { + t.Fatalf("expected flattened len %d, got %d", len(peers), len(flat)) + } + for i := range peers { + if flat[i] != peers[i] { + t.Fatalf("unexpected flatten order at %d: expected %s got %s", i, peers[i], flat[i]) + } + } +} + +func TestSanitizeSpreadClusteringConfig(t *testing.T) { + cfg := sanitizeSpreadClusteringConfig(&SpreadClusteringConfig{ + ClusterPct: 0.4, + NumRings: 5, + }) + if cfg.ClusterPct != 0.4 { + t.Fatalf("expected cluster pct 0.4, got %f", cfg.ClusterPct) + } + if cfg.NumRings != 5 { + t.Fatalf("expected num rings 5, got %d", cfg.NumRings) + } + + cfg = sanitizeSpreadClusteringConfig(&SpreadClusteringConfig{ + ClusterPct: 2, + NumRings: -1, + }) + if cfg.ClusterPct != DefaultSpreadClusterPct { + t.Fatalf("expected default cluster pct %f, got %f", DefaultSpreadClusterPct, cfg.ClusterPct) + } + if cfg.NumRings != DefaultSpreadNumRings { + t.Fatalf("expected default num rings %d, got %d", DefaultSpreadNumRings, cfg.NumRings) + } +} + +func TestSpreadPropagationSkipsEmptyPeerID(t *testing.T) { + sp := NewSpreadPropagation(&SpreadConfig{ + IntraFanout: 1, + InterFanout: 0, + IntraRho: 1, + InterProb: 0, + }) + got := sp.GetPeersForPropagation("from", nil, nil) + if len(got) != 0 { + t.Fatalf("expected no peers, got %v", got) + } +} + +func TestSanitizeSpreadConfigDuplicateRepropagation(t *testing.T) { + cfg := sanitizeSpreadConfig(&SpreadConfig{ + DuplicateRepropagation: 3, + }) + if cfg.DuplicateRepropagation != 3 { + t.Fatalf("expected duplicate repropagation 3, got %d", cfg.DuplicateRepropagation) + } + + cfg = sanitizeSpreadConfig(nil) + if cfg.DuplicateRepropagation != DefaultSpreadDuplicateRepropagation { + t.Fatalf("expected default duplicate repropagation %d, got %d", DefaultSpreadDuplicateRepropagation, cfg.DuplicateRepropagation) + } +} diff --git a/spread_propagation.go b/spread_propagation.go new file mode 100644 index 00000000..8ee2b97e --- /dev/null +++ b/spread_propagation.go @@ -0,0 +1,150 @@ +package pubsub + +import ( + "math/rand" + + "github.com/libp2p/go-libp2p/core/peer" +) + +// SpreadConfig controls SPREAD peer selection: how many intra- and +// inter-cluster peers a message is forwarded to, and with what probability. +type SpreadConfig struct { + // Intra-cluster fanout: number of intra-cluster peers to select + IntraFanout int + // Inter-cluster fanout: number of inter-cluster peers to select + InterFanout int + // Intra-cluster Cobra-walk rho: probability of selecting intra-cluster peers + IntraRho float64 + // Inter-cluster communication probability: probability of selecting inter-cluster peers + InterProb float64 + // Number of times a seen SPREAD message may be re-propagated as a duplicate. + // This is the number of extra forwards after the first forward. + DuplicateRepropagation int + + // UseAngularInterPeers toggles whether inter-cluster peers are drawn using + // angular buckets from the spread clustering state instead of distance-only + // ordering. + UseAngularInterPeers bool +} + +// Default configuration values +const ( + DefaultSpreadIntraFanout = 3 + DefaultSpreadInterFanout = 8 + DefaultSpreadIntraRho = 0.6 + DefaultSpreadInterProb = 0.8 + DefaultSpreadDuplicateRepropagation = 0 +) + +// SpreadPropagation selects the peers a SPREAD message is forwarded to, given +// the local cluster and inter-cluster peer sets computed by SpreadState. +type SpreadPropagation struct { + config *SpreadConfig +} + +// DefaultSpreadConfig returns a SpreadConfig populated with the package defaults. +func DefaultSpreadConfig() *SpreadConfig { + return &SpreadConfig{ + IntraFanout: DefaultSpreadIntraFanout, + InterFanout: DefaultSpreadInterFanout, + IntraRho: DefaultSpreadIntraRho, + InterProb: DefaultSpreadInterProb, + DuplicateRepropagation: DefaultSpreadDuplicateRepropagation, + UseAngularInterPeers: false, + } +} + +func sanitizeSpreadConfig(cfg *SpreadConfig) *SpreadConfig { + out := DefaultSpreadConfig() + if cfg == nil { + return out + } + if cfg.IntraFanout > 0 { + out.IntraFanout = cfg.IntraFanout + } + if cfg.InterFanout > 0 { + out.InterFanout = cfg.InterFanout + } + if cfg.IntraRho >= 0 && cfg.IntraRho <= 1 { + out.IntraRho = cfg.IntraRho + } + if cfg.InterProb >= 0 && cfg.InterProb <= 1 { + out.InterProb = cfg.InterProb + } + if cfg.DuplicateRepropagation >= 0 { + out.DuplicateRepropagation = cfg.DuplicateRepropagation + } + out.UseAngularInterPeers = cfg.UseAngularInterPeers + return out +} + +// NewSpreadPropagation returns a SpreadPropagation using config, falling back +// to the package defaults for any unset or invalid field. +func NewSpreadPropagation(config *SpreadConfig) *SpreadPropagation { + return &SpreadPropagation{ + config: sanitizeSpreadConfig(config), + } +} + +// GetPeersForPropagation returns the peers to forward a SPREAD message to, +// sampling intra- and inter-cluster peers according to the configured fanout +// and probabilities. The sender (from) is never included in the result. +func (sp *SpreadPropagation) GetPeersForPropagation(from peer.ID, clusterPeers []peer.ID, interClusterPeers []peer.ID) []peer.ID { + + selectedPeers := make(map[peer.ID]struct{}) + + // Intra-cluster + + // Add mandatory cluster peer + if mandatory := SelectRandomPeerID(clusterPeers); mandatory != "" { + selectedPeers[mandatory] = struct{}{} + } + // Coin flip to decide whether to select more intra-cluster peers + if rand.Float64() < sp.config.IntraRho { + // Select [intra fanout - 1] more random ones + for i := 0; i < sp.config.IntraFanout-1; i++ { + peerID := SelectRandomPeerID(clusterPeers) + if peerID == "" { + break // No peers in the cluster + } + selectedPeers[peerID] = struct{}{} + } + } + + // Inter-cluster + + // Coin flip to decide whether to select any inter-cluster peers + if rand.Float64() < sp.config.InterProb { + // Select [inter fanout] random ones + if len(interClusterPeers) > 0 { + for i := 0; i < sp.config.InterFanout; i++ { + peerID := SelectRandomPeerID(interClusterPeers) + if peerID == "" { + break // No peers in the cluster + } + selectedPeers[peerID] = struct{}{} + } + } + } + + // Convert map to slice + forwardingPeers := make([]peer.ID, 0, len(selectedPeers)) + for peerID := range selectedPeers { + // don't send to sender + if peerID == from { + continue + } + forwardingPeers = append(forwardingPeers, peerID) + } + + return forwardingPeers +} + +// SelectRandomPeerID returns a uniformly random peer from peers, or the empty +// peer.ID if peers is empty. +func SelectRandomPeerID(peers []peer.ID) peer.ID { + if len(peers) == 0 { + return "" + } + return peers[rand.Intn(len(peers))] +} diff --git a/spread_state.go b/spread_state.go new file mode 100644 index 00000000..d7083f91 --- /dev/null +++ b/spread_state.go @@ -0,0 +1,803 @@ +package pubsub + +import ( + "context" + "math" + "math/rand" + "sort" + "sync" + "time" + + "github.com/libp2p/go-libp2p-pubsub/vivaldi" + "github.com/libp2p/go-libp2p/core/peer" +) + +// SpreadState tracks the SPREAD-capable peers (globally and per topic), holds +// the optional Vivaldi coordinate service, and computes the local cluster and +// inter-cluster peer sets used by SpreadPropagation. It is safe for concurrent +// use. +type SpreadState struct { + lk sync.RWMutex + // peers that advertise SPREAD + peers map[peer.ID]struct{} + // topic -> peers with spread extension that are subscribed to the topic + topics map[string]map[peer.ID]struct{} + + // Vivaldi integration + vivaldiService *vivaldi.Service + vivaldiConfig *VivaldiConfig + runnerStop func() + + // SPREAD clustering configuration and lightweight cache. + clusteringConfig *SpreadClusteringConfig + sortedKnownByRTT []peer.ID + sortedKnownSet map[peer.ID]struct{} + angleBuckets map[int][]peer.ID + angleBucketSize float64 + anglePctPerBucket float64 + cacheDirty bool +} + +// VivaldiConfig holds configurable parameters for Vivaldi/Newton updates and runner. +type VivaldiConfig struct { + Cc float64 + Ce float64 + Newton bool + OutlierThreshold float64 + Samples int + Interval time.Duration + NeighborSetSize int + IN1ThresholdMS float64 + IN2ThresholdMS float64 + IN3MADKRandom float64 + IN3MADKClose float64 + IN3MinSamples int +} + +// SpreadClusteringConfig controls how SPREAD candidates are partitioned. +type SpreadClusteringConfig struct { + // ClusterPct is the percentage of topic peers selected as "local cluster". + // It is applied to find the closest ones over all spread peers in the topic (excluding self). + ClusterPct float64 + // NumRings controls equal-sized ring partitioning of non-cluster known peers. + // Rings are unused now but can be explored in the future. + NumRings int + + // InterAngleDegrees is the bucket size (in degrees) used to pick inter-cluster peers. + InterAngleDegrees float64 + // InterPctPerAngle is the fraction of (topic peers excluding self) to pick per angle bucket. + InterPctPerAngle float64 +} + +const ( + DefaultSpreadClusterPct = 0.1 + DefaultSpreadNumRings = 3 + DefaultSpreadInterAngleDegrees = 45 + DefaultSpreadInterPctPerAngle = 0.10 +) + +// DefaultSpreadClusteringConfig returns a SpreadClusteringConfig populated with +// the package defaults. +func DefaultSpreadClusteringConfig() *SpreadClusteringConfig { + return &SpreadClusteringConfig{ + ClusterPct: DefaultSpreadClusterPct, + NumRings: DefaultSpreadNumRings, + InterAngleDegrees: DefaultSpreadInterAngleDegrees, + InterPctPerAngle: DefaultSpreadInterPctPerAngle, + } +} + +func sanitizeSpreadClusteringConfig(cfg *SpreadClusteringConfig) *SpreadClusteringConfig { + out := DefaultSpreadClusteringConfig() + if cfg == nil { + return out + } + if cfg.ClusterPct > 0 && cfg.ClusterPct <= 1 { + out.ClusterPct = cfg.ClusterPct + } + if cfg.NumRings > 0 { + out.NumRings = cfg.NumRings + } + if cfg.InterAngleDegrees > 0 && cfg.InterAngleDegrees <= 360 { + out.InterAngleDegrees = cfg.InterAngleDegrees + } + if cfg.InterPctPerAngle > 0 && cfg.InterPctPerAngle <= 1 { + out.InterPctPerAngle = cfg.InterPctPerAngle + } + return out +} + +// NewSpreadState returns an empty SpreadState with default clustering config +// and no Vivaldi service configured. +func NewSpreadState() *SpreadState { + return &SpreadState{ + peers: make(map[peer.ID]struct{}), + topics: make(map[string]map[peer.ID]struct{}), + vivaldiService: nil, + vivaldiConfig: nil, + runnerStop: nil, + clusteringConfig: DefaultSpreadClusteringConfig(), + cacheDirty: true, + } +} + +// AddPeer records p as a SPREAD-capable peer. +func (s *SpreadState) AddPeer(p peer.ID) { + s.lk.Lock() + defer s.lk.Unlock() + s.peers[p] = struct{}{} + s.cacheDirty = true +} + +// RemovePeer drops p from the SPREAD peer set and from every topic it was +// tracked under. +func (s *SpreadState) RemovePeer(p peer.ID) { + s.lk.Lock() + defer s.lk.Unlock() + delete(s.peers, p) + // Remove from all topics + for t := range s.topics { + delete(s.topics[t], p) + if len(s.topics[t]) == 0 { + delete(s.topics, t) + } + } + s.cacheDirty = true +} + +// AddPeerTopic records that SPREAD peer p is subscribed to topic. Peers that +// are not known SPREAD peers are ignored. +func (s *SpreadState) AddPeerTopic(topic string, p peer.ID) { + s.lk.Lock() + defer s.lk.Unlock() + if _, ok := s.peers[p]; !ok { + // not a spread peer; ignore + return + } + ps, ok := s.topics[topic] + if !ok { + ps = make(map[peer.ID]struct{}) + s.topics[topic] = ps + } + ps[p] = struct{}{} + s.cacheDirty = true +} + +// RemovePeerTopic stops tracking SPREAD peer p as subscribed to topic. +func (s *SpreadState) RemovePeerTopic(topic string, p peer.ID) { + s.lk.Lock() + defer s.lk.Unlock() + if ps, ok := s.topics[topic]; ok { + delete(ps, p) + if len(ps) == 0 { + delete(s.topics, topic) + } + } + s.cacheDirty = true +} + +// GetSpreadPeers returns a slice of spread-capable peers for the given topic. +func (s *SpreadState) GetSpreadPeers(topic string) []peer.ID { + s.lk.RLock() + defer s.lk.RUnlock() + ps, ok := s.topics[topic] + if !ok { + return nil + } + out := make([]peer.ID, 0, len(ps)) + for p := range ps { + out = append(out, p) + } + return out +} + +// ConfigureVivaldi wires a Vivaldi service and parameters into the SpreadState. +// Passing a nil vsvc disables Vivaldi. +func (s *SpreadState) ConfigureVivaldi(vsvc *vivaldi.Service, cfg *VivaldiConfig) { + s.lk.Lock() + defer s.lk.Unlock() + if s.runnerStop != nil && vsvc == nil { + s.runnerStop() + s.runnerStop = nil + } + s.vivaldiService = vsvc + s.vivaldiConfig = sanitizeVivaldiConfig(cfg) + s.cacheDirty = true +} + +// ConfigureClustering sets the clustering parameters, falling back to defaults +// for unset or invalid fields. +func (s *SpreadState) ConfigureClustering(cfg *SpreadClusteringConfig) { + s.lk.Lock() + defer s.lk.Unlock() + s.clusteringConfig = sanitizeSpreadClusteringConfig(cfg) + s.cacheDirty = true +} + +// UpdatePeerVivaldi performs an ExchangeAndUpdate for a single peer. +func (s *SpreadState) UpdatePeerVivaldi(ctx context.Context, p peer.ID) (*vivaldi.VivaldiState, error) { + s.lk.RLock() + vsvc := s.vivaldiService + vconf := s.vivaldiConfig + s.lk.RUnlock() + if vsvc == nil || vconf == nil { + return nil, nil + } + cfg := vivaldi.UpdateConfig{ + Cc: vconf.Cc, + Ce: vconf.Ce, + Newton: vconf.Newton, + OutlierThreshold: vconf.OutlierThreshold, + Samples: vconf.Samples, + Interval: vconf.Interval, + IN1CentroidThresholdMS: vconf.IN1ThresholdMS, + IN2ProjectionThreshold: vconf.IN2ThresholdMS, + IN3MADKRandom: vconf.IN3MADKRandom, + IN3MADKClose: vconf.IN3MADKClose, + IN3MinSamples: vconf.IN3MinSamples, + } + st, err := vsvc.ExchangeAndUpdate(ctx, p, cfg) + if err == nil { + s.lk.Lock() + s.cacheDirty = true + s.lk.Unlock() + } + return st, err +} + +// StartVivaldiRunner starts periodic exchanges to all known spread peers. +// If a runner is already running, it will be stopped and replaced. +func (s *SpreadState) StartVivaldiRunner() { + s.lk.Lock() + defer s.lk.Unlock() + if s.vivaldiService == nil || s.vivaldiConfig == nil { + return + } + if s.runnerStop != nil { + s.runnerStop() + } + // Start a dynamic runner here in SpreadState so peer list changes are respected. + stopCh := make(chan struct{}) + go func() { + interval := s.vivaldiConfig.Interval + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-stopCh: + return + case <-ticker.C: + // snapshot current peers + s.lk.RLock() + curPeers := make([]peer.ID, 0, len(s.peers)) + for p := range s.peers { + curPeers = append(curPeers, p) + } + vsvc := s.vivaldiService + interval = s.vivaldiConfig.Interval + ucfg := vivaldi.UpdateConfig{ + Cc: s.vivaldiConfig.Cc, + Ce: s.vivaldiConfig.Ce, + Newton: s.vivaldiConfig.Newton, + OutlierThreshold: s.vivaldiConfig.OutlierThreshold, + Samples: s.vivaldiConfig.Samples, + Interval: interval, + IN1CentroidThresholdMS: s.vivaldiConfig.IN1ThresholdMS, + IN2ProjectionThreshold: s.vivaldiConfig.IN2ThresholdMS, + IN3MADKRandom: s.vivaldiConfig.IN3MADKRandom, + IN3MADKClose: s.vivaldiConfig.IN3MADKClose, + IN3MinSamples: s.vivaldiConfig.IN3MinSamples, + } + s.lk.RUnlock() + if vsvc == nil { + continue + } + + selected, closePeers, randomPeers := selectNewtonNeighbors(curPeers, vsvc, s.vivaldiConfig.NeighborSetSize) + vsvc.SetNeighborSets(closePeers, randomPeers) + for _, p := range selected { + ctx, cancel := context.WithTimeout(context.Background(), interval) + _, _ = vsvc.ExchangeAndUpdate(ctx, p, ucfg) + cancel() + } + s.lk.Lock() + s.cacheDirty = true + s.lk.Unlock() + } + } + }() + s.runnerStop = func() { close(stopCh) } +} + +// StopVivaldiRunner stops the background runner if it is running. +func (s *SpreadState) StopVivaldiRunner() { + s.lk.Lock() + defer s.lk.Unlock() + if s.runnerStop != nil { + s.runnerStop() + s.runnerStop = nil + } +} + +// ShutdownVivaldi stops background work and closes the service stream handler. +func (s *SpreadState) ShutdownVivaldi() { + s.lk.Lock() + defer s.lk.Unlock() + if s.runnerStop != nil { + s.runnerStop() + s.runnerStop = nil + } + if s.vivaldiService != nil { + s.vivaldiService.Close() + s.vivaldiService = nil + } + s.cacheDirty = true +} + +// GetPropagationPeers returns cluster peers and inter-cluster peers for topic. +// If useAngularInter is true, inter-cluster peers are chosen using angular buckets +// around the Vivaldi coordinate; otherwise the distance-only ordering is used. +func (s *SpreadState) GetPropagationPeers(topic string, self peer.ID, useAngularInter bool) ([]peer.ID, []peer.ID) { + + // Refresh sorted distance cache before getting propagation peers. + s.refreshDistanceCache(self) + s.refreshAngleCache(self) + + // Get topic peers + s.lk.RLock() + topicPeers, ok := s.topics[topic] + if !ok || len(topicPeers) == 0 { + s.lk.RUnlock() + return nil, nil + } + cfg := s.clusteringConfig + sortedKnown := append([]peer.ID(nil), s.sortedKnownByRTT...) + knownSet := make(map[peer.ID]struct{}, len(s.sortedKnownSet)) + for p := range s.sortedKnownSet { + knownSet[p] = struct{}{} + } + topicSet := make(map[peer.ID]struct{}, len(topicPeers)) + for p := range topicPeers { + topicSet[p] = struct{}{} + } + s.lk.RUnlock() + + // Compute topic size + totalTopicPeers := 0 + for p := range topicSet { + if p != self { + totalTopicPeers++ + } + } + if totalTopicPeers == 0 { + return nil, nil + } + + // Compute cluster size + clusterSize := int(math.Ceil(float64(totalTopicPeers) * cfg.ClusterPct)) + if clusterSize < 1 { + clusterSize = 1 + } + + // Identify known peers in the topic + knownInTopic := make([]peer.ID, 0, len(topicSet)) + for _, p := range sortedKnown { + if p == self { + continue + } + if _, ok := topicSet[p]; ok { + knownInTopic = append(knownInTopic, p) + } + } + + // Get cluster peers + if clusterSize > len(knownInTopic) { + clusterSize = len(knownInTopic) + } + clusterPeers := append([]peer.ID(nil), knownInTopic[:clusterSize]...) + + clusterSet := make(map[peer.ID]struct{}, len(clusterPeers)) + for _, p := range clusterPeers { + clusterSet[p] = struct{}{} + } + + // Distance-based ordering only. + ringInter := func() ([]peer.ID, []peer.ID) { + knownRemainder := append([]peer.ID(nil), knownInTopic[clusterSize:]...) + + unknownInTopic := make([]peer.ID, 0, len(topicSet)) + for p := range topicSet { + if p == self { + continue + } + if _, inCluster := clusterSet[p]; inCluster { + continue + } + if _, known := knownSet[p]; known { + continue + } + unknownInTopic = append(unknownInTopic, p) + } + interPeers := append(knownRemainder, unknownInTopic...) + return clusterPeers, interPeers + } + + if !useAngularInter { + return ringInter() + } + + // Get inter-cluster peers using angle buckets. + numBuckets := int(math.Ceil(360.0 / cfg.InterAngleDegrees)) + if numBuckets < 1 { + numBuckets = 1 + } + + // If angle buckets are not ready or Vivaldi is disabled, fall back + if s.vivaldiService == nil || s.angleBuckets == nil || cfg.InterAngleDegrees <= 0 || cfg.InterPctPerAngle <= 0 { + return ringInter() + } + + interPeers := make([]peer.ID, 0) + used := make(map[peer.ID]struct{}) + for b := 0; b < numBuckets; b++ { + // Count how many topic peers fall into this angle bucket (excluding self). + bucketTopicCount := 0 + for _, p := range s.angleBuckets[b] { + if p == self { + continue + } + if _, ok := topicSet[p]; !ok { + continue + } + bucketTopicCount++ + } + if bucketTopicCount == 0 { + continue + } + kPerAngle := int(math.Ceil(float64(bucketTopicCount) * cfg.InterPctPerAngle)) + if kPerAngle < 1 { + kPerAngle = 1 + } + picked := 0 + for _, p := range s.angleBuckets[b] { + if picked >= kPerAngle { + break + } + if _, ok := topicSet[p]; !ok { + continue + } + if _, inCluster := clusterSet[p]; inCluster { + continue + } + if _, already := used[p]; already { + continue + } + interPeers = append(interPeers, p) + used[p] = struct{}{} + picked++ + } + } + return clusterPeers, interPeers +} + +// splitPeersForTesting is a helper to split peers into rings. Unused for now. +func splitIntoEqualRings(peers []peer.ID, numRings int) [][]peer.ID { + if len(peers) == 0 { + return nil + } + if numRings <= 0 { + numRings = 1 + } + if numRings > len(peers) { + numRings = len(peers) + } + + rings := make([][]peer.ID, 0, numRings) + base := len(peers) / numRings + rem := len(peers) % numRings + start := 0 + for i := 0; i < numRings; i++ { + size := base + if i < rem { + size++ + } + end := start + size + rings = append(rings, append([]peer.ID(nil), peers[start:end]...)) + start = end + } + return rings +} + +// flattenRings is a helper to flatten rings into a single slice. Unused for now. +func flattenRings(rings [][]peer.ID) []peer.ID { + if len(rings) == 0 { + return nil + } + total := 0 + for _, ring := range rings { + total += len(ring) + } + out := make([]peer.ID, 0, total) + for _, ring := range rings { + out = append(out, ring...) + } + return out +} + +func (s *SpreadState) refreshDistanceCache(self peer.ID) { + s.lk.RLock() + + // If cache is clean, no need to refresh. + needsRefresh := s.cacheDirty + if !needsRefresh { + s.lk.RUnlock() + return + } + // Get vivaldi service and peers + vsvc := s.vivaldiService + peers := make([]peer.ID, 0, len(s.peers)) + for p := range s.peers { + peers = append(peers, p) + } + s.lk.RUnlock() + + // If no vivaldi service or no peers, reset cache to empty. + if vsvc == nil { + s.lk.Lock() + s.sortedKnownByRTT = nil + s.sortedKnownSet = nil + s.angleBuckets = nil + s.cacheDirty = false + s.lk.Unlock() + return + } + + // If no local coordinate, we can't compute distances, so reset cache to unsorted. + local := vsvc.GetLocalState() + if local == nil { + s.lk.Lock() + s.sortedKnownByRTT = nil + s.sortedKnownSet = nil + s.angleBuckets = nil + s.cacheDirty = false + s.lk.Unlock() + return + } + + // Compute distances to known peers and sort by distance. + type distanceEntry struct { + id peer.ID + dist float64 + } + known := make([]distanceEntry, 0, len(peers)) + for _, p := range peers { + if p == self { + continue + } + // If peer has no coordinate, we can't compute distance, so treat as unknown + st := vsvc.GetPeerState(p) + if st == nil { + continue + } + known = append(known, distanceEntry{ + id: p, + dist: vivaldi.Distance(local.Coord, st.Coord), + }) + } + // Sort by distance + sort.Slice(known, func(i, j int) bool { + return known[i].dist < known[j].dist + }) + + // Extract sorted peer IDs and sets for quick lookup. + sorted := make([]peer.ID, 0, len(known)) + knownSet := make(map[peer.ID]struct{}, len(known)) + for _, entry := range known { + sorted = append(sorted, entry.id) + knownSet[entry.id] = struct{}{} + } + + // Update cache + s.lk.Lock() + s.sortedKnownByRTT = sorted + s.sortedKnownSet = knownSet + s.cacheDirty = false + s.lk.Unlock() +} + +func (s *SpreadState) refreshAngleCache(self peer.ID) { + s.lk.RLock() + needsRefresh := s.cacheDirty + vsvc := s.vivaldiService + cfg := s.clusteringConfig + peers := make([]peer.ID, 0, len(s.peers)) + for p := range s.peers { + peers = append(peers, p) + } + s.lk.RUnlock() + + if !needsRefresh && s.angleBuckets != nil && s.angleBucketSize == cfg.InterAngleDegrees && s.anglePctPerBucket == cfg.InterPctPerAngle { + return + } + if vsvc == nil || cfg == nil || cfg.InterAngleDegrees <= 0 || cfg.InterAngleDegrees > 360 { + s.lk.Lock() + s.angleBuckets = nil + s.angleBucketSize = 0 + s.anglePctPerBucket = 0 + s.lk.Unlock() + return + } + local := vsvc.GetLocalState() + if local == nil { + s.lk.Lock() + s.angleBuckets = nil + s.angleBucketSize = cfg.InterAngleDegrees + s.anglePctPerBucket = cfg.InterPctPerAngle + s.lk.Unlock() + return + } + + numBuckets := int(math.Ceil(360.0 / cfg.InterAngleDegrees)) + if numBuckets < 1 { + numBuckets = 1 + } + type entry struct { + id peer.ID + dist float64 + } + byBucket := make(map[int][]entry, numBuckets) + + radPerBucket := (cfg.InterAngleDegrees * math.Pi) / 180.0 + for _, p := range peers { + if p == self { + continue + } + st := vsvc.GetPeerState(p) + if st == nil { + continue + } + dx := st.Coord.X - local.Coord.X + dy := st.Coord.Y - local.Coord.Y + theta := math.Atan2(dy, dx) + if theta < 0 { + theta += 2 * math.Pi + } + b := int(theta / radPerBucket) + if b < 0 { + b = 0 + } + if b >= numBuckets { + b = numBuckets - 1 + } + byBucket[b] = append(byBucket[b], entry{ + id: p, + dist: vivaldi.Distance(local.Coord, st.Coord), + }) + } + + out := make(map[int][]peer.ID, numBuckets) + for b := 0; b < numBuckets; b++ { + ents := byBucket[b] + sort.Slice(ents, func(i, j int) bool { return ents[i].dist < ents[j].dist }) + ids := make([]peer.ID, 0, len(ents)) + for _, e := range ents { + ids = append(ids, e.id) + } + out[b] = ids + } + + s.lk.Lock() + s.angleBuckets = out + s.angleBucketSize = cfg.InterAngleDegrees + s.anglePctPerBucket = cfg.InterPctPerAngle + s.lk.Unlock() +} + +func sanitizeVivaldiConfig(cfg *VivaldiConfig) *VivaldiConfig { + out := &VivaldiConfig{ + Cc: 0.25, + Ce: 0.25, + Newton: true, + OutlierThreshold: 0, + Samples: 3, + Interval: 30 * time.Second, + NeighborSetSize: 64, + IN1ThresholdMS: 20, + IN2ThresholdMS: 35, + IN3MADKRandom: 5, + IN3MADKClose: 8, + IN3MinSamples: 8, + } + if cfg == nil { + return out + } + out.Cc = cfg.Cc + out.Ce = cfg.Ce + out.Newton = cfg.Newton + out.OutlierThreshold = cfg.OutlierThreshold + out.Samples = cfg.Samples + out.Interval = cfg.Interval + out.NeighborSetSize = cfg.NeighborSetSize + out.IN1ThresholdMS = cfg.IN1ThresholdMS + out.IN2ThresholdMS = cfg.IN2ThresholdMS + out.IN3MADKRandom = cfg.IN3MADKRandom + out.IN3MADKClose = cfg.IN3MADKClose + out.IN3MinSamples = cfg.IN3MinSamples + if out.Cc <= 0 || out.Cc > 1 { + out.Cc = 0.25 + } + if out.Ce <= 0 || out.Ce > 1 { + out.Ce = 0.25 + } + if out.Samples <= 0 { + out.Samples = 3 + } + if out.Interval <= 0 { + out.Interval = 30 * time.Second + } + if out.NeighborSetSize <= 0 { + out.NeighborSetSize = 64 + } + if out.IN1ThresholdMS <= 0 { + out.IN1ThresholdMS = 20 + } + if out.IN2ThresholdMS <= 0 { + out.IN2ThresholdMS = 35 + } + if out.IN3MADKRandom <= 0 { + out.IN3MADKRandom = 5 + } + if out.IN3MADKClose <= 0 { + out.IN3MADKClose = 8 + } + if out.IN3MinSamples <= 0 { + out.IN3MinSamples = 8 + } + return out +} + +func selectNewtonNeighbors(peers []peer.ID, svc *vivaldi.Service, size int) ([]peer.ID, []peer.ID, []peer.ID) { + if len(peers) == 0 { + return nil, nil, nil + } + if size <= 0 || size > len(peers) { + size = len(peers) + } + type rttEntry struct { + id peer.ID + rtt float64 + } + rttKnown := make([]rttEntry, 0, len(peers)) + for _, p := range peers { + if rtt, ok := svc.GetPeerRTTMS(p); ok { + rttKnown = append(rttKnown, rttEntry{id: p, rtt: rtt}) + } + } + sort.Slice(rttKnown, func(i, j int) bool { return rttKnown[i].rtt < rttKnown[j].rtt }) + + closeTarget := size / 2 + closePeers := make([]peer.ID, 0, closeTarget) + used := make(map[peer.ID]struct{}, size) + + for i := 0; i < len(rttKnown) && len(closePeers) < closeTarget; i++ { + p := rttKnown[i].id + closePeers = append(closePeers, p) + used[p] = struct{}{} + } + + pool := make([]peer.ID, 0, len(peers)-len(closePeers)) + for _, p := range peers { + if _, ok := used[p]; ok { + continue + } + pool = append(pool, p) + } + rand.Shuffle(len(pool), func(i, j int) { pool[i], pool[j] = pool[j], pool[i] }) + randomTarget := size - len(closePeers) + if randomTarget > len(pool) { + randomTarget = len(pool) + } + randomPeers := append([]peer.ID(nil), pool[:randomTarget]...) + selected := append(append([]peer.ID(nil), closePeers...), randomPeers...) + return selected, closePeers, randomPeers +} diff --git a/vivaldi/README.md b/vivaldi/README.md new file mode 100644 index 00000000..10cad4b2 --- /dev/null +++ b/vivaldi/README.md @@ -0,0 +1,143 @@ +# Vivaldi + Newton Module + +## Scope +This package provides: +- A decentralized Vivaldi coordinate service over libp2p streams. +- A periodic runner that exchanges coordinate samples with peers. +- Newton-style safety checks (IN1, IN2, IN3) to reject malicious or inconsistent updates. + +The implementation uses: +- RTT in milliseconds. +- Height-vector coordinates (same as in Vivaldi paper). +- Per-update discard on invariant violation. + +## Runtime Flow +1. `SpreadState.StartVivaldiRunner` snapshots spread peers. +2. It selects a neighbor set each round: +- half low-latency peers (`close`) from observed RTT history. +- half random peers (`random`) from remaining peers. +3. Neighbor classes are installed in `Service.SetNeighborSets`. +4. For each selected peer, `Service.ExchangeAndUpdate`: +- performs `Samples` exchanges. +- picks the median RTT sample. +- computes the Vivaldi update. +- runs Newton checks (IN2, IN1, IN3). +- if checks pass, commits local and peer state. +- if any check fails, discards the update. + +```mermaid +sequenceDiagram + participant Runner as VivaldiRunner + participant Service as VivaldiService + participant Peer as RemotePeer + + Runner->>Service: StartVivaldiRunner() + loop every round + Runner->>Service: SelectNeighborSets() + Service->>Peer: ExchangeAndUpdate() + Peer-->>Service: RTT samples + Service->>Service: ComputeVivaldiUpdate() + Service->>Service: RunNewtonChecks() + alt Checks pass + Service->>Service: CommitState() + else Checks fail + Service->>Service: DiscardUpdate() + end + end +``` + +## Main State +### Coordinate state +- `local`: local `VivaldiState` (`Coord`, `Error`). +- `peerStates`: last known remote state per peer. +- `peerRTTms`: last median RTT per peer (for close/random selection). + +### Newton state +- `closePeers`, `randomPeers`: current round classes. +- `peerMeta`: metadata per peer for Newton checks, including: + - `LastReported`: last reported coordinate for a peer. + - `ExpectedMove`: accumulated expected displacement (IN2). + - `HasLast`: whether previous coordinate exists. + - `forceClose`, `forceRandom`: rolling force magnitude history for IN3. + +## Vivaldi Update (Algorithm 1) +Given local node $i$, remote node $j$, measured RTT $\mathrm{RTT}_{ij}$: + +Height-vector subtraction: +$$ +[x_i, h_i] - [x_j, h_j] = [x_i - x_j,\; h_i + h_j] +$$ + +Height-vector norm: +$$ +\|[x, h]\| = \sqrt{x_x^2 + x_y^2} + h +$$ + +Predicted RTT: +$$ +\widehat{\mathrm{RTT}}_{ij} = \|[x_i, h_i] - [x_j, h_j]\| +$$ + +Sample weight: +$$ +w = \frac{e_i}{e_i + e_j} +$$ + +Sample relative error: +$$ +e_s = \frac{\left|\widehat{\mathrm{RTT}}_{ij} - \mathrm{RTT}_{ij}\right|}{\mathrm{RTT}_{ij}} +$$ + +Error update: +$$ +\alpha = c_e \cdot w,\qquad +e_i' = \alpha e_s + (1-\alpha)e_i +$$ + +Coordinate step: +$$ +\delta = c_c \cdot w,\qquad +x_i' = x_i + \delta\left(\mathrm{RTT}_{ij} - \widehat{\mathrm{RTT}}_{ij}\right)u\!\left([x_i,h_i]-[x_j,h_j]\right) +$$ + +Where $u(\cdot)$ is the unit vector over the height-vector representation. + +## Newton Checks +Checks are run before committing an update. + +### IN1: Centroid consistency (random set) +- Uses only the current `randomPeers` set plus local node. +- Computes centroid and its distance to origin. +- Reject if: +$$ +\|\mathrm{centroid}\| > T_{\mathrm{IN1}} +$$ +- Default $T_{\mathrm{IN1}} = 20\text{ ms}$. + +### IN2: Physically-close projection consistency +- Applies only to peers in current `closePeers`. +- Compares observed displacement vs expected projected displacement. +- Let: +- $\Delta x_{\mathrm{obs}} =$ incoming minus last reported coordinate. +- $\Delta x_{\mathrm{exp}} =$ accumulated projected expected move. +- Reject if: +$$ +\|\Delta x_{\mathrm{obs}} - \Delta x_{\mathrm{exp}}\|_2 > T_{\mathrm{IN2}} +$$ +- Uses plain Euclidean norm over $(X,Y,H)$. +- Default $T_{\mathrm{IN2}} = 35\text{ ms}$. + +### IN3: Force outlier detection (MAD) +- Uses force magnitude history per class (`close` vs `random`). +- Computes median $\tilde{f}$ and MAD $D$ over history. +- Reject if: +$$ +f_j > \tilde{f} + K\cdot D +$$ +- Defaults: +- random peers: $K=5$. +- close peers: $K=8$. + +## Violation Handling +- Any violated invariant causes immediate discard of that update, without state changes. + diff --git a/vivaldi/runner.go b/vivaldi/runner.go new file mode 100644 index 00000000..083868e3 --- /dev/null +++ b/vivaldi/runner.go @@ -0,0 +1,304 @@ +package vivaldi + +import ( + "context" + "errors" + "fmt" + "sort" + "time" + + peerpkg "github.com/libp2p/go-libp2p/core/peer" +) + +// UpdateConfig configures update and Newton guard behavior. +type UpdateConfig struct { + Ce float64 + Cc float64 + Newton bool + OutlierThreshold float64 + Samples int + Interval time.Duration + + // Newton checks + IN1CentroidThresholdMS float64 + IN2ProjectionThreshold float64 + IN3MADKRandom float64 + IN3MADKClose float64 + IN3MinSamples int +} + +type sample struct { + rtt time.Duration + coord Coord + err float64 + errObj error +} + +// ExchangeAndUpdate performs N exchanges with a peer, picks median RTT sample, +// applies Newton checks (if enabled), and updates local coordinates. +func (s *Service) ExchangeAndUpdate(ctx context.Context, pid peerpkg.ID, cfg UpdateConfig) (*VivaldiState, error) { + if cfg.Samples <= 0 { + cfg.Samples = 1 + } + if cfg.Ce <= 0 { + cfg.Ce = 0.25 + } + if cfg.Cc <= 0 { + cfg.Cc = 0.25 + } + if cfg.IN3MinSamples <= 0 { + cfg.IN3MinSamples = 8 + } + + smpls := make([]sample, 0, cfg.Samples) + for i := 0; i < cfg.Samples; i++ { + rtt, coord, remoteErr, err := s.ExchangeOnce(ctx, pid) + smpls = append(smpls, sample{rtt: rtt, coord: coord, err: remoteErr, errObj: err}) + if i+1 < cfg.Samples { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(50 * time.Millisecond): + } + } + } + + good := make([]sample, 0, len(smpls)) + for _, sp := range smpls { + if sp.errObj == nil { + good = append(good, sp) + } + } + if len(good) == 0 { + return nil, errors.New("all exchanges failed") + } + chosen := good[indexOfMedianDuration(good)] + + s.lk.Lock() + updateCount := s.peerUpdateCounts[pid] + s.peerUpdateCounts[pid] = updateCount + 1 + warmup := updateCount < s.invariantWarmupUpdates + + if s.local == nil { + s.local = &VivaldiState{Coord: Coord{0, 0, 0}, Error: 1.0} + } + localCopy := *s.local + s.lk.Unlock() + + params := VivaldiUpdateParams{ + Local: localCopy, + Remote: VivaldiState{Coord: chosen.coord, Error: chosen.err}, + RTT: float64(chosen.rtt.Milliseconds()), + Ce: cfg.Ce, + Cc: cfg.Cc, + Newton: cfg.Newton, + OutlierThreshold: cfg.OutlierThreshold, + } + res, err := UpdateVivaldi(params) + if err != nil { + return nil, err + } + + if cfg.Newton { + // IN1, IN2 and IN3 are treated as soft invariants during warmup (only log), + // and as hard invariants afterwards (log + reject). + if err := s.checkIN2(pid, chosen.coord, cfg); err != nil { + s.logInvariantReject(pid, "IN2", err.Error()) + if !warmup { + return nil, err + } + } + if err := s.checkIN1(pid, chosen.coord, cfg); err != nil { + s.logInvariantReject(pid, "IN1", err.Error()) + if !warmup { + return nil, err + } + } + if err := s.checkIN3(pid, res.ForceMagnitude, cfg); err != nil { + s.logInvariantReject(pid, "IN3", err.Error()) + if !warmup { + return nil, err + } + } + } + + s.lk.Lock() + s.peerStates[pid] = &VivaldiState{Coord: chosen.coord, Error: chosen.err} + s.peerRTTms[pid] = float64(chosen.rtt.Milliseconds()) + s.local = &VivaldiState{Coord: res.NewCoord, Error: res.NewError} + + // Update IN2 expected movement tracking for physically-close peers. + if cfg.Newton { + s.updateIN2Expectations(pid, res.ForceVector) + meta := s.peerMeta[pid] + meta.LastReported = chosen.coord + meta.HasLast = true + meta.ExpectedMove = Coord{} + s.peerMeta[pid] = meta + + cls := s.peerClassLocked(pid) + s.recordForceSampleLocked(cls, res.ForceMagnitude) + } + newLocal := *s.local + s.lk.Unlock() + return &newLocal, nil +} + +func (s *Service) checkIN1(pid peerpkg.ID, incoming Coord, cfg UpdateConfig) error { + if cfg.IN1CentroidThresholdMS <= 0 { + return nil + } + s.lk.RLock() + defer s.lk.RUnlock() + + // IN1: centroid of local + random neighbors should stay close to origin. + count := 1.0 + sum := s.local.Coord + for rp := range s.randomPeers { + ps, ok := s.peerStates[rp] + if !ok { + continue + } + c := ps.Coord + if rp == pid { + c = incoming + } + sum = hvAdd(sum, c) + count++ + } + if count < 2 { + return nil + } + centroid := hvScale(sum, 1.0/count) + if hvNorm(centroid) > cfg.IN1CentroidThresholdMS { + return fmt.Errorf("centroid drift %.2fms exceeds threshold %.2fms", hvNorm(centroid), cfg.IN1CentroidThresholdMS) + } + return nil +} + +func (s *Service) checkIN2(pid peerpkg.ID, incoming Coord, cfg UpdateConfig) error { + if cfg.IN2ProjectionThreshold <= 0 { + return nil + } + s.lk.RLock() + defer s.lk.RUnlock() + if _, ok := s.closePeers[pid]; !ok { + return nil + } + meta := s.peerMeta[pid] + if !meta.HasLast { + return nil + } + observed := hvDiff(incoming, meta.LastReported) + diff := hvDiff(observed, meta.ExpectedMove) + if hvNorm(diff) > cfg.IN2ProjectionThreshold { + return fmt.Errorf("projection mismatch %.2fms exceeds threshold %.2fms", hvNorm(diff), cfg.IN2ProjectionThreshold) + } + return nil +} + +func (s *Service) checkIN3(pid peerpkg.ID, forceMag float64, cfg UpdateConfig) error { + s.lk.RLock() + defer s.lk.RUnlock() + class := s.peerClassLocked(pid) + + var hist []float64 + k := cfg.IN3MADKRandom + if class == "close" { + hist = s.forceClose + k = cfg.IN3MADKClose + } else { + hist = s.forceRandom + } + if len(hist) < cfg.IN3MinSamples || k <= 0 { + return nil + } + + med := median(hist) + madev := mad(hist, med) + // Preserve robustness if all previous samples were identical. + if madev < 1e-6 { + madev = 1e-6 + } + limit := med + k*madev + if forceMag > limit { + return fmt.Errorf("force %.2fms exceeds median+K*MAD %.2fms (median=%.2f, MAD=%.2f, K=%.2f)", forceMag, limit, med, madev, k) + } + return nil +} + +func (s *Service) updateIN2Expectations(source peerpkg.ID, forceVec Coord) { + for k := range s.closePeers { + if k == source { + continue + } + ps, ok := s.peerStates[k] + if !ok { + continue + } + src, ok := s.peerStates[source] + if !ok { + continue + } + dir := hvUnit(hvSub(src.Coord, ps.Coord)) + exp := hvProjection(forceVec, dir) + meta := s.peerMeta[k] + meta.ExpectedMove = hvAdd(meta.ExpectedMove, exp) + s.peerMeta[k] = meta + } +} + +func (s *Service) peerClassLocked(pid peerpkg.ID) string { + if _, ok := s.closePeers[pid]; ok { + return "close" + } + return "random" +} + +func (s *Service) recordForceSampleLocked(class string, v float64) { + const maxHist = 256 + if class == "close" { + s.forceClose = append(s.forceClose, v) + if len(s.forceClose) > maxHist { + s.forceClose = s.forceClose[len(s.forceClose)-maxHist:] + } + return + } + s.forceRandom = append(s.forceRandom, v) + if len(s.forceRandom) > maxHist { + s.forceRandom = s.forceRandom[len(s.forceRandom)-maxHist:] + } +} + +func indexOfMedianDuration(s []sample) int { + n := len(s) + idxs := make([]int, n) + for i := 0; i < n; i++ { + idxs[i] = i + } + sort.Slice(idxs, func(i, j int) bool { return s[idxs[i]].rtt < s[idxs[j]].rtt }) + return idxs[n/2] +} + +// StartPeriodicRunner starts a background goroutine that periodically exchanges with +// the provided peers and updates local coordinates. It returns a stop function. +func (s *Service) StartPeriodicRunner(peers []peerpkg.ID, cfg UpdateConfig) (stop func()) { + done := make(chan struct{}) + go func() { + ticker := time.NewTicker(cfg.Interval) + defer ticker.Stop() + for { + select { + case <-done: + return + case <-ticker.C: + for _, p := range peers { + ctx, cancel := context.WithTimeout(context.Background(), cfg.Interval) + _, _ = s.ExchangeAndUpdate(ctx, p, cfg) + cancel() + } + } + } + }() + return func() { close(done) } +} diff --git a/vivaldi/service.go b/vivaldi/service.go new file mode 100644 index 00000000..c2295707 --- /dev/null +++ b/vivaldi/service.go @@ -0,0 +1,209 @@ +package vivaldi + +import ( + "bufio" + "context" + "encoding/json" + "io" + "log" + "sync" + "time" + + hostpkg "github.com/libp2p/go-libp2p/core/host" + network "github.com/libp2p/go-libp2p/core/network" + peerpkg "github.com/libp2p/go-libp2p/core/peer" + "github.com/libp2p/go-libp2p/core/protocol" +) + +// ProtocolID is the vivaldi protocol used for exchanging coordinates. +const ProtocolID = "/libp2p/vivaldi/1.0.0" + +// Service exposes a small request/response handler for exchanging Coord+Error and measuring RTT by timing the round-trip. +type Service struct { + h hostpkg.Host + protocol protocol.ID + timeout time.Duration + lk sync.RWMutex + local *VivaldiState + peerStates map[peerpkg.ID]*VivaldiState + peerRTTms map[peerpkg.ID]float64 + + // Newton invariant context. + closePeers map[peerpkg.ID]struct{} + randomPeers map[peerpkg.ID]struct{} + peerMeta map[peerpkg.ID]peerInvariantState + forceClose []float64 + forceRandom []float64 + invariantWarmupUpdates int + peerUpdateCounts map[peerpkg.ID]int +} + +type peerInvariantState struct { + LastReported Coord + HasLast bool + ExpectedMove Coord +} + +// Config defines service options. +type Config struct { + Protocol protocol.ID + Timeout time.Duration +} + +// NewService creates a new Vivaldi service bound to the given host. +func NewService(h hostpkg.Host, cfg *Config) *Service { + proto := protocol.ID(ProtocolID) + t := 5 * time.Second + if cfg != nil { + if cfg.Protocol != "" { + proto = cfg.Protocol + } + if cfg.Timeout > 0 { + t = cfg.Timeout + } + } + s := &Service{ + h: h, + protocol: proto, + timeout: t, + peerStates: make(map[peerpkg.ID]*VivaldiState), + peerRTTms: make(map[peerpkg.ID]float64), + closePeers: make(map[peerpkg.ID]struct{}), + randomPeers: make(map[peerpkg.ID]struct{}), + peerMeta: make(map[peerpkg.ID]peerInvariantState), + invariantWarmupUpdates: 0, + peerUpdateCounts: make(map[peerpkg.ID]int), + } + h.SetStreamHandler(proto, s.handleStream) + return s +} + +// Close unregisters the stream handler. +func (s *Service) Close() { + s.h.RemoveStreamHandler(s.protocol) +} + +// response is the payload sent back to requesting peers. +type response struct { + Coord Coord `json:"coord"` + Error float64 `json:"error"` + Timestamp int64 `json:"ts"` // unix nanos +} + +// handleStream answers incoming vivaldi requests with the local coord and error. +func (s *Service) handleStream(st network.Stream) { + defer st.Close() + // Read the (optional) request payload; we don't expect anything large. + _ = st.SetDeadline(time.Now().Add(s.timeout)) + r := bufio.NewReader(st) + // read until EOF or newline; allow empty requests + _, _ = r.ReadBytes('\n') + + // Prepare response from an assumed global/local state provider. + // For now, attempt to obtain state via a package-level accessor. If not set, + // respond with zero coord and high error. + var resp response + s.lk.RLock() + ls := s.local + s.lk.RUnlock() + if ls != nil { + resp = response{Coord: ls.Coord, Error: ls.Error, Timestamp: time.Now().UnixNano()} + } else { + resp = response{Coord: Coord{0, 0, 0}, Error: 1e6, Timestamp: time.Now().UnixNano()} + } + + enc := json.NewEncoder(st) + _ = enc.Encode(&resp) +} + +// ExchangeOnce opens a stream to peer, performs a request/response, measures RTT, +// and returns (rtt, remoteCoord, remoteError, error). +func (s *Service) ExchangeOnce(ctx context.Context, pid peerpkg.ID) (time.Duration, Coord, float64, error) { + start := time.Now() + ctx, cancel := context.WithTimeout(ctx, s.timeout) + defer cancel() + + st, err := s.h.NewStream(ctx, pid, s.protocol) + if err != nil { + return 0, Coord{}, 0, err + } + defer st.Close() + // Set deadlines on the stream so slow peers don't block forever. + _ = st.SetDeadline(time.Now().Add(s.timeout)) + + // Send a small request (newline) to trigger the response. + if _, err := st.Write([]byte("\n")); err != nil { + return 0, Coord{}, 0, err + } + + // Read response fully. + data, err := io.ReadAll(st) + rtt := time.Since(start) + if err != nil { + return rtt, Coord{}, 0, err + } + + var resp response + if err := json.Unmarshal(data, &resp); err != nil { + return rtt, Coord{}, 0, err + } + return rtt, resp.Coord, resp.Error, nil +} + +// SetLocalState sets the service's local Vivaldi state returned to peers. +func (s *Service) SetLocalState(st *VivaldiState) { + s.lk.Lock() + defer s.lk.Unlock() + s.local = st +} + +// GetLocalState returns a copy of the local Vivaldi state. +func (s *Service) GetLocalState() *VivaldiState { + s.lk.RLock() + defer s.lk.RUnlock() + if s.local == nil { + return nil + } + tmp := *s.local + return &tmp +} + +// GetPeerState returns the last known Vivaldi state for a peer, or nil. +func (s *Service) GetPeerState(pid peerpkg.ID) *VivaldiState { + s.lk.RLock() + defer s.lk.RUnlock() + ps := s.peerStates[pid] + if ps == nil { + return nil + } + tmp := *ps + return &tmp +} + +// GetPeerRTTMS returns the most recently measured RTT to pid in milliseconds, +// and whether a measurement exists. +func (s *Service) GetPeerRTTMS(pid peerpkg.ID) (float64, bool) { + s.lk.RLock() + defer s.lk.RUnlock() + v, ok := s.peerRTTms[pid] + return v, ok +} + +// SetNeighborSets records the close and random neighbour sets used to select +// peers for Newton-Vivaldi coordinate exchanges. +func (s *Service) SetNeighborSets(closePeers []peerpkg.ID, randomPeers []peerpkg.ID) { + s.lk.Lock() + defer s.lk.Unlock() + s.closePeers = make(map[peerpkg.ID]struct{}, len(closePeers)) + s.randomPeers = make(map[peerpkg.ID]struct{}, len(randomPeers)) + for _, p := range closePeers { + s.closePeers[p] = struct{}{} + } + for _, p := range randomPeers { + s.randomPeers[p] = struct{}{} + } +} + +func (s *Service) logInvariantReject(pid peerpkg.ID, invariant string, details string) { + log.Printf("vivaldi[newton]: reject update peer=%s invariant=%s reason=%s", pid, invariant, details) +} diff --git a/vivaldi/service_test.go b/vivaldi/service_test.go new file mode 100644 index 00000000..48e76ef3 --- /dev/null +++ b/vivaldi/service_test.go @@ -0,0 +1,49 @@ +package vivaldi + +import ( + "context" + "testing" + "time" + + libp2p "github.com/libp2p/go-libp2p" + "github.com/libp2p/go-libp2p/core/peer" +) + +func TestExchangeOnce(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + h1, err := libp2p.New() + if err != nil { + t.Fatal(err) + } + defer h1.Close() + + h2, err := libp2p.New() + if err != nil { + t.Fatal(err) + } + defer h2.Close() + + s1 := NewService(h1, nil) + s2 := NewService(h2, nil) + defer s1.Close() + defer s2.Close() + + // Set local state on host2 so that host1 receives a non-zero coord + s2.SetLocalState(&VivaldiState{Coord: Coord{X: 1.0, Y: 2.0, H: 0.1}, Error: 0.05}) + if err := h1.Connect(ctx, peer.AddrInfo{ID: h2.ID(), Addrs: h2.Addrs()}); err != nil { + t.Fatal(err) + } + + rtt, coord, _, err := s1.ExchangeOnce(ctx, h2.ID()) + if err != nil { + t.Fatal(err) + } + if rtt <= 0 { + t.Fatalf("expected positive rtt, got %v", rtt) + } + if coord.X != 1.0 || coord.Y != 2.0 { + t.Fatalf("unexpected coord received: %v", coord) + } +} diff --git a/vivaldi/vivaldi.go b/vivaldi/vivaldi.go new file mode 100644 index 00000000..ec7dfbb1 --- /dev/null +++ b/vivaldi/vivaldi.go @@ -0,0 +1,218 @@ +// Package vivaldi implements the Vivaldi/Newton update logic for virtual coordinates. +package vivaldi + +import ( + "fmt" + "math" + "math/rand" +) + +// Coord represents a 2D Euclidean coordinate with a height component. +type Coord struct { + X float64 + Y float64 + H float64 // Height +} + +// VivaldiState holds the local coordinate and error estimate. +type VivaldiState struct { + Coord Coord + Error float64 +} + +// VivaldiUpdateParams holds the parameters for a single update step. +type VivaldiUpdateParams struct { + Local VivaldiState // Our node's state + Remote VivaldiState // State from other peer + RTT float64 // Measured RTT (ms) + Ce float64 // Error correction constant + Cc float64 // Coordinate correction constant + Newton bool // If true, apply Newton-Vivaldi security checks + OutlierThreshold float64 // Optional outlier threshold on force magnitude (ms) +} + +// UpdateResult holds the outcome of a single Vivaldi coordinate update, +// including the new coordinate and error plus diagnostic force and prediction +// values. +type UpdateResult struct { + NewCoord Coord + NewError float64 + Weight float64 + SampleRelError float64 + ForceVector Coord + ForceMagnitude float64 + PredictedRTTMS float64 +} + +// UpdateVivaldi performs the Vivaldi Algorithm 1 update with height-vector operations. +func UpdateVivaldi(params VivaldiUpdateParams) (UpdateResult, error) { + if params.RTT <= 0 { + return UpdateResult{}, fmt.Errorf("vivaldi: non-positive RTT %.4fms", params.RTT) + } + if params.Ce <= 0 || params.Cc <= 0 { + return UpdateResult{}, fmt.Errorf("vivaldi: invalid constants ce=%.4f cc=%.4f", params.Ce, params.Cc) + } + + // Height-vector distance prediction: ||[xi - xj, hi + hj]|| + diff := hvSub(params.Local.Coord, params.Remote.Coord) + predDist := hvNorm(diff) + forceMag := params.RTT - predDist + + // Newton outlier guard if requested. + if params.Newton && params.OutlierThreshold > 0 && math.Abs(forceMag) > params.OutlierThreshold { + return UpdateResult{}, fmt.Errorf( + "newton: outlier detected (|force|=%.4f > threshold=%.4f)", + math.Abs(forceMag), params.OutlierThreshold, + ) + } + + // Algorithm 1: w = ei / (ei + ej) + totalErr := params.Local.Error + params.Remote.Error + if totalErr <= 0 { + totalErr = 1e-6 + } + w := params.Local.Error / totalErr + if w < 0 { + w = 0 + } + if w > 1 { + w = 1 + } + + // e_s = |pred-rtt| / rtt + es := math.Abs(predDist-params.RTT) / params.RTT + + // alpha = ce * w + alpha := params.Ce * w + if alpha < 0 { + alpha = 0 + } + if alpha > 1 { + alpha = 1 + } + + // e_i = alpha*e_s + (1-alpha)*e_i + newErr := alpha*es + (1-alpha)*params.Local.Error + if newErr < 1e-6 { + newErr = 1e-6 + } + + // delta = cc * w + delta := params.Cc * w + + // x_i = x_i + delta * (rtt - ||xi-xj||) * u(xi-xj) + u := hvUnit(diff) + forceVec := hvScale(u, forceMag) + step := hvScale(u, delta*forceMag) + newCoord := hvAdd(params.Local.Coord, step) + if newCoord.H < 0 { + newCoord.H = 0 + } + + return UpdateResult{ + NewCoord: newCoord, + NewError: newErr, + Weight: w, + SampleRelError: es, + ForceVector: forceVec, + ForceMagnitude: math.Abs(forceMag), + PredictedRTTMS: predDist, + }, nil +} + +func euclideanDist(a, b Coord) float64 { + dx := a.X - b.X + dy := a.Y - b.Y + return math.Sqrt(dx*dx + dy*dy) +} + +func hvSub(a, b Coord) Coord { + return Coord{ + X: a.X - b.X, + Y: a.Y - b.Y, + H: a.H + b.H, + } +} + +func hvAdd(a, b Coord) Coord { + return Coord{ + X: a.X + b.X, + Y: a.Y + b.Y, + H: a.H + b.H, + } +} + +func hvScale(a Coord, s float64) Coord { + return Coord{ + X: a.X * s, + Y: a.Y * s, + H: a.H * s, + } +} + +func hvNorm(a Coord) float64 { + return math.Sqrt(a.X*a.X+a.Y*a.Y) + a.H +} + +func hvUnit(a Coord) Coord { + n := hvNorm(a) + if n == 0 { + // As in Vivaldi, break ties randomly if colocated. + theta := rand.Float64() * 2 * math.Pi + return Coord{X: math.Cos(theta), Y: math.Sin(theta), H: 0} + } + return hvScale(a, 1.0/n) +} + +func hvDot(a, b Coord) float64 { + return a.X*b.X + a.Y*b.Y + a.H*b.H +} + +func hvProjection(a, unitDir Coord) Coord { + return hvScale(unitDir, hvDot(a, unitDir)) +} + +func hvDiff(a, b Coord) Coord { + return Coord{X: a.X - b.X, Y: a.Y - b.Y, H: a.H - b.H} +} + +// Distance returns the Vivaldi height-vector distance between two coordinates: +// the Euclidean distance between their planar components plus both heights. +func Distance(a, b Coord) float64 { + return euclideanDist(a, b) + a.H + b.H +} + +func median(vals []float64) float64 { + if len(vals) == 0 { + return 0 + } + tmp := append([]float64(nil), vals...) + sortFloat64s(tmp) + n := len(tmp) + if n%2 == 1 { + return tmp[n/2] + } + return (tmp[n/2-1] + tmp[n/2]) / 2 +} + +func mad(vals []float64, med float64) float64 { + if len(vals) == 0 { + return 0 + } + dev := make([]float64, len(vals)) + for i, v := range vals { + dev[i] = math.Abs(v - med) + } + return median(dev) +} + +func sortFloat64s(v []float64) { + for i := 1; i < len(v); i++ { + x := v[i] + j := i - 1 + for ; j >= 0 && v[j] > x; j-- { + v[j+1] = v[j] + } + v[j+1] = x + } +} diff --git a/vivaldi/vivaldi_update_test.go b/vivaldi/vivaldi_update_test.go new file mode 100644 index 00000000..314e302a --- /dev/null +++ b/vivaldi/vivaldi_update_test.go @@ -0,0 +1,44 @@ +package vivaldi + +import ( + "math" + "testing" + + peerpkg "github.com/libp2p/go-libp2p/core/peer" +) + +func TestUpdateVivaldiAlgorithm1(t *testing.T) { + res, err := UpdateVivaldi(VivaldiUpdateParams{ + Local: VivaldiState{Coord: Coord{X: 0, Y: 0, H: 0}, Error: 1}, + Remote: VivaldiState{Coord: Coord{X: 10, Y: 0, H: 0}, Error: 1}, + RTT: 12, + Ce: 0.25, + Cc: 0.25, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if math.Abs(res.NewCoord.X-(-0.25)) > 1e-9 { + t.Fatalf("unexpected new x: %.12f", res.NewCoord.X) + } + if math.Abs(res.NewCoord.Y) > 1e-9 || math.Abs(res.NewCoord.H) > 1e-9 { + t.Fatalf("unexpected coord: %+v", res.NewCoord) + } + expectedErr := 0.125*(2.0/12.0) + 0.875*1.0 + if math.Abs(res.NewError-expectedErr) > 1e-9 { + t.Fatalf("unexpected new error: got %.12f want %.12f", res.NewError, expectedErr) + } +} + +func TestIN3RejectsOutlierForce(t *testing.T) { + s := &Service{ + forceRandom: []float64{10, 11, 12, 10, 9, 11, 10, 12, 9, 10}, + closePeers: map[peerpkg.ID]struct{}{}, + randomPeers: map[peerpkg.ID]struct{}{}, + } + cfg := UpdateConfig{IN3MADKRandom: 5, IN3MinSamples: 8} + if err := s.checkIN3(peerpkg.ID("peer1"), 100, cfg); err == nil { + t.Fatal("expected IN3 rejection") + } +}