LocalAI/core/services/nodes/probe_cache.go

package nodes

import (
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"
)

// probeCache memoizes recent successful gRPC HealthCheck results for
// (nodeID, addr) tuples so SmartRouter.probeHealth doesn't pay a round-trip
// on every inference request.
//
// Why this exists: with per-request routing (see pkg/model/loader.go), every
// inference call goes through SmartRouter.Route, which probes the backend
// before returning a client. Many gRPC backends (notably llama.cpp's server)
// serialize HealthCheck against active Predict on a shared goroutine, so a
// burst of new requests can stall behind a single long-running stream —
// exactly the "queue stalling" symptom observed in distributed clusters.
//
// The background HealthMonitor (perModelHealthCheck) is still the cluster-wide
// source of truth that reaps actually-dead backends within ~45s; this cache
// only saves the per-request hot path from re-asking when nothing has changed.
//
// TTL matches healthCheckTTL in pkg/model/model.go so the single-process
// IsRecentlyHealthy path and this distributed-mode path share the same
// staleness budget.
type probeCache struct {
	ttl    time.Duration
	mu     sync.Mutex
	seen   map[string]time.Time // key → last successful probe
	flight singleflight.Group   // coalesces concurrent probes for the same key
}

// newProbeCache returns a probeCache with the given TTL. Zero TTL disables
// caching: every call to DoOrCachedResult invokes the probe.
func newProbeCache(ttl time.Duration) *probeCache {
	return &probeCache{
		ttl:  ttl,
		seen: make(map[string]time.Time),
	}
}

// IsFresh reports whether key was successfully probed within TTL.
func (c *probeCache) IsFresh(key string) bool {
	if c.ttl <= 0 {
		return false
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	last, ok := c.seen[key]
	return ok && time.Since(last) < c.ttl
}

// markFresh records key as successfully probed at the current time.
func (c *probeCache) markFresh(key string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.seen[key] = time.Now()
}

// Invalidate drops any cached freshness for key. Used after a probe failure
// (or any other signal that the backend may not be alive) so the next call
// will re-probe instead of trusting stale state.
func (c *probeCache) Invalidate(key string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.seen, key)
}

// InvalidateNode drops every cached probe for nodeID.
//
// Keys are "<nodeID>|<addr>", so this is a prefix scan and not a lookup: a
// departed worker that comes back with recycled ports would otherwise serve one
// request per still-fresh address with no probe at all, which is the one moment
// the cache is asked about a process that certainly is not the one it saw.
func (c *probeCache) InvalidateNode(nodeID string) {
	if nodeID == "" {
		return
	}
	prefix := nodeID + "|"
	c.mu.Lock()
	defer c.mu.Unlock()
	for k := range c.seen {
		if strings.HasPrefix(k, prefix) {
			delete(c.seen, k)
		}
	}
}

// DoOrCachedResult returns true if key is fresh; otherwise it runs probe
// (coalescing concurrent callers via singleflight) and caches a successful
// result. Failed probes invalidate the cache, so a transient miss does not pin
// every subsequent request to a re-probe.
//
// It is the ONLY entry point. A boolean-only sibling, DoOrCached, stood beside
// it until probeHealth stopped using it, after which it was production code
// held green by nothing but its own specs; the shim that reads it as a boolean
// now lives in probe_cache_test.go, where its one caller is.
//
// The second result is the reason the probe never reached the backend, or nil
// when it did.
//
// The second result travels through the SINGLEFLIGHT, which is the whole reason
// it is not simply a variable the caller closes over. A closed-over variable is
// only written by the goroutine that actually runs the probe; every other
// caller coalesced into that flight adopts the leader's boolean and sees its own
// unset variable, so the leader would correctly decline to reap while its
// joiners reaped on the very same observation. Carrying it in singleflight's
// error slot hands every joiner the leader's reason as well as its answer.
//
// A probe that could not reach the backend is NOT cached either way. Caching it
// as fresh would hide a genuinely dead backend behind a network blip, and
// caching it as a failure is what Invalidate already does.
func (c *probeCache) DoOrCachedResult(key string, probe func() (bool, error)) (bool, error) {
	if c.IsFresh(key) {
		return true, nil
	}
	v, unreached, _ := c.flight.Do(key, func() (any, error) {
		// Double-check after potentially waiting: another caller in this
		// flight may have just populated the cache.
		if c.IsFresh(key) {
			return true, nil
		}
		ok, unreached := probe()
		if ok {
			c.markFresh(key)
		} else {
			c.Invalidate(key)
		}
		return ok, unreached
	})
	return v.(bool), unreached
}