redo pubsub

pkg/actor/grain_pool.go

@@ -26,7 +26,6 @@ type GrainPool[V any] interface {
 	AddRemoteHost(host string)
 	IsHealthy() bool
 	IsKnown(string) bool
-	GetPubSub() *PubSub
 	Close()
 }
 

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)
 	}
 }

pkg/actor/simple_grain_pool.go

@@ -17,12 +17,12 @@ type SimpleGrainPool[V any] struct {
 	grains           map[uint64]Grain[V]
 	mutationRegistry MutationRegistry
 	spawn            func(ctx context.Context, id uint64) (Grain[V], error)
+	destroy          func(grain Grain[V]) error
 	spawnHost        func(host string) (Host, error)
 	listeners        []LogListener
 	storage          LogStorage[V]
 	ttl              time.Duration
 	poolSize         int
-	pubsub           *PubSub
 
 	// Cluster coordination --------------------------------------------------
 	hostname     string
@@ -39,6 +39,7 @@ type GrainPoolConfig[V any] struct {
 	Hostname         string
 	Spawn            func(ctx context.Context, id uint64) (Grain[V], error)
 	SpawnHost        func(host string) (Host, error)
+	Destroy          func(grain Grain[V]) error
 	TTL              time.Duration
 	PoolSize         int
 	MutationRegistry MutationRegistry
@@ -52,6 +53,7 @@ func NewSimpleGrainPool[V any](config GrainPoolConfig[V]) (*SimpleGrainPool[V],
 		storage:          config.Storage,
 		spawn:            config.Spawn,
 		spawnHost:        config.SpawnHost,
+		destroy:          config.Destroy,
 		ttl:              config.TTL,
 		poolSize:         config.PoolSize,
 		hostname:         config.Hostname,
@@ -89,9 +91,10 @@ func (p *SimpleGrainPool[V]) purge() {
 	for id, grain := range p.grains {
 		if grain.GetLastAccess().Before(purgeLimit) {
 			purgedIds = append(purgedIds, id)
-			if p.pubsub != nil {
-				p.pubsub.UnsubscribeAll(id)
+			if err := p.destroy(grain); err != nil {
+				log.Printf("failed to destroy grain %d: %v", id, err)
 			}
+
 			delete(p.grains, id)
 		}
 	}
@@ -417,11 +420,7 @@ func (p *SimpleGrainPool[V]) Apply(ctx context.Context, id uint64, mutation ...p
 	if err != nil {
 		return nil, err
 	}
-	if p.pubsub != nil {
-		if sub, ok := any(grain).(Subscribable); ok {
-			sub.UpdateSubscriptions(p.pubsub)
-		}
-	}
+
 	return &MutationResult[*V]{
 		Result:    result,
 		Mutations: mutations,
@@ -450,21 +449,6 @@ func (p *SimpleGrainPool[V]) Hostname() string {
 	return p.hostname
 }
 
-// GetPubSub returns the pubsub instance.
-func (p *SimpleGrainPool[V]) GetPubSub() *PubSub {
-	return p.pubsub
-}
-
-func (p *SimpleGrainPool[V]) SetPubSub(pubsub *PubSub) {
-	p.pubsub = pubsub
-}
-
-func (p *SimpleGrainPool[V]) Publish(event Event) {
-	if p.pubsub != nil {
-		p.pubsub.Publish(event)
-	}
-}
-
 // Close notifies remotes that this host is shutting down.
 func (p *SimpleGrainPool[V]) Close() {