mirror of
https://github.com/therootcompany/golib.git
synced 2026-07-10 15:20:21 +00:00
feat(sync/dataset): generic hot-swap refresh with atomic.Pointer
Coordinates a Set of Fetchers (each producing a typed dataset) so a running server can swap to fresh data without locks on the read path. Readers go through atomic.Pointer; Load/Close are mutex-serialized so concurrent reloads can't race. Add and AddInitial register fetchers before the first Load.
This commit is contained in:
parent
584061744b
commit
a408dc890b
295
sync/dataset/dataset.go
Normal file
295
sync/dataset/dataset.go
Normal file
@ -0,0 +1,295 @@
|
||||
// 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 Fetchers, 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 Fetcher
|
||||
// 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); err != nil { ... } // initial populate
|
||||
// go s.Tick(ctx, 47*time.Minute, onError) // background refresh
|
||||
// current := inbound.Value() // lock-free read
|
||||
package dataset
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Fetcher 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 Set.Tick can cancel mid-fetch on shutdown.
|
||||
type Fetcher interface {
|
||||
Fetch(ctx context.Context) (updated bool, err error)
|
||||
}
|
||||
|
||||
// FetcherFunc adapts a plain function to Fetcher.
|
||||
type FetcherFunc func(ctx context.Context) (bool, error)
|
||||
|
||||
func (f FetcherFunc) Fetch(ctx context.Context) (bool, error) { return f(ctx) }
|
||||
|
||||
// NopFetcher always reports no update. Use for sets whose source never
|
||||
// changes (test fixtures, embedded data).
|
||||
type NopFetcher struct{}
|
||||
|
||||
func (NopFetcher) Fetch(ctx context.Context) (bool, error) { return false, nil }
|
||||
|
||||
// PollFiles returns a Fetcher 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.Tick to pick up
|
||||
// changes automatically.
|
||||
func PollFiles(paths ...string) Fetcher {
|
||||
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) Fetch(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 Fetchers 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 Load and Close with a mutex so concurrent callers (e.g. an
|
||||
// /admin/reload handler firing while Tick also fires) don't race the views.
|
||||
// Add and AddInitial mutate s.views without locking — they MUST be called
|
||||
// before the first Load.
|
||||
type Set struct {
|
||||
mu sync.Mutex
|
||||
closed atomic.Bool // set by Close; Load returns early when true
|
||||
fetchers []Fetcher
|
||||
views []reloader
|
||||
loaded atomic.Bool
|
||||
}
|
||||
|
||||
// reloader is a type-erased handle to a View's reload function.
|
||||
type reloader interface {
|
||||
reload(ctx context.Context) error
|
||||
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 ...Fetcher) *Set {
|
||||
return &Set{fetchers: fetchers}
|
||||
}
|
||||
|
||||
// Loaded reports whether Load has completed successfully at least once.
|
||||
func (s *Set) Loaded() bool {
|
||||
return s.loaded.Load()
|
||||
}
|
||||
|
||||
// Load fetches upstream and, on the first call or whenever any fetcher
|
||||
// reports a change, reloads every view and atomically installs the new values.
|
||||
// Returns context.Canceled if the set has been closed.
|
||||
func (s *Set) Load(ctx context.Context) error {
|
||||
if s.closed.Load() {
|
||||
return context.Canceled
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
updated := false
|
||||
for _, f := range s.fetchers {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
u, err := f.Fetch(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if u {
|
||||
updated = true
|
||||
}
|
||||
}
|
||||
if s.loaded.Load() && !updated {
|
||||
return nil
|
||||
}
|
||||
for _, v := range s.views {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := v.reload(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
s.loaded.Store(true)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close closes every view's currently-held value and clears all view
|
||||
// pointers. Call on shutdown to release any OS resources held by the
|
||||
// final snapshots (file handles, network connections). Safe to call on a
|
||||
// Set that hasn't been loaded or whose views hold pure in-memory values —
|
||||
// non-Closer values are skipped. Idempotent.
|
||||
//
|
||||
// It is safe to call Close while Tick is still running; subsequent Load
|
||||
// calls will return context.Canceled.
|
||||
func (s *Set) Close() error {
|
||||
if s.closed.Swap(true) {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
var errs []error
|
||||
for _, v := range s.views {
|
||||
if c, ok := v.(io.Closer); ok {
|
||||
if err := c.Close(); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}
|
||||
v.clear()
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
// Tick calls Load every interval until ctx is done. Load errors are passed to
|
||||
// onError (if non-nil) and do not stop the loop; callers choose whether to log,
|
||||
// count, page, or ignore. Run in a goroutine: `go s.Tick(ctx, d, onError)`.
|
||||
func (s *Set) Tick(ctx context.Context, interval time.Duration, onError func(error)) {
|
||||
t := time.NewTicker(interval)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
if err := s.Load(ctx); err != nil && onError != nil {
|
||||
onError(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// View is a read-only handle to one dataset inside a Set.
|
||||
type View[T any] struct {
|
||||
loader func(ctx context.Context) (*T, error)
|
||||
ptr atomic.Pointer[T]
|
||||
loadedAt atomic.Pointer[time.Time] // nil until first successful reload
|
||||
}
|
||||
|
||||
// Value returns the current snapshot. Nil before the Set is first loaded
|
||||
// unless the view was registered via AddInitial.
|
||||
func (v *View[T]) Value() *T {
|
||||
return v.ptr.Load()
|
||||
}
|
||||
|
||||
// LoadedAt returns the time of the most recent successful reload, or the
|
||||
// zero time if the view has never loaded.
|
||||
func (v *View[T]) LoadedAt() time.Time {
|
||||
if t := v.loadedAt.Load(); t != nil {
|
||||
return *t
|
||||
}
|
||||
return time.Time{}
|
||||
}
|
||||
|
||||
// Close clears the view and closes the currently-held value if it
|
||||
// implements io.Closer. Idempotent — subsequent calls are no-ops.
|
||||
func (v *View[T]) Close() error {
|
||||
prev := v.ptr.Swap(nil)
|
||||
if prev == nil {
|
||||
return nil
|
||||
}
|
||||
if closer, ok := any(prev).(io.Closer); ok {
|
||||
return closer.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (v *View[T]) reload(ctx context.Context) error {
|
||||
t, err := v.loader(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
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)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (v *View[T]) clear() { v.ptr.Swap(nil) }
|
||||
|
||||
// Add registers a new view in s and returns it. Call after NewSet and before
|
||||
// the first Load. View.Value() 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]{loader: loader}
|
||||
s.views = append(s.views, v)
|
||||
return v
|
||||
}
|
||||
|
||||
// AddInitial is like Add but pre-populates the view with initial, so
|
||||
// View.Value() 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]{loader: loader}
|
||||
v.ptr.Store(initial)
|
||||
s.views = append(s.views, v)
|
||||
return v
|
||||
}
|
||||
273
sync/dataset/dataset_test.go
Normal file
273
sync/dataset/dataset_test.go
Normal file
@ -0,0 +1,273 @@
|
||||
package dataset_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/therootcompany/golib/sync/dataset"
|
||||
)
|
||||
|
||||
type countFetcher struct {
|
||||
calls atomic.Int32
|
||||
updated bool
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *countFetcher) Fetch(_ context.Context) (bool, error) {
|
||||
f.calls.Add(1)
|
||||
return f.updated, f.err
|
||||
}
|
||||
|
||||
func TestSet_LoadPopulatesAllViews(t *testing.T) {
|
||||
f := &countFetcher{}
|
||||
g := dataset.NewSet(f)
|
||||
|
||||
var aCalls, bCalls int
|
||||
a := dataset.Add(g, func(_ context.Context) (*string, error) {
|
||||
aCalls++
|
||||
v := "a"
|
||||
return &v, nil
|
||||
})
|
||||
b := dataset.Add(g, func(_ context.Context) (*int, error) {
|
||||
bCalls++
|
||||
v := 42
|
||||
return &v, nil
|
||||
})
|
||||
|
||||
if err := g.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.calls.Load() != 1 {
|
||||
t.Errorf("Fetch called %d times, want 1", f.calls.Load())
|
||||
}
|
||||
if aCalls != 1 || bCalls != 1 {
|
||||
t.Errorf("loaders called (%d,%d), want (1,1)", aCalls, bCalls)
|
||||
}
|
||||
if got := a.Value(); got == nil || *got != "a" {
|
||||
t.Errorf("a.Value() = %v", got)
|
||||
}
|
||||
if got := b.Value(); got == nil || *got != 42 {
|
||||
t.Errorf("b.Value() = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSet_SecondLoadSkipsUnchanged(t *testing.T) {
|
||||
f := &countFetcher{updated: false}
|
||||
g := dataset.NewSet(f)
|
||||
calls := 0
|
||||
dataset.Add(g, func(_ context.Context) (*string, error) {
|
||||
calls++
|
||||
v := "x"
|
||||
return &v, nil
|
||||
})
|
||||
if err := g.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if calls != 1 {
|
||||
t.Fatalf("initial load ran loader %d times, want 1", calls)
|
||||
}
|
||||
if err := g.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if calls != 1 {
|
||||
t.Errorf("second load ran loader %d times, want 1 (no upstream change)", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSet_LoadOnUpdateSwaps(t *testing.T) {
|
||||
f := &countFetcher{updated: true}
|
||||
g := dataset.NewSet(f)
|
||||
n := 0
|
||||
v := dataset.Add(g, func(_ context.Context) (*int, error) {
|
||||
n++
|
||||
return &n, nil
|
||||
})
|
||||
if err := g.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := g.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := v.Value(); got == nil || *got != 2 {
|
||||
t.Errorf("v.Value() = %v, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSet_FetchError(t *testing.T) {
|
||||
f := &countFetcher{err: errors.New("offline")}
|
||||
g := dataset.NewSet(f)
|
||||
dataset.Add(g, func(_ context.Context) (*string, error) {
|
||||
s := "x"
|
||||
return &s, nil
|
||||
})
|
||||
if err := g.Load(t.Context()); err == nil {
|
||||
t.Error("expected fetch error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSet_LoaderError(t *testing.T) {
|
||||
g := dataset.NewSet(dataset.NopFetcher{})
|
||||
dataset.Add(g, func(_ context.Context) (*string, error) {
|
||||
return nil, errors.New("parse fail")
|
||||
})
|
||||
if err := g.Load(t.Context()); err == nil {
|
||||
t.Error("expected loader error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPollFiles(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
a := dir + "/a.txt"
|
||||
b := dir + "/b.txt"
|
||||
if err := os.WriteFile(a, []byte("1"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(b, []byte("2"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
p := dataset.PollFiles(a, b)
|
||||
|
||||
if u, err := p.Fetch(t.Context()); err != nil || !u {
|
||||
t.Fatalf("first Fetch: updated=%v err=%v, want true/nil", u, err)
|
||||
}
|
||||
if u, err := p.Fetch(t.Context()); err != nil || u {
|
||||
t.Fatalf("unchanged Fetch: updated=%v err=%v, want false/nil", u, err)
|
||||
}
|
||||
|
||||
// Bump mtime + change contents on b.
|
||||
future := time.Now().Add(2 * time.Second)
|
||||
if err := os.WriteFile(b, []byte("22"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.Chtimes(b, future, future); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if u, err := p.Fetch(t.Context()); err != nil || !u {
|
||||
t.Errorf("after change: updated=%v err=%v, want true/nil", u, err)
|
||||
}
|
||||
if u, err := p.Fetch(t.Context()); err != nil || u {
|
||||
t.Errorf("steady Fetch: updated=%v err=%v, want false/nil", u, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSet_ClosePreventsTick(t *testing.T) {
|
||||
f := &countFetcher{updated: true}
|
||||
s := dataset.NewSet(f)
|
||||
v := dataset.Add(s, func(_ context.Context) (*int, error) {
|
||||
n := 1
|
||||
return &n, nil
|
||||
})
|
||||
if err := s.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
go s.Tick(ctx, 10*time.Millisecond, nil)
|
||||
time.Sleep(30 * time.Millisecond) // let a few ticks fire
|
||||
|
||||
// Close while Tick is still running.
|
||||
if err := s.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cancel()
|
||||
time.Sleep(30 * time.Millisecond) // let more ticks fire
|
||||
|
||||
// Value should be nil after Close.
|
||||
if v.Value() != nil {
|
||||
t.Error("Value() should be nil after Close")
|
||||
}
|
||||
// After Close, no further fetches should occur. The count should not
|
||||
// have increased between the Close and the final sleep.
|
||||
finalCalls := f.calls.Load()
|
||||
// Give the scheduler a moment to drain any in-flight Load calls.
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
if f.calls.Load() != finalCalls {
|
||||
t.Errorf("fetch count increased after Close (%d → %d)", finalCalls, f.calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSet_CloseIsIdempotent(t *testing.T) {
|
||||
s := dataset.NewSet(dataset.NopFetcher{})
|
||||
dataset.Add(s, func(_ context.Context) (*string, error) {
|
||||
s := "x"
|
||||
return &s, nil
|
||||
})
|
||||
if err := s.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Close(); err != nil {
|
||||
t.Error("second Close should be a no-op")
|
||||
}
|
||||
if err := s.Load(t.Context()); err == nil {
|
||||
t.Error("Load after Close should error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSet_CloseNilsViews(t *testing.T) {
|
||||
s := dataset.NewSet(dataset.NopFetcher{})
|
||||
v := dataset.Add(s, func(_ context.Context) (*int, error) {
|
||||
n := 42
|
||||
return &n, nil
|
||||
})
|
||||
if err := s.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if v.Value() == nil || *v.Value() != 42 {
|
||||
t.Fatal("Value() should be 42 after Load")
|
||||
}
|
||||
if err := s.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if v.Value() != nil {
|
||||
t.Error("Value() should be nil after Close")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSet_AddAfterLoadPanic(t *testing.T) {
|
||||
defer func() {
|
||||
if r := recover(); r == nil {
|
||||
t.Error("Add after Load should panic")
|
||||
}
|
||||
}()
|
||||
s := dataset.NewSet(dataset.NopFetcher{})
|
||||
dataset.Add(s, func(_ context.Context) (*string, error) {
|
||||
s := "x"
|
||||
return &s, nil
|
||||
})
|
||||
if err := s.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
dataset.Add(s, func(_ context.Context) (*string, error) {
|
||||
s := "y"
|
||||
return &s, nil
|
||||
})
|
||||
}
|
||||
|
||||
func TestSet_AddInitialAfterLoadPanic(t *testing.T) {
|
||||
defer func() {
|
||||
if r := recover(); r == nil {
|
||||
t.Error("AddInitial after Load should panic")
|
||||
}
|
||||
}()
|
||||
s := dataset.NewSet(dataset.NopFetcher{})
|
||||
dataset.Add(s, func(_ context.Context) (*string, error) {
|
||||
s := "x"
|
||||
return &s, nil
|
||||
})
|
||||
if err := s.Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var x string
|
||||
dataset.AddInitial(s, &x, func(_ context.Context) (*string, error) {
|
||||
s := "y"
|
||||
return &s, nil
|
||||
})
|
||||
}
|
||||
3
sync/dataset/go.mod
Normal file
3
sync/dataset/go.mod
Normal file
@ -0,0 +1,3 @@
|
||||
module github.com/therootcompany/golib/sync/dataset
|
||||
|
||||
go 1.26.0
|
||||
Loading…
x
Reference in New Issue
Block a user