go-cart-actor/pkg/actor/simple_grain_pool.go

package actor

import (
	"context"
	"fmt"
	"log"
	"maps"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"
)

type SimpleGrainPool[V any] struct {
	// fields and methods
	localMu          sync.RWMutex
	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[V], error)
	listeners        []LogListener
	storage          LogStorage[V]
	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

	// metrics is the prometheus instrumentation wired in by NewMetrics
	// at composition time. Nil is fine — every touch site nil-checks
	// first, so the existing tests (which construct a pool without
	// metrics) keep working.
	metrics *Metrics

	// Cluster coordination --------------------------------------------------
	hostname     string
	remoteMu     sync.RWMutex
	remoteOwners map[uint64]Host[V]
	remoteHosts  map[string]Host[V]
	//discardedHostHandler *DiscardedHostHandler

	// House-keeping ---------------------------------------------------------
	purgeTicker *time.Ticker
}

// noLocalStamp is the sentinel value used in the parallel []int64
// lastChanges slice to mark an id for which the broadcaster has no
// local grain (e.g. TakeOwnership called before anyone has loaded the
// grain, or a custom Host implementation that doesn't carry a stamp).
// Receivers treat this sentinel as "no arbitration possible" and fall
// back to the pre-arbitration "always accept remote" verdict so the
// behaviour is identical to running an older binary that didn't carry
// stamps at all.
const noLocalStamp int64 = -1

type GrainPoolConfig[V any] struct {
	Hostname         string
	Spawn            func(ctx context.Context, id uint64) (Grain[V], error)
	SpawnHost        func(host string) (Host[V], error)
	// Destroy is an optional cleanup hook called when a grain is evicted on TTL
	// (e.g. to detach inventory subscriptions wired into its mutations). Leave
	// nil when there is nothing to clean up — purge() skips a nil Destroy.
	Destroy          func(grain Grain[V]) error
	TTL              time.Duration
	PoolSize         int
	MutationRegistry MutationRegistry
	Storage          LogStorage[V]
	// Metrics is the prometheus instrumentation to register pool-level
	// counters, gauges, and histograms with. Nil-safe: when unset, all
	// touch sites are no-ops. Use actor.NewMetrics("cart") (or the
	// matching service prefix) to construct the per-service metrics
	// struct, and pair it with DiskStorage.SetMetrics so the same
	// Metrics instance covers both pool and event-log metrics.
	Metrics *Metrics
}

func NewSimpleGrainPool[V any](config GrainPoolConfig[V]) (*SimpleGrainPool[V], error) {
	p := &SimpleGrainPool[V]{
		grains:           make(map[uint64]Grain[V]),
		mutationRegistry: config.MutationRegistry,
		storage:          config.Storage,
		spawn:            config.Spawn,
		spawnHost:        config.SpawnHost,
		destroy:          config.Destroy,
		ttl:              config.TTL,
		poolSize:         config.PoolSize,
		hostname:         config.Hostname,
		metrics:          config.Metrics,
		remoteOwners:     make(map[uint64]Host[V]),
		remoteHosts:      make(map[string]Host[V]),
		grainLocks:       newKeyedMutex(),
	}

	p.purgeTicker = time.NewTicker(time.Minute)
	go func() {
		for range p.purgeTicker.C {
			p.purge()
		}
	}()

	// Publish the initial pool stats so the dashboard has a number to
	// draw on first scrape, rather than waiting for the first grain
	// spawn / eviction to bump the gauges.
	p.metrics.SetPoolStats(0, p.poolSize)

	return p, nil
}

func (p *SimpleGrainPool[V]) AddListener(listener LogListener) {
	p.listeners = append(p.listeners, listener)
}

func (p *SimpleGrainPool[V]) RemoveListener(listener LogListener) {
	for i, l := range p.listeners {
		if l == listener {
			p.listeners = append(p.listeners[:i], p.listeners[i+1:]...)
			break
		}
	}
}

func (p *SimpleGrainPool[V]) purge() {
	purgeLimit := time.Now().Add(-p.ttl)
	purgedIds := make([]uint64, 0, len(p.grains))
	purgedStamps := make([]int64, 0, len(p.grains))
	p.localMu.Lock()
	evicted := 0
	for id, grain := range p.grains {
		if grain.GetLastAccess().Before(purgeLimit) {
			if p.destroy != nil {
				if err := p.destroy(grain); err != nil {
					log.Printf("failed to destroy grain %d: %v", id, err)
				}
			}
			// Capture the lastChange stamp BEFORE delete so peers
			// receive an honest per-id eviction record on the wire.
			purgedIds = append(purgedIds, id)
			purgedStamps = append(purgedStamps, grain.GetLastChange().UnixNano())
			delete(p.grains, id)
			evicted++
		}
	}
	remaining := len(p.grains)
	p.localMu.Unlock()
	if evicted > 0 {
		p.metrics.SetPoolStats(remaining, p.poolSize)
	}
	p.forAllHosts(func(remote Host[V]) {
		remote.AnnounceExpiry(purgedIds, purgedStamps)
	})

}

// LocalUsage returns the number of resident grains and configured capacity.
func (p *SimpleGrainPool[V]) LocalUsage() (int, int) {
	p.localMu.RLock()
	defer p.localMu.RUnlock()
	return len(p.grains), p.poolSize
}

// LocalCartIDs returns the currently owned cart ids (for control-plane RPCs).
func (p *SimpleGrainPool[V]) GetLocalIds() []uint64 {
	p.localMu.RLock()
	defer p.localMu.RUnlock()
	ids := make([]uint64, 0, len(p.grains))
	for _, g := range p.grains {
		if g == nil {
			continue
		}
		ids = append(ids, uint64(g.GetId()))
	}
	return ids
}

func (p *SimpleGrainPool[V]) HandleRemoteExpiry(host string, ids []uint64, _ []int64) error {
	// lastChanges is informational on this path; expiry is unilateral
	// from the announcer's perspective — we just drop the id from any
	// remoteOwners mapping so future cross-pod forwards stop hitting a
	// dead peer. The parallel stamp slice is kept in the signature for
	// wire symmetry with AnnounceOwnership and so a future hardening
	// pass can use stamps to detect ghost-evictions (e.g. peer evicted
	// with a stamp newer than our last seen Apply result).
	p.remoteMu.Lock()
	defer p.remoteMu.Unlock()
	for _, id := range ids {
		delete(p.remoteOwners, id)
	}
	return nil
}

func (p *SimpleGrainPool[V]) HandleOwnershipChange(host string, ids []uint64, lastChanges []int64) error {
	p.remoteMu.RLock()
	remoteHost, exists := p.remoteHosts[host]
	p.remoteMu.RUnlock()
	if !exists {
		createdHost, err := p.AddRemote(host)
		if err != nil {
			return err
		}
		remoteHost = createdHost
	}
	// Lock order: remoteMu before localMu. Unlock in reverse order so
	// the localMu gauge read below cannot deadlock against a concurrent
	// p.localMu holder.
	p.remoteMu.Lock()
	p.localMu.Lock()
	var accepted, kept, legacy int
	for i, id := range ids {
		// remoteStamp is the announcer's lastChange at broadcast time,
		// or noLocalStamp (-1) if the announcer didn't have a local
		// grain (legacy / TakeOwnership path). The
		// 0..len(lastChanges)-1 slice is indexed parallel to ids; out
		// of range falls back to noLocalStamp (= "no arbitration").
		var remoteStamp int64 = noLocalStamp
		if i < len(lastChanges) {
			remoteStamp = lastChanges[i]
		}
		var localStamp int64 = noLocalStamp
		var hasLocal bool
		if g, ok := p.grains[id]; ok {
			localStamp = g.GetLastChange().UnixNano()
			hasLocal = true
		}

		// First-spawn-wins arbitration. Five distinct cases, ordered:
		//   1) remoteStamp == noLocalStamp  → announcer didn't carry a
		//      real stamp (legacy or TakeOwnership). Keep current
		//      "always accept remote" behaviour so mixed-version
		//      rollouts and unbroadcast-advertised TakeOwnership don't
		//      regress.
		//   2) localStamp  == noLocalStamp  → we have no grain at all.
		//      Accept remote so the next read forwards to the right
		//      owner instead of us spawning a parallel projection.
		//   3) stamps differ → first-spawn-wins. The pod that spawned
		//      the grain earlier (smaller stamp) owns it; its broadcast
		//      is authoritative. The newer pod's projection is
		//      redundant — we drop our local copy and defer to remote
		//      so future OwnerHost lookups route correctly.
		//   4) stamps equal  → deterministic tie-break: lexicographically
		//      smaller hostname wins (both pods compute the same
		//      verdict, so the cold-cache first-touch cannot flip-flop).
		var keepLocal bool
		switch {
		case remoteStamp == noLocalStamp:
			keepLocal = false
			if hasLocal {
				legacy++
			}
		case localStamp == noLocalStamp:
			keepLocal = false
		case localStamp != remoteStamp:
			keepLocal = localStamp < remoteStamp
		default:
			keepLocal = p.hostname < host
		}
		if keepLocal {
			kept++
			continue
		}
		delete(p.grains, id)
		p.remoteOwners[id] = remoteHost
		accepted++
	}
	remaining := len(p.grains)
	p.localMu.Unlock()
	p.remoteMu.Unlock()
	p.metrics.SetPoolStats(remaining, p.poolSize)
	if kept > 0 || accepted > 0 || legacy > 0 {
		log.Printf("HandleOwnershipChange from %s: accepted_remote=%d kept_local=%d legacy_no_stamp=%d",
			host, accepted, kept, legacy)
	}
	return nil
}

// TakeOwnership takes ownership of a grain.
func (p *SimpleGrainPool[V]) TakeOwnership(id uint64) {
	p.broadcastOwnership([]uint64{id})
}

func (p *SimpleGrainPool[V]) AddRemoteHost(host string) {
	p.AddRemote(host)
}

func (p *SimpleGrainPool[V]) AddRemote(host string) (Host[V], error) {
	if host == "" {
		return nil, fmt.Errorf("host is empty")
	}
	if host == p.hostname {
		return nil, fmt.Errorf("same host, this should not happen")
	}
	p.remoteMu.RLock()
	existing, found := p.remoteHosts[host]
	p.remoteMu.RUnlock()
	if found {
		return existing, nil
	}

	remote, err := p.spawnHost(host)
	if err != nil {
		log.Printf("AddRemote %s failed: %v", host, err)

		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
	count := len(p.remoteHosts)
	p.remoteMu.Unlock()
	p.metrics.SetConnectedRemotes(count)

	log.Printf("Connected to remote host %s", host)
	go p.pingLoop(remote)
	go p.initializeRemote(remote)
	go p.SendNegotiation()
	return remote, nil
}

func (p *SimpleGrainPool[V]) initializeRemote(remote Host[V]) {

	remotesIds := remote.GetActorIds()

	p.remoteMu.Lock()
	for _, id := range remotesIds {
		p.localMu.Lock()
		delete(p.grains, id)
		p.localMu.Unlock()
		if _, exists := p.remoteOwners[id]; !exists {
			p.remoteOwners[id] = remote
		}
	}
	p.remoteMu.Unlock()

}

func (p *SimpleGrainPool[V]) RemoveHost(host string) {
	p.remoteMu.Lock()
	remote, exists := p.remoteHosts[host]

	if exists {
		delete(p.remoteHosts, host)
	}
	count := 0
	for id, owner := range p.remoteOwners {
		if owner.Name() == host {
			count++
			delete(p.remoteOwners, id)
		}
	}
	remoteCount := len(p.remoteHosts)
	log.Printf("Removing host %s, grains: %d", host, count)
	p.remoteMu.Unlock()

	// Close once, outside the lock.
	if exists {
		remote.Close()
	}
	p.metrics.SetConnectedRemotes(remoteCount)
}

func (p *SimpleGrainPool[V]) RemoteCount() int {
	p.remoteMu.RLock()
	defer p.remoteMu.RUnlock()
	return len(p.remoteHosts)
}

// RemoteHostNames returns a snapshot of connected remote host identifiers.
func (p *SimpleGrainPool[V]) RemoteHostNames() []string {
	p.remoteMu.RLock()
	defer p.remoteMu.RUnlock()
	hosts := make([]string, 0, len(p.remoteHosts))
	for host := range p.remoteHosts {
		hosts = append(hosts, host)
	}
	return hosts
}

func (p *SimpleGrainPool[V]) IsKnown(host string) bool {
	if host == p.hostname {
		return true
	}
	p.remoteMu.RLock()
	defer p.remoteMu.RUnlock()
	_, ok := p.remoteHosts[host]
	return ok
}

func (p *SimpleGrainPool[V]) pingLoop(remote Host[V]) {
	remote.Ping()
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		if !remote.Ping() {
			if !remote.IsHealthy() {
				log.Printf("Remote %s unhealthy, removing", remote.Name())
				// 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
			}
			continue
		}
	}
}

func (p *SimpleGrainPool[V]) IsHealthy() bool {
	p.remoteMu.RLock()
	defer p.remoteMu.RUnlock()
	for _, r := range p.remoteHosts {
		if !r.IsHealthy() {
			return false
		}
	}
	return true
}

func (p *SimpleGrainPool[V]) Negotiate(otherHosts []string) {

	for _, host := range otherHosts {
		if host != p.hostname {
			p.remoteMu.RLock()
			_, ok := p.remoteHosts[host]
			p.remoteMu.RUnlock()
			if !ok {
				go p.AddRemote(host)
			}
		}
	}
}

func (p *SimpleGrainPool[V]) SendNegotiation() {
	p.metrics.IncNegotiation()

	p.remoteMu.RLock()
	hosts := make([]string, 0, len(p.remoteHosts)+1)
	hosts = append(hosts, p.hostname)
	remotes := make([]Host[V], 0, len(p.remoteHosts))
	for h, r := range p.remoteHosts {
		hosts = append(hosts, h)
		remotes = append(remotes, r)
	}
	p.remoteMu.RUnlock()

	p.forAllHosts(func(remote Host[V]) {
		knownByRemote, err := remote.Negotiate(hosts)

		if err != nil {
			log.Printf("Negotiate with %s failed: %v", remote.Name(), err)
			return
		}
		for _, h := range knownByRemote {
			if !p.IsKnown(h) {
				go p.AddRemote(h)
			}
		}
	})
}

func (p *SimpleGrainPool[V]) forAllHosts(fn func(Host[V])) {
	p.remoteMu.RLock()
	rh := maps.Clone(p.remoteHosts)
	p.remoteMu.RUnlock()

	wg := sync.WaitGroup{}
	for _, host := range rh {
		wg.Go(func() { fn(host) })
	}
	wg.Wait()

	for name, host := range rh {
		if !host.IsHealthy() {
			host.Close()
			p.remoteMu.Lock()
			delete(p.remoteHosts, name)
			p.remoteMu.Unlock()
		}

	}
}

func (p *SimpleGrainPool[V]) broadcastOwnership(ids []uint64) {
	if len(ids) == 0 {
		return
	}

	// Capture each id's grain.lastChange UnixNano stamp under the local
	// read-lock so peers receive an honest oracle for the
	// first-spawn-wins arbitration. For ids that don't have a local
	// grain (e.g. TakeOwnership called before anyone has loaded the
	// grain, or a caller pushing ids into the pool from outside the
	// spawn path), substitute the noLocalStamp sentinel — receivers see
	// -1 and fall back to "always accept remote", preserving the
	// pre-arbitration semantics for that case.
	p.localMu.RLock()
	stamps := make([]int64, 0, len(ids))
	for _, id := range ids {
		if g, ok := p.grains[id]; ok {
			stamps = append(stamps, g.GetLastChange().UnixNano())
		} else {
			stamps = append(stamps, noLocalStamp)
		}
	}
	p.localMu.RUnlock()

	p.forAllHosts(func(rh Host[V]) {
		rh.AnnounceOwnership(p.hostname, ids, stamps)
	})
	log.Printf("%s taking ownership of %d ids", p.hostname, len(ids))
	// go p.statsUpdate()
}

func (p *SimpleGrainPool[V]) getOrClaimGrain(ctx context.Context, id uint64) (Grain[V], error) {
	p.localMu.RLock()
	grain, exists := p.grains[id]
	p.localMu.RUnlock()
	if exists && grain != nil {
		return grain, nil
	}

	grain, err := p.spawn(ctx, id)
	if err != nil {
		return nil, err
	}
	p.metrics.IncSpawn()
	go p.broadcastOwnership([]uint64{id})
	p.localMu.Lock()
	p.grains[id] = grain
	remaining := len(p.grains)
	p.localMu.Unlock()
	p.metrics.SetPoolStats(remaining, p.poolSize)

	return grain, nil
}

// // ErrNotOwner is returned when a cart belongs to another host.
// var ErrNotOwner = fmt.Errorf("not owner")

// Apply applies a mutation to a grain.
func (p *SimpleGrainPool[V]) Apply(ctx context.Context, id uint64, mutation ...proto.Message) (result *MutationResult[V], err error) {
	// The metric closure fires exactly once on every return path: the
	// latency / counters cover the full Apply (lock acquire + spawn +
	// handler + storage enqueue), and the failure counter only bumps
	// when the call returned a non-nil error. `len(mutation)` is the
	// number of proto messages in this Apply call — a single call can
	// batch several messages (SetCartItems sends ClearCart + N AddItem
	// together), and the counter is incremented by that count so the
	// dashboard's `rate(cart_mutations_total[1m])` reports messages
	// per second rather than calls per second, matching the original
	// `grainMutations.Add(len(data.Mutations))` semantics.
	done := p.metrics.MutationStarted(len(mutation))
	defer func() { done(err) }()

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

	mutations, err := p.mutationRegistry.Apply(ctx, grain, mutation...)
	if err != nil {
		return nil, err
	}
	if p.storage != nil {
		go func() {
			if err := p.storage.AppendMutations(id, mutation...); err != nil {
				log.Printf("failed to store mutation for grain %d: %v", id, err)
			}
		}()
	}
	for _, listener := range p.listeners {
		go listener.AppendMutations(id, mutations...)
	}
	state, err := grain.GetCurrentState()
	if err != nil {
		return nil, err
	}

	return &MutationResult[V]{
		Result:    *state,
		Mutations: mutations,
	}, nil
}

// Get returns the current state of a grain.
func (p *SimpleGrainPool[V]) Get(ctx context.Context, id uint64) (result *V, err error) {
	p.metrics.IncLookup()
	unlock := p.grainLocks.lock(id)
	defer unlock()

	grain, err := p.getOrClaimGrain(ctx, id)
	if err != nil {
		return nil, err
	}
	return grain.GetCurrentState()
}

// OwnerHost reports the remote owner (if any) for the supplied cart id.
func (p *SimpleGrainPool[V]) OwnerHost(id uint64) (Host[V], bool) {
	p.remoteMu.RLock()
	defer p.remoteMu.RUnlock()
	owner, ok := p.remoteOwners[id]
	return owner, ok
}

// Hostname returns the local hostname (pod IP).
func (p *SimpleGrainPool[V]) Hostname() string {
	return p.hostname
}

// Close notifies remotes that this host is shutting down.
func (p *SimpleGrainPool[V]) Close() {

	p.forAllHosts(func(rh Host[V]) {
		rh.Close()
	})

	if p.purgeTicker != nil {
		p.purgeTicker.Stop()
	}
}