fixes

pkg/actor/keyed_mutex.go

@@ -0,0 +1,50 @@
+package actor
+
+import "sync"
+
+// keyedMutex provides one logical mutex per key (grain id). It is the
+// mechanism that enforces the core actor guarantee: a single grain only ever
+// processes one message (spawn / mutation / read) at a time, while distinct
+// grains run fully in parallel.
+//
+// Locks are reference counted and removed once no caller holds or waits on
+// them, so memory stays proportional to in-flight grains rather than the total
+// number of grains ever touched.
+type keyedMutex struct {
+	mu    sync.Mutex
+	locks map[uint64]*keyedMutexEntry
+}
+
+type keyedMutexEntry struct {
+	mu   sync.Mutex
+	refs int
+}
+
+func newKeyedMutex() *keyedMutex {
+	return &keyedMutex{locks: make(map[uint64]*keyedMutexEntry)}
+}
+
+// lock acquires the mutex for id and returns an unlock function. The returned
+// function must be called exactly once.
+func (k *keyedMutex) lock(id uint64) func() {
+	k.mu.Lock()
+	entry, ok := k.locks[id]
+	if !ok {
+		entry = &keyedMutexEntry{}
+		k.locks[id] = entry
+	}
+	entry.refs++
+	k.mu.Unlock()
+
+	entry.mu.Lock()
+
+	return func() {
+		entry.mu.Unlock()
+		k.mu.Lock()
+		entry.refs--
+		if entry.refs == 0 {
+			delete(k.locks, id)
+		}
+		k.mu.Unlock()
+	}
+}

pkg/actor/simple_grain_pool.go

@@ -24,6 +24,10 @@ type SimpleGrainPool[V any] struct {
 	ttl              time.Duration
 	poolSize         int
 
+	// grainLocks serializes spawn + mutation + read per grain id so that each
+	// grain processes a single message at a time (the actor guarantee).
+	grainLocks *keyedMutex
+
 	// Cluster coordination --------------------------------------------------
 	hostname     string
 	remoteMu     sync.RWMutex
@@ -59,6 +63,7 @@ func NewSimpleGrainPool[V any](config GrainPoolConfig[V]) (*SimpleGrainPool[V],
 		hostname:         config.Hostname,
 		remoteOwners:     make(map[uint64]Host[V]),
 		remoteHosts:      make(map[string]Host[V]),
+		grainLocks:       newKeyedMutex(),
 	}
 
 	p.purgeTicker = time.NewTicker(time.Minute)
@@ -187,7 +192,17 @@ func (p *SimpleGrainPool[V]) AddRemote(host string) (Host[V], error) {
 
 		return nil, err
 	}
+
+	// Re-check under the write lock: spawnHost opened a real connection above
+	// without holding the lock, so a concurrent AddRemote for the same host may
+	// have already registered one. Keep the existing entry and close ours to
+	// avoid leaking the gRPC connection / HTTP transport (and its file handles).
 	p.remoteMu.Lock()
+	if existing, found := p.remoteHosts[host]; found {
+		p.remoteMu.Unlock()
+		go remote.Close()
+		return existing, nil
+	}
 	p.remoteHosts[host] = remote
 	p.remoteMu.Unlock()
 	// connectedRemotes.Set(float64(p.RemoteCount()))
@@ -221,7 +236,6 @@ func (p *SimpleGrainPool[V]) RemoveHost(host string) {
 	remote, exists := p.remoteHosts[host]
 
 	if exists {
-		go remote.Close()
 		delete(p.remoteHosts, host)
 	}
 	count := 0
@@ -234,6 +248,7 @@ func (p *SimpleGrainPool[V]) RemoveHost(host string) {
 	log.Printf("Removing host %s, grains: %d", host, count)
 	p.remoteMu.Unlock()
 
+	// Close once, outside the lock.
 	if exists {
 		remote.Close()
 	}
@@ -275,7 +290,9 @@ func (p *SimpleGrainPool[V]) pingLoop(remote Host[V]) {
 		if !remote.Ping() {
 			if !remote.IsHealthy() {
 				log.Printf("Remote %s unhealthy, removing", remote.Name())
-				p.Close()
+				// Remove only this host. Previously this called p.Close(),
+				// which tore down every remote connection and stopped the
+				// purge ticker for the whole pool.
 				p.RemoveHost(remote.Name())
 				return
 			}
@@ -397,6 +414,12 @@ func (p *SimpleGrainPool[V]) getOrClaimGrain(ctx context.Context, id uint64) (Gr
 
 // Apply applies a mutation to a grain.
 func (p *SimpleGrainPool[V]) Apply(ctx context.Context, id uint64, mutation ...proto.Message) (*MutationResult[V], error) {
+	// Serialize all access to this grain: spawn, mutation handlers and the
+	// final state read happen atomically with respect to other callers of the
+	// same id. Different ids never contend.
+	unlock := p.grainLocks.lock(id)
+	defer unlock()
+
 	grain, err := p.getOrClaimGrain(ctx, id)
 	if err != nil {
 		return nil, err
@@ -429,6 +452,9 @@ func (p *SimpleGrainPool[V]) Apply(ctx context.Context, id uint64, mutation ...p
 
 // Get returns the current state of a grain.
 func (p *SimpleGrainPool[V]) Get(ctx context.Context, id uint64) (*V, error) {
+	unlock := p.grainLocks.lock(id)
+	defer unlock()
+
 	grain, err := p.getOrClaimGrain(ctx, id)
 	if err != nil {
 		return nil, err