diff --git a/sync/dataset/dataset.go b/sync/dataset/dataset.go new file mode 100644 index 0000000..6911efa --- /dev/null +++ b/sync/dataset/dataset.go @@ -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 +} diff --git a/sync/dataset/dataset_test.go b/sync/dataset/dataset_test.go new file mode 100644 index 0000000..88fa0d3 --- /dev/null +++ b/sync/dataset/dataset_test.go @@ -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 + }) +} diff --git a/sync/dataset/go.mod b/sync/dataset/go.mod new file mode 100644 index 0000000..b22cfc7 --- /dev/null +++ b/sync/dataset/go.mod @@ -0,0 +1,3 @@ +module github.com/therootcompany/golib/sync/dataset + +go 1.26.0