Add deterministic routing

Signed-off-by: Andrew Poydence <[email protected]>

Author: thepeterstone

internal/node/node.go

@@ -2,20 +2,26 @@ package node
 
 type Node struct {
 	children      map[uint64]*Node
-	subscriptions map[string][]SubscriptionEnvelope
+	subscriptions map[string]subscriptionInfo
 	shards        map[int64]string
 	rand          func(int64) int64
 }
 
+type subscriptionInfo struct {
+	deterministicRoutingCount int
+	envelopes                 []SubscriptionEnvelope
+}
+
 type SubscriptionEnvelope struct {
 	Subscription func(interface{})
 	id           int64
+	dName        string
 }
 
 func New(int63n func(n int64) int64) *Node {
 	return &Node{
 		children:      make(map[uint64]*Node),
-		subscriptions: make(map[string][]SubscriptionEnvelope),
+		subscriptions: make(map[string]subscriptionInfo),
 		shards:        make(map[int64]string),
 		rand:          int63n,
 	}
@@ -59,16 +65,24 @@ func (n *Node) ChildLen() int {
 	return len(n.children)
 }
 
-func (n *Node) AddSubscription(s func(interface{}), shardID string) int64 {
+func (n *Node) AddSubscription(s func(interface{}), shardID, deterministicRoutingName string) int64 {
 	if n == nil {
 		return 0
 	}
 
 	id := n.createAndSetID(shardID)
-	n.subscriptions[shardID] = append(n.subscriptions[shardID], SubscriptionEnvelope{
+
+	si := n.subscriptions[shardID]
+	si.envelopes = append(si.envelopes, SubscriptionEnvelope{
 		Subscription: s,
 		id:           id,
+		dName:        deterministicRoutingName,
 	})
+	if deterministicRoutingName != "" {
+		si.deterministicRoutingCount++
+	}
+	n.subscriptions[shardID] = si
+
 	return id
 }
 
@@ -85,15 +99,21 @@ func (n *Node) DeleteSubscription(id int64) {
 	delete(n.shards, id)
 
 	s := n.subscriptions[shardID]
-	for i, ss := range s {
+	for i, ss := range s.envelopes {
 		if ss.id != id {
 			continue
 		}
 
-		n.subscriptions[shardID] = append(s[:i], s[i+1:]...)
+		if ss.dName != "" {
+			s.deterministicRoutingCount--
+		}
+
+		s.envelopes = append(s.envelopes[:i], s.envelopes[i+1:]...)
+		break
 	}
+	n.subscriptions[shardID] = s
 
-	if len(n.subscriptions[shardID]) == 0 {
+	if len(n.subscriptions[shardID].envelopes) == 0 {
 		delete(n.subscriptions, shardID)
 	}
 }
@@ -102,13 +122,13 @@ func (n *Node) SubscriptionLen() int {
 	return len(n.shards)
 }
 
-func (n *Node) ForEachSubscription(f func(shardID string, s []SubscriptionEnvelope)) {
+func (n *Node) ForEachSubscription(f func(shardID string, isDeterministic bool, s []SubscriptionEnvelope)) {
 	if n == nil {
 		return
 	}
 
 	for shardID, s := range n.subscriptions {
-		f(shardID, s)
+		f(shardID, s.deterministicRoutingCount > 0, s.envelopes)
 	}
 }
 

internal/node/node_test.go

@@ -50,26 +50,46 @@ func TestNode(t *testing.T) {
 	})
 
 	o.Spec("returns all subscriptions", func(t TN) {
-		id1 := t.n.AddSubscription(func(interface{}) {}, "")
+		id1 := t.n.AddSubscription(func(interface{}) {}, "", "")
 
-		t.n.AddSubscription(func(interface{}) {}, "")
-		t.n.AddSubscription(func(interface{}) {}, "")
+		t.n.AddSubscription(func(interface{}) {}, "", "")
+		t.n.AddSubscription(func(interface{}) {}, "", "")
 		t.n.DeleteSubscription(id1)
 
 		var ss []func(interface{})
-		t.n.ForEachSubscription(func(id string, s []node.SubscriptionEnvelope) {
+		t.n.ForEachSubscription(func(id string, isD bool, s []node.SubscriptionEnvelope) {
 			for _, x := range s {
 				ss = append(ss, x.Subscription)
 			}
+			Expect(t, isD).To(Equal(false))
 		})
 		Expect(t, ss).To(HaveLen(2))
 		Expect(t, t.n.SubscriptionLen()).To(Equal(2))
 	})
 
+	o.Spec("returns is deterministic if a single route has deterministic name", func(t TN) {
+		t.n.AddSubscription(func(interface{}) {}, "a", "")
+		t.n.AddSubscription(func(interface{}) {}, "a", "some-name")
+
+		t.n.ForEachSubscription(func(id string, isD bool, s []node.SubscriptionEnvelope) {
+			Expect(t, isD).To(Equal(true))
+		})
+	})
+
+	o.Spec("returns is not deterministic if all deterministic names have been deleted", func(t TN) {
+		t.n.AddSubscription(func(interface{}) {}, "a", "")
+		id := t.n.AddSubscription(func(interface{}) {}, "a", "some-name")
+		t.n.DeleteSubscription(id)
+
+		t.n.ForEachSubscription(func(id string, isD bool, s []node.SubscriptionEnvelope) {
+			Expect(t, isD).To(Equal(false))
+		})
+	})
+
 	o.Spec("it handles ID collisions", func(t TN) {
 		n := node.New(func(int64) int64 { return 0 })
-		id1 := n.AddSubscription(func(interface{}) {}, "")
-		id2 := n.AddSubscription(func(interface{}) {}, "")
+		id1 := n.AddSubscription(func(interface{}) {}, "", "")
+		id2 := n.AddSubscription(func(interface{}) {}, "", "")
 
 		Expect(t, id1).To(Not(Equal(id2)))
 	})

pubsub.go

@@ -24,25 +24,32 @@ import (
 // of PubSub's methods safe to access concurrently. PubSub should be
 // constructed with New().
 type PubSub struct {
-	mu   rlocker
-	n    *node.Node
-	rand func(n int64) int64
+	mu                         rlocker
+	n                          *node.Node
+	rand                       func(n int64) int64
+	deterministicRoutingHasher func(interface{}) uint64
 }
 
 // New constructs a new PubSub.
 func New(opts ...PubSubOption) *PubSub {
-	p := &PubSub{
+	s := &PubSub{
 		mu:   &sync.RWMutex{},
 		rand: rand.Int63n,
 	}
 
 	for _, o := range opts {
-		o.configure(p)
+		o.configure(s)
 	}
 
-	p.n = node.New(p.rand)
+	if s.deterministicRoutingHasher == nil {
+		s.deterministicRoutingHasher = func(_ interface{}) uint64 {
+			return uint64(s.rand(0x7FFFFFFFFFFFFFFF))
+		}
+	}
 
-	return p
+	s.n = node.New(s.rand)
+
+	return s
 }
 
 // PubSubOption is used to configure a PubSub.
@@ -52,25 +59,34 @@ type PubSubOption interface {
 
 type pubsubConfigFunc func(*PubSub)
 
-func (f pubsubConfigFunc) configure(p *PubSub) {
-	f(p)
+func (f pubsubConfigFunc) configure(s *PubSub) {
+	f(s)
 }
 
 // WithNoMutex configures a PubSub that does not have any internal mutexes.
 // This is useful if more complex or custom locking is required. For example,
 // if a subscription needs to subscribe while being published to.
 func WithNoMutex() PubSubOption {
-	return pubsubConfigFunc(func(p *PubSub) {
-		p.mu = nopLock{}
+	return pubsubConfigFunc(func(s *PubSub) {
+		s.mu = nopLock{}
 	})
 }
 
 // WithRand configures a PubSub that will use the given function to make
 // sharding decisions. The given function has to match the symantics of
 // math/rand.Int63n.
 func WithRand(int63 func(max int64) int64) PubSubOption {
-	return pubsubConfigFunc(func(p *PubSub) {
-		p.rand = int63
+	return pubsubConfigFunc(func(s *PubSub) {
+		s.rand = int63
+	})
+}
+
+// WithDeterministicHashing configures a PubSub that will use the given
+// function to hash each published data point. The hash is used only for a
+// subscription that has set its deterministic routing name.
+func WithDeterministicHashing(hashFunction func(interface{}) uint64) PubSubOption {
+	return pubsubConfigFunc(func(s *PubSub) {
+		s.deterministicRoutingHasher = hashFunction
 	})
 }
 
@@ -106,9 +122,19 @@ func WithPath(path []uint64) SubscribeOption {
 	})
 }
 
+// WithDeterministicRouting configures a subscription to have a deterministic
+// routing name. A PubSub configured to use deterministic hashing will use
+// this name and the subscription's shard ID to maintain consistent routing.
+func WithDeterministicRouting(name string) SubscribeOption {
+	return subscribeConfigFunc(func(c *subscribeConfig) {
+		c.deterministicRoutingName = name
+	})
+}
+
 type subscribeConfig struct {
-	shardID string
-	path    []uint64
+	shardID                  string
+	deterministicRoutingName string
+	path                     []uint64
 }
 
 type subscribeConfigFunc func(*subscribeConfig)
@@ -133,7 +159,7 @@ func (s *PubSub) Subscribe(sub Subscription, opts ...SubscribeOption) Unsubscrib
 	for _, p := range c.path {
 		n = n.AddChild(p)
 	}
-	id := n.AddSubscription(sub, c.shardID)
+	id := n.AddSubscription(sub, c.shardID, c.deterministicRoutingName)
 
 	return func() {
 		s.mu.Lock()
@@ -261,15 +287,15 @@ func (s *PubSub) traversePublish(d, next interface{}, a TreeTraverser, n *node.N
 	if n == nil {
 		return
 	}
-	n.ForEachSubscription(func(shardID string, ss []node.SubscriptionEnvelope) {
+	n.ForEachSubscription(func(shardID string, isDeterministic bool, ss []node.SubscriptionEnvelope) {
 		if shardID == "" {
 			for _, x := range ss {
 				x.Subscription(d)
 			}
 			return
 		}
 
-		idx := s.rand(int64(len(ss)))
+		idx := s.determineIdx(d, len(ss), isDeterministic)
 		ss[idx].Subscription(d)
 	})
 
@@ -291,6 +317,13 @@ func (s *PubSub) traversePublish(d, next interface{}, a TreeTraverser, n *node.N
 	}
 }
 
+func (s *PubSub) determineIdx(d interface{}, l int, isDeterministic bool) int64 {
+	if isDeterministic {
+		return int64(s.deterministicRoutingHasher(d) % uint64(l))
+	}
+	return s.rand(int64(l))
+}
+
 // rlocker is used to hold either a real sync.RWMutex or a nop lock.
 // This is used to turn off locking.
 type rlocker interface {

pubsub_test.go

@@ -172,6 +172,79 @@ func TestPubSubWithShardID(t *testing.T) {
 	})
 }
 
+func TestPubSubWithShardingScheme(t *testing.T) {
+	t.Parallel()
+	o := onpar.New()
+	defer o.Run(t)
+	o.BeforeEach(func(t *testing.T) TPS {
+		s, f := newSpySubscrption()
+
+		return TPS{
+			T:            t,
+			subscription: s,
+			sub:          f,
+			p: pubsub.New(pubsub.WithDeterministicHashing(func(data interface{}) uint64 {
+				return uint64(data.(*testStruct).a)
+			})),
+		}
+	})
+
+	o.Spec("it splits data between same shardIDs", func(t TPS) {
+		sub1, f1 := newSpySubscrption()
+		sub2, f2 := newSpySubscrption()
+		sub3, f3 := newSpySubscrption()
+		sub4, f4 := newSpySubscrption()
+		sub5, f5 := newSpySubscrption()
+
+		t.p.Subscribe(f1,
+			pubsub.WithShardID("1"),
+			pubsub.WithDeterministicRouting("black"),
+			pubsub.WithPath(testStructTravCreatePath(&testStructFilter{})),
+		)
+
+		t.p.Subscribe(f2,
+			pubsub.WithShardID("1"),
+			pubsub.WithDeterministicRouting("blue"),
+			pubsub.WithPath(testStructTravCreatePath(&testStructFilter{})),
+		)
+		t.p.Subscribe(f3,
+			pubsub.WithShardID("2"),
+			pubsub.WithPath(testStructTravCreatePath(&testStructFilter{})),
+		)
+
+		t.p.Subscribe(f4,
+			pubsub.WithPath(testStructTravCreatePath(&testStructFilter{})),
+		)
+
+		t.p.Subscribe(f5,
+			pubsub.WithPath(testStructTravCreatePath(&testStructFilter{})),
+		)
+
+		for i := 0; i < 100; i++ {
+			t.p.Publish(&testStruct{a: 1, b: 2}, testStructTravTraverse)
+			t.p.Publish(&testStruct{a: 2, b: 3}, testStructTravTraverse)
+		}
+
+		Expect(t, len(sub1.data)).To(Equal(100))
+		Expect(t, len(sub2.data)).To(Equal(100))
+
+		Expect(t, len(sub3.data)).To(Equal(200))
+		Expect(t, len(sub4.data)).To(Equal(200))
+		Expect(t, len(sub5.data)).To(Equal(200))
+
+		Expect(t, sub1.data[0]).To(Or(
+			Equal(&testStruct{a: 1, b: 2}),
+			Equal(&testStruct{a: 2, b: 3}),
+		))
+
+		Expect(t, sub2.data[0]).To(Or(
+			Equal(&testStruct{a: 1, b: 2}),
+			Equal(&testStruct{a: 2, b: 3}),
+		))
+		Expect(t, sub1.data[0]).To(Not(Equal(sub2.data[0])))
+	})
+}
+
 type spySubscription struct {
 	mu   sync.Mutex
 	data []interface{}