redo pubsub

pkg/actor/pubsub.go

@@ -1,93 +1,79 @@
 package actor
 
 import (
+	"iter"
+	"log"
+	"slices"
 	"sync"
 )
 
-// Event represents an event to be published.
-type Event struct {
-	Topic   string
-	Payload interface{}
-}
+type ReceiverFunc[V any] func(event V)
 
-// NotifyFunc is a function to notify a grain of an event.
-type NotifyFunc func(grainID uint64, event Event)
-
-// Subscribable is an interface for grains that can receive notifications.
-type Subscribable interface {
-	Notify(event Event)
-	UpdateSubscriptions(pubsub *PubSub)
-}
-
-// PubSub manages subscriptions for grains to topics.
-// Topics are strings, e.g., "sku:12345"
-// Subscribers are grain IDs (uint64)
-type PubSub struct {
-	subscribers map[string][]uint64
+type PubSub[V any] struct {
+	subscribers []*ReceiverFunc[V]
 	mu          sync.RWMutex
-	notify      NotifyFunc
 }
 
 // NewPubSub creates a new PubSub instance.
-func NewPubSub(notify NotifyFunc) *PubSub {
-	return &PubSub{
-		subscribers: make(map[string][]uint64),
-		notify:      notify,
+func NewPubSub[V any]() *PubSub[V] {
+	return &PubSub[V]{
+		subscribers: make([]*ReceiverFunc[V], 0),
 	}
 }
 
 // Subscribe adds a grain ID to the subscribers of a topic.
-func (p *PubSub) Subscribe(topic string, grainID uint64) {
+func (p *PubSub[V]) Subscribe(receiver ReceiverFunc[V]) {
+	if receiver == nil {
+		return
+	}
 	p.mu.Lock()
 	defer p.mu.Unlock()
-	p.subscribers[topic] = append(p.subscribers[topic], grainID)
+	log.Printf("adding subscriber")
+	p.subscribers = append(p.subscribers, &receiver)
 }
 
 // Unsubscribe removes a grain ID from the subscribers of a topic.
-func (p *PubSub) Unsubscribe(topic string, grainID uint64) {
+func (p *PubSub[V]) Unsubscribe(receiver ReceiverFunc[V]) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
-	list := p.subscribers[topic]
-	for i, id := range list {
-		if id == grainID {
-			p.subscribers[topic] = append(list[:i], list[i+1:]...)
+	list := p.subscribers
+	prt := &receiver
+	for i, sub := range list {
+		if sub == nil {
+			continue
+		}
+		if sub == prt {
+			log.Printf("removing subscriber")
+			p.subscribers = append(list[:i], list[i+1:]...)
 			break
 		}
 	}
+	p.subscribers = slices.DeleteFunc(p.subscribers, func(fn *ReceiverFunc[V]) bool {
+		return fn == nil
+	})
 	// If list is empty, could delete, but not necessary
 }
 
-// UnsubscribeAll removes the grain ID from all topics.
-func (p *PubSub) UnsubscribeAll(grainID uint64) {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	for topic, list := range p.subscribers {
-		newList := make([]uint64, 0, len(list))
-		for _, id := range list {
-			if id != grainID {
-				newList = append(newList, id)
-			}
-		}
-		if len(newList) == 0 {
-			delete(p.subscribers, topic)
-		} else {
-			p.subscribers[topic] = newList
-		}
-	}
-}
-
 // GetSubscribers returns a copy of the subscriber IDs for a topic.
-func (p *PubSub) GetSubscribers(topic string) []uint64 {
+func (p *PubSub[V]) GetSubscribers() iter.Seq[ReceiverFunc[V]] {
 	p.mu.RLock()
 	defer p.mu.RUnlock()
-	list := p.subscribers[topic]
-	return append([]uint64(nil), list...)
+
+	return func(yield func(ReceiverFunc[V]) bool) {
+		for _, sub := range p.subscribers {
+			if sub == nil {
+				continue
+			}
+			if !yield(*sub) {
+				return
+			}
+		}
+	}
 }
 
 // Publish sends an event to all subscribers of the topic.
-func (p *PubSub) Publish(event Event) {
-	subs := p.GetSubscribers(event.Topic)
-	for _, id := range subs {
-		p.notify(id, event)
+func (p *PubSub[V]) Publish(event V) {
+	for notify := range p.GetSubscribers() {
+		notify(event)
 	}
 }