// Package dataset manages values that are periodically re-fetched from an // upstream source and hot-swapped behind atomic pointers. Consumers read via // View.Value (lock-free); a single Load drives any number of views off a // shared set of Upstreams, so upstreams (one git pull, one tar.gz download) // don't get re-fetched per view. // // Typical lifecycle: // // s := dataset.NewSet(repo) // *gitshallow.Repo satisfies Upstream // inbound := dataset.Add(s, func(ctx context.Context) (*ipcohort.Cohort, error) { ... }) // outbound := dataset.Add(s, func(ctx context.Context) (*ipcohort.Cohort, error) { ... }) // if err := s.Load(ctx, true); err != nil { ... } // initial populate // inbound.Start(ctx, 47*time.Minute) // optional refresh // current := inbound.Current() // lock-free read package dataset import ( "context" "io" "os" "sync" "sync/atomic" "time" "github.com/therootcompany/golib/sync/cacheable" ) // Upstream reports whether an upstream source has changed since the last call. // Implementations should dedup rapid-fire calls internally (e.g. gitshallow // skips redundant pulls within a short window; httpcache uses ETag) and must // honor ctx so asynchronous updates can cancel on shutdown. type Upstream interface { Update(ctx context.Context) (updated bool, err error) } // UpstreamFunc adapts a plain function to Upstream. type UpstreamFunc func(ctx context.Context) (bool, error) func (f UpstreamFunc) Update(ctx context.Context) (bool, error) { return f(ctx) } // NopUpstream always reports no update. Use for sets whose source never // changes (test fixtures, embedded data). type NopUpstream struct{} func (NopUpstream) Update(ctx context.Context) (bool, error) { return false, nil } // PollFiles returns a Upstream that stat's the given paths and reports // "updated" whenever any file's size or modtime has changed since the last // call. The first call always reports updated=true. // // Use for Sets whose source is local files that may be edited out of band // (e.g. a user-provided --inbound list) — pair with Set.Start to pick up // changes automatically. func PollFiles(paths ...string) Upstream { return &filePoller{paths: paths, stats: make(map[string]fileStat, len(paths))} } type fileStat struct { size int64 modTime time.Time } type filePoller struct { mu sync.Mutex paths []string stats map[string]fileStat } func (p *filePoller) Update(ctx context.Context) (bool, error) { p.mu.Lock() defer p.mu.Unlock() changed := false for _, path := range p.paths { if err := ctx.Err(); err != nil { return false, err } info, err := os.Stat(path) if err != nil { return false, err } cur := fileStat{size: info.Size(), modTime: info.ModTime()} if prev, ok := p.stats[path]; !ok || prev != cur { changed = true p.stats[path] = cur } } return changed, nil } // Set ties one or more Upstreams to one or more views. A Load call fetches // each source and, on the first call or when any source reports a change, // reloads every view and atomically swaps its current value. Use multiple // fetchers when a single logical dataset is spread across several archives // (e.g. GeoLite2 City + ASN); a single fetcher is the common case (one git // repo, one tar.gz). // Set serializes snapshot publication so concurrent callers do not race views. // Add and AddInitial mutate s.views without locking — they MUST be called // before the first Load. type Set struct { publishMu sync.Mutex fetchers []Upstream views []reloader loaded atomic.Bool loadedAt atomic.Pointer[time.Time] start sync.Once cancel context.CancelFunc refresh cacheable.Refresh generation atomic.Uint64 RefreshTimeout time.Duration // RefreshInterval gates how often Revalidate will start a refresh. Zero // means always due (freshness is delegated to the upstreams' own dedup // windows). Set to a positive duration to match the other cacheable // implementations' time-based Due. RefreshInterval time.Duration } // reloader is a type-erased handle to a View's reload function. type reloader interface { prepare(ctx context.Context) (any, error) publish(value any) clear() } // NewSet creates a Set backed by fetchers. All fetchers are called on every // Load; the set reloads its views whenever any one of them reports a change. func NewSet(fetchers ...Upstream) *Set { return &Set{fetchers: fetchers} } // Loaded reports whether Load has completed successfully at least once. func (s *Set) Loaded() bool { return s.loaded.Load() } // Due reports whether the set should be revalidated. With RefreshInterval > 0 // it returns false until StaleAt (loadedAt + RefreshInterval) has passed, so a // hot path calling Load(false) does not start a refresh on every request. // With RefreshInterval == 0 it always returns true: the upstreams own their // own dedup windows (httpcache ETag, gitshallow pull throttling) and a refresh // is a cheap "nothing changed" check. func (s *Set) Due(ctx context.Context) (bool, error) { if err := ctx.Err(); err != nil { return false, err } if s.RefreshInterval <= 0 { return true, nil } if !s.loaded.Load() { return true, nil } loadedAt := s.loadedAt.Load() if loadedAt == nil { return true, nil } return !time.Now().Before(loadedAt.Add(s.RefreshInterval)), nil } // Revalidate checks whether work is due and starts one asynchronous update. // Concurrent calls coalesce onto the same update. The update runs with a // detached context bounded by RefreshTimeout (default 5 minutes) so a // cancelled caller context does not kill background work. func (s *Set) Revalidate(ctx context.Context) (started bool, err error) { due, err := s.Due(ctx) if err != nil || !due { return false, err } if !s.refresh.Begin() { return false, nil } go func() { timeout := s.RefreshTimeout if timeout <= 0 { timeout = 5 * time.Minute } refreshCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) defer cancel() s.refresh.Finish(s.update(refreshCtx)) }() return true, nil } func (s *Set) update(ctx context.Context) error { generation := s.generation.Load() updated := false for _, f := range s.fetchers { if err := ctx.Err(); err != nil { return err } u, err := f.Update(ctx) if err != nil { return err } if u { updated = true } } if s.loaded.Load() && !updated { return nil } prepared := make([]any, len(s.views)) for i, v := range s.views { if err := ctx.Err(); err != nil { return err } value, err := v.prepare(ctx) if err != nil { return err } prepared[i] = value } s.publishMu.Lock() defer s.publishMu.Unlock() if s.generation.Load() != generation { return nil } for i, v := range s.views { v.publish(prepared[i]) } now := time.Now() s.loadedAt.Store(&now) s.loaded.Store(true) return nil } // Load performs the initial load, or revalidates an already-loaded set. It // keeps the last good snapshot when a later revalidation fails. func (s *Set) Load(ctx context.Context, wait bool) error { _, err := s.Revalidate(ctx) if err != nil { return err } if !s.loaded.Load() || wait { return s.refresh.Wait(ctx) } return nil } // Start starts at most one optional background revalidation ticker. func (s *Set) Start(ctx context.Context, interval time.Duration) { s.start.Do(func() { if interval <= 0 { interval = time.Hour } tickCtx, cancel := context.WithCancel(ctx) s.cancel = cancel go func() { ticker := time.NewTicker(interval) defer ticker.Stop() for { select { case <-tickCtx.Done(): return case <-ticker.C: _, _ = s.Revalidate(tickCtx) } } }() }) } // Stop stops the optional background ticker and waits for any in-flight // refresh to finish or reach RefreshTimeout. It does not clear snapshots. func (s *Set) Stop() error { if s.cancel != nil { s.cancel() } timeout := s.RefreshTimeout if timeout <= 0 { timeout = 5 * time.Minute } ctx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() return s.refresh.Stop(ctx) } // Clear removes every published view and closes values that own resources. // The set remains usable and may be loaded again. func (s *Set) Clear() error { s.generation.Add(1) s.loaded.Store(false) s.loadedAt.Store(nil) s.publishMu.Lock() defer s.publishMu.Unlock() for _, v := range s.views { v.clear() } return nil } // View is a read-only handle to one dataset inside a Set. type View[T any] struct { set *Set loader func(ctx context.Context) (*T, error) ptr atomic.Pointer[T] loadedAt atomic.Pointer[time.Time] // nil until first successful reload } var ( _ cacheable.Cacheable[any] = (*View[any])(nil) _ cacheable.Mutable[any] = (*View[any])(nil) _ cacheable.Inspectable = (*View[any])(nil) _ cacheable.Tickable = (*View[any])(nil) ) // Current returns the current snapshot without blocking. It is nil before // the Set is first loaded unless the view was registered via AddInitial. func (v *View[T]) Current() *T { return v.ptr.Load() } // Due reports whether the owning set should be revalidated. It delegates to // Set.Due, which honors RefreshInterval when set and otherwise treats the set // as always due (upstreams own their own dedup windows). func (v *View[T]) Due(ctx context.Context) (bool, error) { return v.set.Due(ctx) } // Revalidate asks the owning set to check its upstreams and returns whether // this caller started the refresh. func (v *View[T]) Revalidate(ctx context.Context) (bool, error) { return v.set.Revalidate(ctx) } // Load revalidates the owning set and returns this view's newest snapshot. // On refresh failure, the last good snapshot is returned with the error. func (v *View[T]) Load(ctx context.Context, wait bool) (*T, error) { if err := v.set.Load(ctx, wait); err != nil { return v.Current(), err } return v.Current(), nil } func (v *View[T]) Set(value *T) error { v.set.publishMu.Lock() defer v.set.publishMu.Unlock() v.set.generation.Add(1) previous := v.ptr.Swap(value) if previous != nil { if closer, ok := any(previous).(io.Closer); ok { _ = closer.Close() } } if value != nil { now := time.Now() v.loadedAt.Store(&now) v.set.loadedAt.Store(&now) } return nil } func (v *View[T]) Clear() error { return v.Set(nil) } // Status reports the view's published snapshot. StaleAt is populated from the // set's RefreshInterval (when set); with a zero interval StaleAt stays zero // because freshness is delegated to the upstreams. func (v *View[T]) Status() cacheable.Status { status := cacheable.Status{HasValue: v.Current() != nil, Refreshing: v.set.refresh.Running()} if loadedAt := v.loadedAt.Load(); loadedAt != nil { status.LoadedAt = *loadedAt if v.set.RefreshInterval > 0 { status.StaleAt = status.LoadedAt.Add(v.set.RefreshInterval) } } return status } // Start starts the owning set's optional background ticker. func (v *View[T]) Start(ctx context.Context, interval time.Duration) { v.set.Start(ctx, interval) } // Stop stops the owning set's optional background ticker and waits for any // in-flight refresh to finish or reach RefreshTimeout. func (v *View[T]) Stop() error { return v.set.Stop() } func (v *View[T]) prepare(ctx context.Context) (any, error) { return v.loader(ctx) } func (v *View[T]) publish(value any) { t := value.(*T) prev := v.ptr.Swap(t) // Close the replaced value if it holds OS resources (open file handles, // network connections). Geoip readers and similar wrappers implement // io.Closer; cohort and other pure-in-memory values don't — the type // assertion filters to only the ones that need it. The `prev != nil` // guard is required because any(typedNilPtr) is a non-nil interface // wrapping a nil pointer — the type assertion succeeds and we'd // call Close on a nil receiver. if prev != nil { if closer, ok := any(prev).(io.Closer); ok { _ = closer.Close() } } now := time.Now() v.loadedAt.Store(&now) } func (v *View[T]) clear() { prev := v.ptr.Swap(nil) if prev != nil { if closer, ok := any(prev).(io.Closer); ok { _ = closer.Close() } } } // Add registers a new view in s and returns it. Call after NewSet and before // the first Load. View.Current() returns nil until Set.Load succeeds. // The loader receives the ctx passed to Set.Load, so long-running parses // should honor ctx.Err() to support graceful shutdown. // // Panics if called after the first Load. func Add[T any](s *Set, loader func(ctx context.Context) (*T, error)) *View[T] { if s.loaded.Load() { panic("dataset: Add called after Load") } v := &View[T]{set: s, loader: loader} s.views = append(s.views, v) return v } // AddInitial is like Add but pre-populates the view with initial, so // View.Current() returns a usable (possibly empty) value before the first // Load completes. Use when the initial state is benign (e.g. an empty // cohort matches nothing) and you want to start serving before the // first load finishes. // // Panics if called after the first Load. func AddInitial[T any](s *Set, initial *T, loader func(ctx context.Context) (*T, error)) *View[T] { if s.loaded.Load() { panic("dataset: AddInitial called after Load") } v := &View[T]{set: s, loader: loader} v.ptr.Store(initial) s.views = append(s.views, v) return v }