ai pubsub

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() {