avoid limits on the number of parallel trace/bucket notifications listeners (#14799)

Simplifies overall limits on the incoming listeners for notifications.

Fixes #14566

internal/pubsub/pubsub.go

@@ -18,6 +18,7 @@
 package pubsub
 
 import (
+	"fmt"
 	"sync"
 	"sync/atomic"
 )
@@ -32,6 +33,7 @@ type Sub struct {
 type PubSub struct {
 	subs           []*Sub
 	numSubscribers int32
+	maxSubscribers int32
 	sync.RWMutex
 }
 
@@ -53,13 +55,18 @@ func (ps *PubSub) Publish(item interface{}) {
 }
 
 // Subscribe - Adds a subscriber to pubsub system
-func (ps *PubSub) Subscribe(subCh chan interface{}, doneCh <-chan struct{}, filter func(entry interface{}) bool) {
+func (ps *PubSub) Subscribe(subCh chan interface{}, doneCh <-chan struct{}, filter func(entry interface{}) bool) error {
+	totalSubs := atomic.AddInt32(&ps.numSubscribers, 1)
+	if ps.maxSubscribers > 0 && totalSubs > ps.maxSubscribers {
+		atomic.AddInt32(&ps.numSubscribers, -1)
+		return fmt.Errorf("the limit of `%d` subscribers is reached", ps.maxSubscribers)
+	}
+
 	ps.Lock()
 	defer ps.Unlock()
 
 	sub := &Sub{subCh, filter}
 	ps.subs = append(ps.subs, sub)
-	atomic.AddInt32(&ps.numSubscribers, 1)
 
 	go func() {
 		<-doneCh
@@ -74,6 +81,8 @@ func (ps *PubSub) Subscribe(subCh chan interface{}, doneCh <-chan struct{}, filt
 		}
 		atomic.AddInt32(&ps.numSubscribers, -1)
 	}()
+
+	return nil
 }
 
 // NumSubscribers returns the number of current subscribers
@@ -81,7 +90,8 @@ func (ps *PubSub) NumSubscribers() int32 {
 	return atomic.LoadInt32(&ps.numSubscribers)
 }
 
-// New inits a PubSub system
-func New() *PubSub {
-	return &PubSub{}
+// New inits a PubSub system with a limit of maximum
+// subscribers unless zero is specified
+func New(maxSubscribers int32) *PubSub {
+	return &PubSub{maxSubscribers: maxSubscribers}
 }

internal/pubsub/pubsub_test.go

@@ -24,68 +24,100 @@ import (
 )
 
 func TestSubscribe(t *testing.T) {
-	ps := New()
+	ps := New(2)
 	ch1 := make(chan interface{}, 1)
 	ch2 := make(chan interface{}, 1)
 	doneCh := make(chan struct{})
 	defer close(doneCh)
-	ps.Subscribe(ch1, doneCh, nil)
-	ps.Subscribe(ch2, doneCh, nil)
+	if err := ps.Subscribe(ch1, doneCh, nil); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if err := ps.Subscribe(ch2, doneCh, nil); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
 	ps.Lock()
 	defer ps.Unlock()
 	if len(ps.subs) != 2 {
-		t.Errorf("expected 2 subscribers")
+		t.Fatalf("expected 2 subscribers")
+	}
+}
+
+func TestSubscribeExceedingLimit(t *testing.T) {
+	ps := New(2)
+	ch1 := make(chan interface{}, 1)
+	ch2 := make(chan interface{}, 1)
+	ch3 := make(chan interface{}, 1)
+	doneCh := make(chan struct{})
+	defer close(doneCh)
+	if err := ps.Subscribe(ch1, doneCh, nil); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if err := ps.Subscribe(ch2, doneCh, nil); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if err := ps.Subscribe(ch3, doneCh, nil); err == nil {
+		t.Fatalf("unexpected nil err")
 	}
 }
 
 func TestUnsubscribe(t *testing.T) {
-	ps := New()
+	ps := New(2)
 	ch1 := make(chan interface{}, 1)
 	ch2 := make(chan interface{}, 1)
 	doneCh1 := make(chan struct{})
 	doneCh2 := make(chan struct{})
-	ps.Subscribe(ch1, doneCh1, nil)
-	ps.Subscribe(ch2, doneCh2, nil)
+	if err := ps.Subscribe(ch1, doneCh1, nil); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if err := ps.Subscribe(ch2, doneCh2, nil); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
 
 	close(doneCh1)
 	// Allow for the above statement to take effect.
 	time.Sleep(100 * time.Millisecond)
 	ps.Lock()
 	if len(ps.subs) != 1 {
-		t.Errorf("expected 1 subscriber")
+		t.Fatal("expected 1 subscriber")
 	}
 	ps.Unlock()
 	close(doneCh2)
 }
 
 func TestPubSub(t *testing.T) {
-	ps := New()
+	ps := New(1)
 	ch1 := make(chan interface{}, 1)
 	doneCh1 := make(chan struct{})
 	defer close(doneCh1)
-	ps.Subscribe(ch1, doneCh1, func(entry interface{}) bool { return true })
+	if err := ps.Subscribe(ch1, doneCh1, func(entry interface{}) bool { return true }); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
 	val := "hello"
 	ps.Publish(val)
 	msg := <-ch1
 	if msg != "hello" {
-		t.Errorf(fmt.Sprintf("expected %s , found %s", val, msg))
+		t.Fatalf(fmt.Sprintf("expected %s , found %s", val, msg))
 	}
 }
 
 func TestMultiPubSub(t *testing.T) {
-	ps := New()
+	ps := New(2)
 	ch1 := make(chan interface{}, 1)
 	ch2 := make(chan interface{}, 1)
 	doneCh := make(chan struct{})
 	defer close(doneCh)
-	ps.Subscribe(ch1, doneCh, func(entry interface{}) bool { return true })
-	ps.Subscribe(ch2, doneCh, func(entry interface{}) bool { return true })
+	if err := ps.Subscribe(ch1, doneCh, func(entry interface{}) bool { return true }); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if err := ps.Subscribe(ch2, doneCh, func(entry interface{}) bool { return true }); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
 	val := "hello"
 	ps.Publish(val)
 
 	msg1 := <-ch1
 	msg2 := <-ch2
 	if msg1 != "hello" && msg2 != "hello" {
-		t.Errorf(fmt.Sprintf("expected both subscribers to have%s , found %s and  %s", val, msg1, msg2))
+		t.Fatalf(fmt.Sprintf("expected both subscribers to have%s , found %s and  %s", val, msg1, msg2))
 	}
 }