amadan

carlos / platform Public

geo: Swappable reader with one-lock snapshots (geoip spec §5.4)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
claude 6b7fe6675b4df754ef514392dd3571bbea083868 parent ecbbf2e
3 files changed, +192 −1
  • internal/geo/mmdb.go +9 −1
  • internal/geo/swappable.go +75 −0
  • internal/geo/swappable_test.go +108 −0
diff --git a/internal/geo/mmdb.go b/internal/geo/mmdb.go
index e085d1eda..406ce6de1 100644
--- a/internal/geo/mmdb.go
+++ b/internal/geo/mmdb.go
@@ -27,6 +27,8 @@ type MMDB struct {
source string
build time.Time
nodes uint64
+
+ closeHook func() // test seam: called by Close before the mapping is released
}
// OpenMMDB opens and verifies the file. An error means "IP step off"; the
@@ -226,7 +228,13 @@ func isSubdivisionCode(s string) bool {
// Close releases the mapping.
func (m *MMDB) Close() error {
- if m == nil || m.r == nil {
+ if m == nil {
+ return nil
+ }
+ if m.closeHook != nil {
+ m.closeHook()
+ }
+ if m.r == nil {
return nil
}
return m.r.Close()
diff --git a/internal/geo/swappable.go b/internal/geo/swappable.go
new file mode 100644
index 000000000..c24948e2c
--- /dev/null
+++ b/internal/geo/swappable.go
@@ -0,0 +1,75 @@
+package geo
+
+import (
+ "net"
+ "sync"
+ "time"
+)
+
+// Snapshot is one answer from one reader: the result, that reader's source,
+// and whether any reader was loaded at all (geoip spec §5.4, codex R2-4).
+// Headers, analytics and attribution all come from one Snapshot, so a swap
+// between the lookup and the response can never relabel the data.
+type Snapshot struct {
+ Result Result
+ Found bool
+ Source string // empty only when Loaded is false
+ Loaded bool
+}
+
+// Swappable is a reader that can be replaced while lookups run. The old
+// mapping is closed only after the write lock has been taken and released,
+// so no lookup is ever in flight on an unmapped file. There is deliberately
+// no Source() method: a source read apart from its lookup is the race this
+// type exists to prevent.
+type Swappable struct {
+ mu sync.RWMutex
+ cur *MMDB
+}
+
+func NewSwappable() *Swappable { return &Swappable{} }
+
+// Lookup answers under one read lock.
+func (s *Swappable) Lookup(ip net.IP) Snapshot {
+ s.mu.RLock()
+ defer s.mu.RUnlock()
+ if s.cur == nil {
+ return Snapshot{}
+ }
+ r, ok := s.cur.Lookup(ip)
+ return Snapshot{Result: r, Found: ok, Source: s.cur.Source(), Loaded: true}
+}
+
+// Locate makes Swappable a Locator, so the console's region guess can hold
+// it without ever seeing the *MMDB underneath (codex R3-2).
+func (s *Swappable) Locate(ip net.IP) (Point, bool) {
+ s.mu.RLock()
+ defer s.mu.RUnlock()
+ if s.cur == nil {
+ return Point{}, false
+ }
+ return s.cur.Locate(ip)
+}
+
+// Loaded reports the current reader's source and build time.
+func (s *Swappable) Loaded() (source string, build time.Time, ok bool) {
+ s.mu.RLock()
+ defer s.mu.RUnlock()
+ if s.cur == nil {
+ return "", time.Time{}, false
+ }
+ return s.cur.Source(), s.cur.BuildEpoch(), true
+}
+
+// Swap installs next (nil = none) and closes the previous reader after
+// every lookup that could see it has finished. Go's RWMutex blocks new
+// readers while a writer waits, so lookups cannot starve a swap.
+func (s *Swappable) Swap(next *MMDB) {
+ s.mu.Lock()
+ old := s.cur
+ s.cur = next
+ s.mu.Unlock()
+ if old != nil && old != next {
+ old.Close()
+ }
+}
diff --git a/internal/geo/swappable_test.go b/internal/geo/swappable_test.go
new file mode 100644
index 000000000..ab1a70da7
--- /dev/null
+++ b/internal/geo/swappable_test.go
@@ -0,0 +1,108 @@
+package geo
+
+import (
+ "net"
+ "sync"
+ "sync/atomic"
+ "testing"
+ "time"
+)
+
+func TestSwappableEmptyAnswersNotLoaded(t *testing.T) {
+ s := NewSwappable()
+ if snap := s.Lookup(net.ParseIP("81.2.69.142")); snap.Loaded || snap.Found || snap.Source != "" {
+ t.Fatalf("empty swappable = %+v", snap)
+ }
+ if _, ok := s.Locate(net.ParseIP("81.2.69.142")); ok {
+ t.Fatal("empty swappable located")
+ }
+ var _ Locator = s
+}
+
+func TestSwappableSnapshotCarriesItsReadersSource(t *testing.T) {
+ s := NewSwappable()
+ m, err := OpenMMDB("testdata/GeoIP2-City-Test.mmdb")
+ if err != nil {
+ t.Fatal(err)
+ }
+ s.Swap(m)
+ snap := s.Lookup(net.ParseIP("81.2.69.142"))
+ if !snap.Loaded || !snap.Found || snap.Source != "GeoIP2-City" || snap.Result.Country != "GB" {
+ t.Fatalf("snapshot = %+v", snap)
+ }
+ miss := s.Lookup(net.ParseIP("10.0.0.1"))
+ if !miss.Loaded || miss.Found || miss.Source != "GeoIP2-City" {
+ t.Fatalf("a miss with a loaded reader must say Loaded and name the source: %+v", miss)
+ }
+ s.Swap(nil)
+ if snap := s.Lookup(net.ParseIP("81.2.69.142")); snap.Loaded {
+ t.Fatalf("after Swap(nil) = %+v", snap)
+ }
+}
+
+// Review Focus 3: no lookup ever runs on a closed reader, and a snapshot's
+// source always belongs to the reader that answered.
+func TestSwappableNeverLooksUpAClosedReader(t *testing.T) {
+ s := NewSwappable()
+ var closedWhileInUse atomic.Bool
+ mk := func(src string) *MMDB {
+ var inUse atomic.Int32
+ var closed atomic.Bool
+ m := &MMDB{source: src}
+ m.lookup = func(_ net.IP, result any) error {
+ inUse.Add(1)
+ defer inUse.Add(-1)
+ if closed.Load() {
+ closedWhileInUse.Store(true)
+ }
+ rec := result.(*mmdbRecord)
+ rec.Country.ISOCode = src[:2]
+ time.Sleep(time.Microsecond)
+ return nil
+ }
+ m.closeHook = func() {
+ if inUse.Load() != 0 {
+ closedWhileInUse.Store(true)
+ }
+ closed.Store(true)
+ }
+ return m
+ }
+ s.Swap(mk("AAsrc"))
+ var wg sync.WaitGroup
+ stop := make(chan struct{})
+ var mismatched atomic.Bool
+ for i := 0; i < 32; i++ {
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ for {
+ select {
+ case <-stop:
+ return
+ default:
+ }
+ snap := s.Lookup(net.ParseIP("81.2.69.142"))
+ if snap.Found && snap.Result.Country != snap.Source[:2] {
+ mismatched.Store(true)
+ }
+ s.Locate(net.ParseIP("81.2.69.142"))
+ }
+ }()
+ }
+ for i := 0; i < 100; i++ {
+ if i%2 == 0 {
+ s.Swap(mk("BBsrc"))
+ } else {
+ s.Swap(mk("AAsrc"))
+ }
+ }
+ close(stop)
+ wg.Wait()
+ if closedWhileInUse.Load() {
+ t.Fatal("a reader was closed while a lookup was using it")
+ }
+ if mismatched.Load() {
+ t.Fatal("a snapshot carried a different reader's source")
+ }
+}