ai pubsub

pkg/actor/grain_pool.go

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

pkg/actor/pubsub.go

@@ -0,0 +1,93 @@
+package actor
+
+import (
+	"sync"
+)
+
+// Event represents an event to be published.
+type Event struct {
+	Topic   string
+	Payload interface{}
+}
+
+// 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
+	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,
+	}
+}
+
+// Subscribe adds a grain ID to the subscribers of a topic.
+func (p *PubSub) Subscribe(topic string, grainID uint64) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.subscribers[topic] = append(p.subscribers[topic], grainID)
+}
+
+// Unsubscribe removes a grain ID from the subscribers of a topic.
+func (p *PubSub) Unsubscribe(topic string, grainID uint64) {
+	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:]...)
+			break
+		}
+	}
+	// 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 {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	list := p.subscribers[topic]
+	return append([]uint64(nil), list...)
+}
+
+// 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)
+	}
+}

pkg/actor/simple_grain_pool.go

@@ -22,6 +22,7 @@ type SimpleGrainPool[V any] struct {
 	storage          LogStorage[V]
 	ttl              time.Duration
 	poolSize         int
+	pubsub           *PubSub
 
 	// Cluster coordination --------------------------------------------------
 	hostname     string
@@ -88,7 +89,9 @@ 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)
+			}
 			delete(p.grains, id)
 		}
 	}
@@ -414,6 +417,11 @@ 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,
@@ -442,6 +450,21 @@ 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() {