ollama source for Momentry Core verification

x/mlxrunner/status_memory.go

@@ -0,0 +1,109 @@
+package mlxrunner
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+	"time"
+)
+
+const statusMemoryRefreshWait = 50 * time.Millisecond
+
+type statusMemoryRefreshFunc func() (uint64, error)
+
+// statusMemoryCache keeps health checks from depending synchronously on the
+// serialized MLX worker while still refreshing memory telemetry opportunistically.
+type statusMemoryCache struct {
+	done    <-chan struct{}
+	wait    time.Duration
+	refresh statusMemoryRefreshFunc
+
+	mu          sync.Mutex
+	memory      uint64
+	refreshedAt time.Time
+	inFlight    chan struct{}
+}
+
+func newStatusMemoryCache(ctx context.Context, memory uint64, refreshedAt time.Time, wait time.Duration, refresh statusMemoryRefreshFunc) *statusMemoryCache {
+	return &statusMemoryCache{
+		done:        ctx.Done(),
+		wait:        wait,
+		refresh:     refresh,
+		memory:      memory,
+		refreshedAt: refreshedAt,
+	}
+}
+
+func (c *statusMemoryCache) Memory() uint64 {
+	done := c.startRefresh()
+	if c.wait <= 0 {
+		<-done
+		memory, _ := c.snapshot()
+		return memory
+	}
+
+	timer := time.NewTimer(c.wait)
+	defer timer.Stop()
+
+	select {
+	case <-done:
+	case <-timer.C:
+		memory, refreshedAt := c.snapshot()
+		if refreshedAt.IsZero() {
+			slog.Debug("using cached MLX memory status before first refresh")
+		} else {
+			slog.Debug("using cached MLX memory status", "stale", time.Since(refreshedAt))
+		}
+		return memory
+	case <-c.done:
+	}
+
+	memory, _ := c.snapshot()
+	return memory
+}
+
+func (c *statusMemoryCache) startRefresh() chan struct{} {
+	c.mu.Lock()
+	if c.inFlight != nil {
+		done := c.inFlight
+		c.mu.Unlock()
+		return done
+	}
+
+	refreshDone := make(chan struct{})
+	c.inFlight = refreshDone
+	refresh := c.refresh
+	lifecycleDone := c.done
+	c.mu.Unlock()
+
+	go func() {
+		memory, err := refresh()
+		now := time.Now()
+
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		defer close(refreshDone)
+
+		if err != nil {
+			select {
+			case <-lifecycleDone:
+			default:
+				slog.Debug("failed to refresh MLX memory status", "error", err)
+			}
+			c.inFlight = nil
+			return
+		}
+
+		c.memory = memory
+		c.refreshedAt = now
+		c.inFlight = nil
+	}()
+
+	return refreshDone
+}
+
+func (c *statusMemoryCache) snapshot() (uint64, time.Time) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.memory, c.refreshedAt
+}