mark pubsub type safe via generics (#15961)

internal/pubsub/pubsub.go

@@ -24,28 +24,28 @@ import (
 )
 
 // Sub - subscriber entity.
-type Sub struct {
-	ch     chan Maskable
+type Sub[T Maskable] struct {
+	ch     chan T
 	types  Mask
-	filter func(entry Maskable) bool
+	filter func(entry T) bool
 }
 
 // PubSub holds publishers and subscribers
-type PubSub struct {
+type PubSub[T Maskable, M Maskable] struct {
 	// atomics, keep at top:
 	types          uint64
 	numSubscribers int32
 	maxSubscribers int32
 
 	// not atomics:
-	subs []*Sub
+	subs []*Sub[T]
 	sync.RWMutex
 }
 
 // Publish message to the subscribers.
 // Note that publish is always nob-blocking send so that we don't block on slow receivers.
 // Hence receivers should use buffered channel so as not to miss the published events.
-func (ps *PubSub) Publish(item Maskable) {
+func (ps *PubSub[T, M]) Publish(item T) {
 	ps.RLock()
 	defer ps.RUnlock()
 	for _, sub := range ps.subs {
@@ -59,7 +59,7 @@ func (ps *PubSub) Publish(item Maskable) {
 }
 
 // Subscribe - Adds a subscriber to pubsub system
-func (ps *PubSub) Subscribe(mask Mask, subCh chan Maskable, doneCh <-chan struct{}, filter func(entry Maskable) bool) error {
+func (ps *PubSub[T, M]) Subscribe(mask M, subCh chan T, doneCh <-chan struct{}, filter func(entry T) bool) error {
 	totalSubs := atomic.AddInt32(&ps.numSubscribers, 1)
 	if ps.maxSubscribers > 0 && totalSubs > ps.maxSubscribers {
 		atomic.AddInt32(&ps.numSubscribers, -1)
@@ -68,12 +68,12 @@ func (ps *PubSub) Subscribe(mask Mask, subCh chan Maskable, doneCh <-chan struct
 	ps.Lock()
 	defer ps.Unlock()
 
-	sub := &Sub{ch: subCh, types: mask, filter: filter}
+	sub := &Sub[T]{ch: subCh, types: Mask(mask.Mask()), filter: filter}
 	ps.subs = append(ps.subs, sub)
 
 	// We hold a lock, so we are safe to update
 	combined := Mask(atomic.LoadUint64(&ps.types))
-	combined.Merge(mask)
+	combined.Merge(Mask(mask.Mask()))
 	atomic.StoreUint64(&ps.types, uint64(combined))
 
 	go func() {
@@ -97,21 +97,23 @@ func (ps *PubSub) Subscribe(mask Mask, subCh chan Maskable, doneCh <-chan struct
 }
 
 // NumSubscribers returns the number of current subscribers,
-// If t is non-nil, the type is checked against the active subscribed types,
-// and 0 will be returned if nobody is subscribed for the type,
-// otherwise the *total* number of subscribers is returned.
-func (ps *PubSub) NumSubscribers(m Maskable) int32 {
-	if m != nil {
-		types := Mask(atomic.LoadUint64(&ps.types))
-		if !types.Overlaps(Mask(m.Mask())) {
-			return 0
-		}
+// The mask is checked against the active subscribed types,
+// and 0 will be returned if nobody is subscribed for the type(s).
+func (ps *PubSub[T, M]) NumSubscribers(mask M) int32 {
+	types := Mask(atomic.LoadUint64(&ps.types))
+	if !types.Overlaps(Mask(mask.Mask())) {
+		return 0
 	}
 	return atomic.LoadInt32(&ps.numSubscribers)
 }
 
+// Subscribers returns the number of current subscribers for all types.
+func (ps *PubSub[T, M]) Subscribers() int32 {
+	return atomic.LoadInt32(&ps.numSubscribers)
+}
+
 // New inits a PubSub system with a limit of maximum
 // subscribers unless zero is specified
-func New(maxSubscribers int32) *PubSub {
-	return &PubSub{maxSubscribers: maxSubscribers}
+func New[T Maskable, M Maskable](maxSubscribers int32) *PubSub[T, M] {
+	return &PubSub[T, M]{maxSubscribers: maxSubscribers}
 }

internal/pubsub/pubsub_test.go

@@ -24,7 +24,7 @@ import (
 )
 
 func TestSubscribe(t *testing.T) {
-	ps := New(2)
+	ps := New[Maskable, Mask](2)
 	ch1 := make(chan Maskable, 1)
 	ch2 := make(chan Maskable, 1)
 	doneCh := make(chan struct{})
@@ -38,21 +38,21 @@ func TestSubscribe(t *testing.T) {
 	ps.Lock()
 	defer ps.Unlock()
 
-	if len(ps.subs) != 2 || ps.NumSubscribers(nil) != 2 {
+	if len(ps.subs) != 2 || ps.NumSubscribers(MaskAll) != 2 || ps.Subscribers() != 2 {
 		t.Fatalf("expected 2 subscribers")
 	}
 }
 
 func TestNumSubscribersMask(t *testing.T) {
-	ps := New(2)
+	ps := New[Maskable, Mask](2)
 	ch1 := make(chan Maskable, 1)
 	ch2 := make(chan Maskable, 1)
 	doneCh := make(chan struct{})
 	defer close(doneCh)
-	if err := ps.Subscribe(1, ch1, doneCh, nil); err != nil {
+	if err := ps.Subscribe(Mask(1), ch1, doneCh, nil); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
-	if err := ps.Subscribe(2, ch2, doneCh, nil); err != nil {
+	if err := ps.Subscribe(Mask(2), ch2, doneCh, nil); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 	ps.Lock()
@@ -70,7 +70,7 @@ func TestNumSubscribersMask(t *testing.T) {
 	if want, got := int32(2), ps.NumSubscribers(Mask(1|2)); got != want {
 		t.Fatalf("want %d subscribers, got %d", want, got)
 	}
-	if want, got := int32(2), ps.NumSubscribers(nil); got != want {
+	if want, got := int32(2), ps.NumSubscribers(MaskAll); got != want {
 		t.Fatalf("want %d subscribers, got %d", want, got)
 	}
 	if want, got := int32(0), ps.NumSubscribers(Mask(4)); got != want {
@@ -79,7 +79,7 @@ func TestNumSubscribersMask(t *testing.T) {
 }
 
 func TestSubscribeExceedingLimit(t *testing.T) {
-	ps := New(2)
+	ps := New[Maskable, Maskable](2)
 	ch1 := make(chan Maskable, 1)
 	ch2 := make(chan Maskable, 1)
 	ch3 := make(chan Maskable, 1)
@@ -97,7 +97,7 @@ func TestSubscribeExceedingLimit(t *testing.T) {
 }
 
 func TestUnsubscribe(t *testing.T) {
-	ps := New(2)
+	ps := New[Maskable, Maskable](2)
 	ch1 := make(chan Maskable, 1)
 	ch2 := make(chan Maskable, 1)
 	doneCh1 := make(chan struct{})
@@ -127,7 +127,7 @@ func (m maskString) Mask() uint64 {
 }
 
 func TestPubSub(t *testing.T) {
-	ps := New(1)
+	ps := New[Maskable, Maskable](1)
 	ch1 := make(chan Maskable, 1)
 	doneCh1 := make(chan struct{})
 	defer close(doneCh1)
@@ -143,7 +143,7 @@ func TestPubSub(t *testing.T) {
 }
 
 func TestMultiPubSub(t *testing.T) {
-	ps := New(2)
+	ps := New[Maskable, Maskable](2)
 	ch1 := make(chan Maskable, 1)
 	ch2 := make(chan Maskable, 1)
 	doneCh := make(chan struct{})
@@ -165,7 +165,7 @@ func TestMultiPubSub(t *testing.T) {
 }
 
 func TestMultiPubSubMask(t *testing.T) {
-	ps := New(3)
+	ps := New[Maskable, Maskable](3)
 	ch1 := make(chan Maskable, 1)
 	ch2 := make(chan Maskable, 1)
 	ch3 := make(chan Maskable, 1)