Add pubsub package

Author: rosylilly

pubsub/pubsub.go

@@ -0,0 +1,88 @@
+package pubsub
+
+import (
+	"context"
+	"sync"
+)
+
+type PubSub struct {
+	Capacity int
+
+	mu sync.RWMutex
+	ch []chan interface{}
+}
+
+func NewPubSub() *PubSub {
+	return &PubSub{
+		Capacity: 10,
+		mu:       sync.RWMutex{},
+		ch:       []chan interface{}{},
+	}
+}
+
+func (p *PubSub) Publish(payload interface{}) {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+
+	for _, ch := range p.ch {
+		ch <- payload
+	}
+}
+
+func (p *PubSub) Subscribe(ctx context.Context, f func(interface{})) <-chan bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	ch := make(chan interface{}, p.Capacity)
+	p.ch = append(p.ch, ch)
+
+	sub := &Subscription{
+		pubsub: p,
+		f:      f,
+		ch:     ch,
+	}
+
+	waiter := make(chan bool)
+
+	go func() {
+	L:
+		for ctx.Err() == nil {
+			select {
+			case payload, ok := <-sub.ch:
+				if ok {
+					sub.f(payload)
+				}
+			case <-ctx.Done():
+				sub.close()
+				break L
+			}
+		}
+		close(waiter)
+	}()
+
+	return waiter
+}
+
+type Subscription struct {
+	pubsub *PubSub
+	f      func(interface{})
+	ch     chan interface{}
+}
+
+func (s *Subscription) close() {
+	s.pubsub.mu.Lock()
+	defer s.pubsub.mu.Unlock()
+
+	for idx, ch := range s.pubsub.ch {
+		if ch == s.ch {
+			deleted := append(s.pubsub.ch[:idx], s.pubsub.ch[idx+1:]...)
+			channels := make([]chan interface{}, len(deleted))
+			copy(channels, deleted)
+			s.pubsub.ch = channels
+		}
+	}
+
+	close(s.ch)
+
+	s.ch = nil
+}

pubsub/pubsub_test.go

@@ -0,0 +1,50 @@
+package pubsub
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestPubSub(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	pubsub := NewPubSub()
+
+	wg := sync.WaitGroup{}
+
+	result1 := int32(0)
+	pubsub.Subscribe(ctx, func(payload interface{}) {
+		atomic.AddInt32(&result1, int32(payload.(int)))
+		wg.Done()
+	})
+	result2 := int32(0)
+	pubsub.Subscribe(ctx, func(payload interface{}) {
+		atomic.AddInt32(&result2, int32(payload.(int)))
+		wg.Done()
+	})
+
+	wg.Add(4)
+	pubsub.Publish(1)
+	pubsub.Publish(2)
+
+	wg.Wait()
+
+	if result1 != 3 {
+		t.Fatalf("invalid 1: %v", result1)
+	}
+	if result2 != 3 {
+		t.Fatalf("invalid 2: %v", result2)
+	}
+}
+
+func TestPubSubUnsubscribe(t *testing.T) {
+	pubsub := NewPubSub()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
+	defer cancel()
+	<-pubsub.Subscribe(ctx, func(payload interface{}) {})
+}