| 1 | package geo |
| 2 | |
| 3 | import ( |
| 4 | "net" |
| 5 | "sync" |
| 6 | "sync/atomic" |
| 7 | "testing" |
| 8 | "time" |
| 9 | ) |
| 10 | |
| 11 | func TestSwappableEmptyAnswersNotLoaded(t *testing.T) { |
| 12 | s := NewSwappable() |
| 13 | if snap := s.Lookup(net.ParseIP("81.2.69.142")); snap.Loaded || snap.Found || snap.Source != "" { |
| 14 | t.Fatalf("empty swappable = %+v", snap) |
| 15 | } |
| 16 | if _, ok := s.Locate(net.ParseIP("81.2.69.142")); ok { |
| 17 | t.Fatal("empty swappable located") |
| 18 | } |
| 19 | var _ Locator = s |
| 20 | } |
| 21 | |
| 22 | func TestSwappableSnapshotCarriesItsReadersSource(t *testing.T) { |
| 23 | s := NewSwappable() |
| 24 | m, err := OpenMMDB("testdata/GeoIP2-City-Test.mmdb") |
| 25 | if err != nil { |
| 26 | t.Fatal(err) |
| 27 | } |
| 28 | s.Swap(m) |
| 29 | snap := s.Lookup(net.ParseIP("81.2.69.142")) |
| 30 | if !snap.Loaded || !snap.Found || snap.Source != "GeoIP2-City" || snap.Result.Country != "GB" { |
| 31 | t.Fatalf("snapshot = %+v", snap) |
| 32 | } |
| 33 | miss := s.Lookup(net.ParseIP("10.0.0.1")) |
| 34 | if !miss.Loaded || miss.Found || miss.Source != "GeoIP2-City" { |
| 35 | t.Fatalf("a miss with a loaded reader must say Loaded and name the source: %+v", miss) |
| 36 | } |
| 37 | s.Swap(nil) |
| 38 | if snap := s.Lookup(net.ParseIP("81.2.69.142")); snap.Loaded { |
| 39 | t.Fatalf("after Swap(nil) = %+v", snap) |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | // Review Focus 3: no lookup ever runs on a closed reader, and a snapshot's |
| 44 | // source always belongs to the reader that answered. |
| 45 | func TestSwappableNeverLooksUpAClosedReader(t *testing.T) { |
| 46 | s := NewSwappable() |
| 47 | var closedWhileInUse atomic.Bool |
| 48 | mk := func(src string) *MMDB { |
| 49 | var inUse atomic.Int32 |
| 50 | var closed atomic.Bool |
| 51 | m := &MMDB{source: src} |
| 52 | m.lookup = func(_ net.IP, result any) error { |
| 53 | inUse.Add(1) |
| 54 | defer inUse.Add(-1) |
| 55 | if closed.Load() { |
| 56 | closedWhileInUse.Store(true) |
| 57 | } |
| 58 | rec := result.(*mmdbRecord) |
| 59 | rec.Country.ISOCode = src[:2] |
| 60 | time.Sleep(time.Microsecond) |
| 61 | return nil |
| 62 | } |
| 63 | m.closeHook = func() { |
| 64 | if inUse.Load() != 0 { |
| 65 | closedWhileInUse.Store(true) |
| 66 | } |
| 67 | closed.Store(true) |
| 68 | } |
| 69 | return m |
| 70 | } |
| 71 | s.Swap(mk("AAsrc")) |
| 72 | var wg sync.WaitGroup |
| 73 | stop := make(chan struct{}) |
| 74 | var mismatched atomic.Bool |
| 75 | for i := 0; i < 32; i++ { |
| 76 | wg.Add(1) |
| 77 | go func() { |
| 78 | defer wg.Done() |
| 79 | for { |
| 80 | select { |
| 81 | case <-stop: |
| 82 | return |
| 83 | default: |
| 84 | } |
| 85 | snap := s.Lookup(net.ParseIP("81.2.69.142")) |
| 86 | if snap.Found && snap.Result.Country != snap.Source[:2] { |
| 87 | mismatched.Store(true) |
| 88 | } |
| 89 | s.Locate(net.ParseIP("81.2.69.142")) |
| 90 | } |
| 91 | }() |
| 92 | } |
| 93 | for i := 0; i < 100; i++ { |
| 94 | if i%2 == 0 { |
| 95 | s.Swap(mk("BBsrc")) |
| 96 | } else { |
| 97 | s.Swap(mk("AAsrc")) |
| 98 | } |
| 99 | } |
| 100 | close(stop) |
| 101 | wg.Wait() |
| 102 | if closedWhileInUse.Load() { |
| 103 | t.Fatal("a reader was closed while a lookup was using it") |
| 104 | } |
| 105 | if mismatched.Load() { |
| 106 | t.Fatal("a snapshot carried a different reader's source") |
| 107 | } |
| 108 | } |
| 109 | |
amadan