Simplify HTTP trace related code (#7833)

pkg/pubsub/pubsub.go

@@ -20,64 +20,65 @@ import (
 	"sync"
 )
 
+// Sub - subscriber entity.
+type Sub struct {
+	ch     chan interface{}
+	filter func(entry interface{}) bool
+}
+
 // PubSub holds publishers and subscribers
 type PubSub struct {
-	subs  []chan interface{}
-	pub   chan interface{}
-	mutex sync.Mutex
+	subs []*Sub
+	sync.RWMutex
 }
 
-// process item to subscribers.
-func (ps *PubSub) process() {
-	for item := range ps.pub {
-		ps.mutex.Lock()
-		for _, sub := range ps.subs {
-			go func(s chan interface{}) {
-				s <- item
-			}(sub)
-		}
-		ps.mutex.Unlock()
-	}
-}
-
-// Publish message to pubsub system
+// 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 interface{}) {
-	ps.pub <- item
+	ps.RLock()
+	defer ps.RUnlock()
+
+	for _, sub := range ps.subs {
+		if sub.filter(item) {
+			select {
+			case sub.ch <- item:
+			default:
+			}
+		}
+	}
 }
 
 // Subscribe - Adds a subscriber to pubsub system
-func (ps *PubSub) Subscribe() chan interface{} {
-	ps.mutex.Lock()
-	defer ps.mutex.Unlock()
-	ch := make(chan interface{})
-	ps.subs = append(ps.subs, ch)
-	return ch
-}
+func (ps *PubSub) Subscribe(subCh chan interface{}, doneCh chan struct{}, filter func(entry interface{}) bool) {
+	ps.Lock()
+	defer ps.Unlock()
 
-// Unsubscribe removes current subscriber
-func (ps *PubSub) Unsubscribe(ch chan interface{}) {
-	ps.mutex.Lock()
-	defer ps.mutex.Unlock()
+	sub := &Sub{subCh, filter}
+	ps.subs = append(ps.subs, sub)
 
-	for i, sub := range ps.subs {
-		if sub == ch {
-			close(ch)
-			ps.subs = append(ps.subs[:i], ps.subs[i+1:]...)
+	go func() {
+		<-doneCh
+
+		ps.Lock()
+		defer ps.Unlock()
+
+		for i, s := range ps.subs {
+			if s == sub {
+				ps.subs = append(ps.subs[:i], ps.subs[i+1:]...)
+			}
 		}
-	}
+	}()
 }
 
 // HasSubscribers returns true if pubsub system has subscribers
 func (ps *PubSub) HasSubscribers() bool {
-	ps.mutex.Lock()
-	defer ps.mutex.Unlock()
+	ps.RLock()
+	defer ps.RUnlock()
 	return len(ps.subs) > 0
 }
 
 // New inits a PubSub system
 func New() *PubSub {
-	ps := &PubSub{}
-	ps.pub = make(chan interface{})
-	go ps.process()
-	return ps
+	return &PubSub{}
 }

pkg/pubsub/pubsub_test.go

@@ -19,12 +19,19 @@ package pubsub
 import (
 	"fmt"
 	"testing"
+	"time"
 )
 
 func TestSubscribe(t *testing.T) {
 	ps := New()
-	ps.Subscribe()
-	ps.Subscribe()
+	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)
+	ps.Lock()
+	defer ps.Unlock()
 	if len(ps.subs) != 2 {
 		t.Errorf("expected 2 subscribers")
 	}
@@ -32,20 +39,33 @@ func TestSubscribe(t *testing.T) {
 
 func TestUnsubscribe(t *testing.T) {
 	ps := New()
-	c1 := ps.Subscribe()
-	ps.Subscribe()
-	ps.Unsubscribe(c1)
+	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)
+
+	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")
 	}
+	ps.Unlock()
+	close(doneCh2)
 }
 
 func TestPubSub(t *testing.T) {
 	ps := New()
-	c1 := ps.Subscribe()
+	ch1 := make(chan interface{}, 1)
+	doneCh1 := make(chan struct{})
+	defer close(doneCh1)
+	ps.Subscribe(ch1, doneCh1, func(entry interface{}) bool { return true })
 	val := "hello"
 	ps.Publish(val)
-	msg := <-c1
+	msg := <-ch1
 	if msg != "hello" {
 		t.Errorf(fmt.Sprintf("expected %s , found %s", val, msg))
 	}
@@ -53,13 +73,17 @@ func TestPubSub(t *testing.T) {
 
 func TestMultiPubSub(t *testing.T) {
 	ps := New()
-	c1 := ps.Subscribe()
-	c2 := ps.Subscribe()
+	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 })
 	val := "hello"
 	ps.Publish(val)
 
-	msg1 := <-c1
-	msg2 := <-c2
+	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))
 	}