platform/catalog/snapshot_manager.go

// SnapshotHolder provides atomic hot-reload for SnapshotReader and solves the
// classic "Get concurrent with Close/munmap" race that would otherwise produce
// uncatchable SIGBUS / SIGSEGV under reload.
//
// Safety is an RWMutex (readMu): Get holds RLock across the load-and-read of a
// reader, and a reader is only Close()d (munmap) under Lock, so a reader can
// never be unmapped while a Get is reading it — regardless of how long that
// Get's goroutine was descheduled. (A drain-grace timer alone, the original
// approach, is only best-effort: a goroutine paused longer than the grace can
// still read freed pages. TestSnapshotHolder_SwapRace reproduces that crash.)
//
// The drain grace is retained as a close-BATCHING optimization: it defers and
// coalesces reader closes onto the reaper tick instead of taking the exclusive
// Lock on every swap. Pattern source: slask-finder's mmap+reload stores.
// See docs/inventory-shape.md for the prior art.
//
// Use:
//
//	holder := catalog.NewSnapshotHolder(30 * time.Second) // drain grace
//	if err := holder.Open("/var/lib/catalog/bygglagret.bin"); err != nil { ... }
//	// later, when a fresh snapshot lands:
//	if err := holder.Reload("/var/lib/catalog/bygglagret.bin"); err != nil { ... }
//	// pulls the projection; always hits the current reader:
//	p, ok := holder.Get(sku)
//
// holder.Get / Len are safe for parallel use. Open and Reload are also
// safe for parallel use. Close stops the reaper goroutine and shuts the
// current reader + drains down.
package catalog

import (
	"sync"
	"sync/atomic"
	"time"
)

// DefaultDrainGrace is the default drain period for SnapshotHolder. Finder's
// analogous reload path uses ~30s; we mirror that value.
const DefaultDrainGrace = 30 * time.Second

// SnapshotHolder swaps a SnapshotReader under load without unmapping pages
// out from under in-flight Get()s. Set retains the previous reader behind a
// drain timer; after the timer elapses the reaper goroutine closes the old
// reader (and runs its munmap).
type SnapshotHolder[T any] struct {
	current atomic.Pointer[SnapshotReader[T]]

	// readMu makes reclamation race-free: Get holds RLock across the
	// Load-and-read of a reader, and a reader is only Close()d (munmap) under
	// Lock — so a reader can never be unmapped while a Get is reading it, even
	// if that Get's goroutine was descheduled past the drain grace. The grace
	// (below) is now only close-BATCHING — defer + coalesce closes off the swap
	// path — NOT the safety mechanism it was before (which was best-effort and
	// could SIGSEGV under a long enough scheduler pause).
	readMu sync.RWMutex

	drainsMu   sync.Mutex
	drains     []*pendingClose[T]
	drainGrace time.Duration

	stopCh chan struct{}
	once   sync.Once
}

type pendingClose[T any] struct {
	r        *SnapshotReader[T]
	expireAt time.Time
}

// NewSnapshotHolder returns a holder of T with the given drain grace (the time
// the previous reader lives behind Set before being closed). Zero or
// negative grace falls back to DefaultDrainGrace.
func NewSnapshotHolder[T any](grace time.Duration) *SnapshotHolder[T] {
	if grace <= 0 {
		grace = DefaultDrainGrace
	}
	h := &SnapshotHolder[T]{
		drainGrace: grace,
		stopCh:     make(chan struct{}),
	}
	go h.sweepLoop()
	return h
}

// Open is shorthand for Set(ReadSnapshotFile(path)). Errors leave the
// holder's current reader unchanged.
func (h *SnapshotHolder[T]) Open(path string) error {
	r, err := ReadSnapshotFile[T](path)
	if err != nil {
		return err
	}
	h.Set(r)
	return nil
}

// Reload atomically swaps the current reader for one freshly opened from
// path. The old reader (if any) is queued for close after the drain grace
// elapses; in-flight Get(s) on the old reader continue to work.
func (h *SnapshotHolder[T]) Reload(path string) error {
	r, err := ReadSnapshotFile[T](path)
	if err != nil {
		return err
	}
	h.Set(r)
	return nil
}

// Set replaces the current reader with r. The previous reader is queued
// for close after the drain grace elapses. r may be nil to clear.
func (h *SnapshotHolder[T]) Set(r *SnapshotReader[T]) {
	if r == nil {
		h.current.Store(nil)
		return
	}
	old := h.current.Swap(r)
	if old == nil {
		return
	}
	h.drainsMu.Lock()
	h.drains = append(h.drains, &pendingClose[T]{
		r:        old,
		expireAt: time.Now().Add(h.drainGrace),
	})
	h.drainsMu.Unlock()
}

// Get is the concurrent-safe shim over the current reader's Get. Returns
// (zero, false) when the holder is empty. The RLock is held across the Load AND
// the read so a concurrent reaper can't munmap this reader mid-read.
func (h *SnapshotHolder[T]) Get(sku string) (T, bool) {
	h.readMu.RLock()
	defer h.readMu.RUnlock()
	r := h.current.Load()
	if r == nil {
		var zero T
		return zero, false
	}
	return r.Get(sku)
}

// Len returns the count of the current reader (0 if empty).
func (h *SnapshotHolder[T]) Len() int {
	h.readMu.RLock()
	defer h.readMu.RUnlock()
	r := h.current.Load()
	if r == nil {
		return 0
	}
	return r.Len()
}

// Close stops the reaper goroutine and closes the current reader + any
// pending drains without waiting for the grace period. Safe to call
// multiple times.
func (h *SnapshotHolder[T]) Close() error {
	var firstErr error
	h.once.Do(func() {
		close(h.stopCh)
		// Close under Lock so no in-flight Get is reading these readers' mmap
		// when we munmap them.
		h.readMu.Lock()
		r := h.current.Swap(nil)
		if r != nil {
			if err := r.Close(); err != nil {
				firstErr = err
			}
		}
		h.readMu.Unlock()
		h.drainsMu.Lock()
		drained := h.drains
		h.drains = nil
		h.drainsMu.Unlock()
		for _, d := range drained {
			h.readMu.Lock()
			err := d.r.Close()
			h.readMu.Unlock()
			if err != nil && firstErr == nil {
				firstErr = err
			}
		}
	})
	return firstErr
}

// DrainCount returns the number of readers currently in the drain queue.
// Useful for tests and operational metrics.
func (h *SnapshotHolder[T]) DrainCount() int {
	h.drainsMu.Lock()
	defer h.drainsMu.Unlock()
	return len(h.drains)
}

// reapNow scans the drain queue and Closes any readers whose grace has
// elapsed. Called by the sweeper goroutine on a tick.
func (h *SnapshotHolder[T]) reapNow() {
	now := time.Now()
	var alive []*pendingClose[T]
	var reaped []*pendingClose[T]
	h.drainsMu.Lock()
	for _, d := range h.drains {
		if !d.expireAt.After(now) {
			reaped = append(reaped, d)
		} else {
			alive = append(alive, d)
		}
	}
	h.drains = alive
	h.drainsMu.Unlock()
	for _, d := range reaped {
		// Lock briefly so the munmap can't race an in-flight Get on this reader.
		// Gets that loaded this reader hold RLock and finish before the Lock is
		// granted; Gets starting after see the swapped-in reader, never this one.
		h.readMu.Lock()
		_ = d.r.Close()
		h.readMu.Unlock()
	}
}

func (h *SnapshotHolder[T]) sweepLoop() {
	interval := h.drainGrace / 4
	if interval < 10*time.Millisecond {
		interval = 10 * time.Millisecond
	}
	t := time.NewTicker(interval)
	defer t.Stop()
	for {
		select {
		case <-h.stopCh:
			return
		case <-t.C:
			h.reapNow()
		}
	}
}