diff --git a/config.toml.example b/config.toml.example index b4de6852..a3b10a15 100644 --- a/config.toml.example +++ b/config.toml.example @@ -77,6 +77,7 @@ write_behind_batch_size = 50 # Number of entries per batch write write_behind_batch_timeout = 100 # Max wait time in ms before flushing partial batch profile_routes = false # Turn on debugging endpoints profile_contention = false # Collect data for contention (use with above) - has a perf impact +s2_cell_lookup = false # Pre-compute S2 cell lookup for faster geofence matching. Trades memory (~60x geofence file size) for ~7x faster lookups. (default: false) # When enabled, reduce_updates will make fort update debounce windows much longer # to reduce database churn. Specifically, gym/pokestop/station debounce will be diff --git a/config/config.go b/config/config.go index 4649fc55..b9af097e 100644 --- a/config/config.go +++ b/config/config.go @@ -133,6 +133,7 @@ type tuning struct { WriteBehindWorkerCount int `koanf:"write_behind_worker_count"` // concurrent writers, default: 50 WriteBehindBatchSize int `koanf:"write_behind_batch_size"` // entries per batch, default: 50 WriteBehindBatchTimeoutMs int `koanf:"write_behind_batch_timeout"` // max wait for batch in ms, default: 100 + S2CellLookup bool `koanf:"s2_cell_lookup"` // Pre-compute S2 cell lookup for faster geofence matching. Trades memory (~60x geofence file size) for ~7x faster lookups, default: false } type scanRule struct { diff --git a/decoder/fort.go b/decoder/fort.go index 6247fbb5..36fe2fbe 100644 --- a/decoder/fort.go +++ b/decoder/fort.go @@ -26,6 +26,7 @@ type FortWebhook struct { Description *string `json:"description"` ImageUrl *string `json:"image_url"` Location Location `json:"location"` + CellId uint64 `json:"-"` // internal use only, not sent in webhook } type FortChangeWebhook struct { @@ -83,6 +84,7 @@ func InitWebHookFortFromGym(gym *Gym) *FortWebhook { fort.ImageUrl = gym.Url.Ptr() fort.Description = gym.Description.Ptr() fort.Location = Location{Latitude: gym.Lat, Longitude: gym.Lon} + fort.CellId = uint64(gym.CellId.ValueOrZero()) return fort } @@ -97,6 +99,7 @@ func InitWebHookFortFromPokestop(stop *Pokestop) *FortWebhook { fort.ImageUrl = stop.Url.Ptr() fort.Description = stop.Description.Ptr() fort.Location = Location{Latitude: stop.Lat, Longitude: stop.Lon} + fort.CellId = uint64(stop.CellId.ValueOrZero()) return fort } @@ -111,7 +114,7 @@ func CreateFortChangeWebhooks(fort *FortWebhook, change FortChange) { func CreateFortWebHooks(old *FortWebhook, new *FortWebhook, change FortChange) { if change == NEW { - areas := MatchStatsGeofence(new.Location.Latitude, new.Location.Longitude) + areas := MatchStatsGeofenceWithCell(new.Location.Latitude, new.Location.Longitude, new.CellId) hook := FortChangeWebhook{ ChangeType: change.String(), New: new, @@ -119,7 +122,7 @@ func CreateFortWebHooks(old *FortWebhook, new *FortWebhook, change FortChange) { webhooksSender.AddMessage(webhooks.FortUpdate, hook, areas) statsCollector.UpdateFortCount(areas, new.Type, "addition") } else if change == REMOVAL { - areas := MatchStatsGeofence(old.Location.Latitude, old.Location.Longitude) + areas := MatchStatsGeofenceWithCell(old.Location.Latitude, old.Location.Longitude, old.CellId) hook := FortChangeWebhook{ ChangeType: change.String(), Old: old, @@ -127,7 +130,7 @@ func CreateFortWebHooks(old *FortWebhook, new *FortWebhook, change FortChange) { webhooksSender.AddMessage(webhooks.FortUpdate, hook, areas) statsCollector.UpdateFortCount(areas, old.Type, "removal") } else if change == EDIT { - areas := MatchStatsGeofence(new.Location.Latitude, new.Location.Longitude) + areas := MatchStatsGeofenceWithCell(new.Location.Latitude, new.Location.Longitude, new.CellId) var editTypes []string // Check if Name has changed diff --git a/decoder/geography.go b/decoder/geography.go index ba452e34..49c73812 100644 --- a/decoder/geography.go +++ b/decoder/geography.go @@ -2,13 +2,17 @@ package decoder import ( "encoding/json" - "github.com/tidwall/rtree" - "golbat/geo" "io/ioutil" "net/http" + "sync/atomic" + + "golbat/config" + "golbat/geo" + "github.com/golang/geo/s2" "github.com/paulmach/orb/geojson" log "github.com/sirupsen/logrus" + "github.com/tidwall/rtree" ) type KojiResponse struct { @@ -19,8 +23,9 @@ type KojiResponse struct { // Stats KojiStats `json:"stats"` } -var statsTree *rtree.RTreeG[*geojson.Feature] -var nestTree *rtree.RTreeG[*geojson.Feature] +var statsTree atomic.Value +var nestTree atomic.Value +var statsS2Lookup atomic.Value var kojiUrl = "" var kojiBearerToken = "" @@ -105,15 +110,34 @@ func ReadGeofences() error { statsFeatureCollection = fc } - statsTree = geo.LoadRtree(statsFeatureCollection) + newStatsTree := geo.LoadRtree(statsFeatureCollection) + var newStatsS2Lookup *geo.S2CellLookup + if config.Config.Tuning.S2CellLookup { + newStatsS2Lookup = geo.BuildS2LookupFromFeatures(statsFeatureCollection) + } + + statsTree.Store(newStatsTree) + statsS2Lookup.Store(newStatsS2Lookup) return nil } func MatchStatsGeofence(lat, lon float64) []geo.AreaName { - return geo.MatchGeofencesRtree(statsTree, lat, lon) + return MatchStatsGeofenceWithCell(lat, lon, 0) +} + +func MatchStatsGeofenceWithCell(lat, lon float64, cellId uint64) []geo.AreaName { + lookup, _ := statsS2Lookup.Load().(*geo.S2CellLookup) + if cellId != 0 && lookup != nil { + if areas := lookup.Lookup(s2.CellID(cellId)); len(areas) > 0 { + return areas + } + } + tree, _ := statsTree.Load().(*rtree.RTreeG[*geojson.Feature]) + return geo.MatchGeofencesRtree(tree, lat, lon) } func MatchNestGeofence(lat, lon float64) []geo.AreaName { - return geo.MatchGeofencesRtree(nestTree, lat, lon) + tree, _ := nestTree.Load().(*rtree.RTreeG[*geojson.Feature]) + return geo.MatchGeofencesRtree(tree, lat, lon) } diff --git a/decoder/geography_test.go b/decoder/geography_test.go new file mode 100644 index 00000000..f93ed7af --- /dev/null +++ b/decoder/geography_test.go @@ -0,0 +1,129 @@ +//go:build ignore + +package decoder + +import ( + "encoding/json" + "math/rand" + "os" + "reflect" + "sort" + "testing" + "time" + + "golbat/geo" + + "github.com/golang/geo/s2" + "github.com/paulmach/orb/geojson" +) + +type KojiTestResponse struct { + Data geojson.FeatureCollection `json:"data"` +} + +func TestMatchStatsGeofenceWithCellVsRtree(t *testing.T) { + data, err := os.ReadFile("../cache/geofences.txt") + if err != nil { + t.Fatalf("Failed to read geofences.txt: %v", err) + } + + var response KojiTestResponse + if err := json.Unmarshal(data, &response); err != nil { + t.Fatalf("Failed to parse geofences.txt: %v", err) + } + + fc := &response.Data + t.Logf("Loaded %d features from geofences.txt", len(fc.Features)) + + testStatsTree := geo.LoadRtree(fc) + t.Logf("Built rtree") + testS2Lookup := geo.BuildS2LookupFromFeatures(fc) + t.Logf("Built S2 lookup with %d cells, size: %.2f MB", testS2Lookup.CellCount(), float64(testS2Lookup.SizeBytes())/(1024*1024)) + + minLat, maxLat := 49.002043, 54.854478 + minLon, maxLon := 14.120178, 24.145783 + + const numPoints = 50000 + rng := rand.New(rand.NewSource(42)) + + var mismatches int + var s2Hits int + var rtreeHits int + var rtreeTime time.Duration + var lookupTime time.Duration + + for i := 0; i < numPoints; i++ { + lat := minLat + rng.Float64()*(maxLat-minLat) + lon := minLon + rng.Float64()*(maxLon-minLon) + + cellID := s2.CellIDFromLatLng(s2.LatLngFromDegrees(lat, lon)).Parent(geo.S2LookupLevel) + + start := time.Now() + areasRtree := geo.MatchGeofencesRtree(testStatsTree, lat, lon) + rtreeTime += time.Since(start) + + start = time.Now() + areasS2 := testS2Lookup.Lookup(cellID) + var areasWithCell []geo.AreaName + if len(areasS2) > 0 { + areasWithCell = areasS2 + s2Hits++ + } else { + areasWithCell = geo.MatchGeofencesRtree(testStatsTree, lat, lon) + rtreeHits++ + } + lookupTime += time.Since(start) + + if !areasEqual(areasRtree, areasWithCell) { + mismatches++ + if mismatches <= 10 { + t.Logf("Mismatch at point %d (%.6f, %.6f): rtree=%v, withCell=%v", + i, lat, lon, areasRtree, areasWithCell) + } + } + } + + t.Logf("Results: %d points tested", numPoints) + t.Logf("S2 lookup hits: %d (%.2f%%)", s2Hits, float64(s2Hits)/float64(numPoints)*100) + t.Logf("Rtree fallback: %d (%.2f%%)", rtreeHits, float64(rtreeHits)/float64(numPoints)*100) + t.Logf("Mismatches: %d (%.2f%%)", mismatches, float64(mismatches)/float64(numPoints)*100) + + // Mismatches should be 0 — edge cells are excluded from the S2 lookup, + // so the rtree fallback always provides the complete result. + if mismatches > 0 { + t.Errorf("Expected 0 mismatches, got %d. Edge cells should be excluded from S2 lookup.", mismatches) + } + + t.Logf("Timing: rtree only: %v (%.2f µs/call)", rtreeTime, float64(rtreeTime.Microseconds())/float64(numPoints)) + t.Logf("Timing: lookup+fallback: %v (%.2f µs/call)", lookupTime, float64(lookupTime.Microseconds())/float64(numPoints)) + t.Logf("Speedup: %.2fx", float64(rtreeTime)/float64(lookupTime)) +} + +func areasEqual(a, b []geo.AreaName) bool { + if len(a) != len(b) { + return false + } + if len(a) == 0 { + return true + } + + aCopy := make([]geo.AreaName, len(a)) + bCopy := make([]geo.AreaName, len(b)) + copy(aCopy, a) + copy(bCopy, b) + + sort.Slice(aCopy, func(i, j int) bool { + if aCopy[i].Parent != aCopy[j].Parent { + return aCopy[i].Parent < aCopy[j].Parent + } + return aCopy[i].Name < aCopy[j].Name + }) + sort.Slice(bCopy, func(i, j int) bool { + if bCopy[i].Parent != bCopy[j].Parent { + return bCopy[i].Parent < bCopy[j].Parent + } + return bCopy[i].Name < bCopy[j].Name + }) + + return reflect.DeepEqual(aCopy, bCopy) +} diff --git a/decoder/gym_state.go b/decoder/gym_state.go index c60097ba..ff395751 100644 --- a/decoder/gym_state.go +++ b/decoder/gym_state.go @@ -401,7 +401,7 @@ func saveGymRecord(ctx context.Context, db db.DbDetails, gym *Gym) { fortRtreeUpdateGymOnSave(gym) } - areas := MatchStatsGeofence(gym.Lat, gym.Lon) + areas := MatchStatsGeofenceWithCell(gym.Lat, gym.Lon, uint64(gym.CellId.ValueOrZero())) createGymWebhooks(gym, areas) createGymFortWebhooks(gym) updateRaidStats(gym, areas) diff --git a/decoder/incident_state.go b/decoder/incident_state.go index f122bb2f..a4fd5916 100644 --- a/decoder/incident_state.go +++ b/decoder/incident_state.go @@ -146,13 +146,15 @@ func saveIncidentRecord(ctx context.Context, db db.DbDetails, incident *Incident createIncidentWebhooks(ctx, db, incident) var stopLat, stopLon float64 + var stopCellId uint64 stop, unlock, _ := getPokestopRecordReadOnly(ctx, db, incident.PokestopId) if stop != nil { stopLat, stopLon = stop.Lat, stop.Lon + stopCellId = uint64(stop.CellId.ValueOrZero()) unlock() } - areas := MatchStatsGeofence(stopLat, stopLon) + areas := MatchStatsGeofenceWithCell(stopLat, stopLon, stopCellId) updateIncidentStats(incident, areas) if config.Config.FortInMemory { @@ -217,12 +219,14 @@ func createIncidentWebhooks(ctx context.Context, db db.DbDetails, incident *Inci var pokestopName, stopUrl string var stopLat, stopLon float64 var stopEnabled bool + var stopCellId uint64 stop, unlock, _ := getPokestopRecordReadOnly(ctx, db, incident.PokestopId) if stop != nil { pokestopName = stop.Name.ValueOrZero() stopLat, stopLon = stop.Lat, stop.Lon stopUrl = stop.Url.ValueOrZero() stopEnabled = stop.Enabled.ValueOrZero() + stopCellId = uint64(stop.CellId.ValueOrZero()) unlock() } if pokestopName == "" { @@ -270,7 +274,7 @@ func createIncidentWebhooks(ctx context.Context, db db.DbDetails, incident *Inci Lineup: lineup, } - areas := MatchStatsGeofence(stopLat, stopLon) + areas := MatchStatsGeofenceWithCell(stopLat, stopLon, stopCellId) webhooksSender.AddMessage(webhooks.Invasion, incidentHook, areas) statsCollector.UpdateIncidentCount(areas) } diff --git a/decoder/pokemon_state.go b/decoder/pokemon_state.go index c1c6582b..3db16224 100644 --- a/decoder/pokemon_state.go +++ b/decoder/pokemon_state.go @@ -297,7 +297,7 @@ func savePokemonRecordAsAtTime(ctx context.Context, db db.DbDetails, pokemon *Po updatePokemonLookup(pokemon, changePvpField, pvpResults) // Webhooks and stats happen immediately (not queued) - areas := MatchStatsGeofence(pokemon.Lat, pokemon.Lon) + areas := MatchStatsGeofenceWithCell(pokemon.Lat, pokemon.Lon, uint64(pokemon.CellId.ValueOrZero())) if webhook { createPokemonWebhooks(ctx, db, pokemon, areas) } diff --git a/decoder/pokestop_process.go b/decoder/pokestop_process.go index e2792fd2..ca196eeb 100644 --- a/decoder/pokestop_process.go +++ b/decoder/pokestop_process.go @@ -53,7 +53,7 @@ func UpdatePokestopWithQuest(ctx context.Context, db db.DbDetails, quest *pogo.F updatePokestopGetMapFortCache(pokestop) savePokestopRecord(ctx, db, pokestop) - areas := MatchStatsGeofence(pokestop.Lat, pokestop.Lon) + areas := MatchStatsGeofenceWithCell(pokestop.Lat, pokestop.Lon, uint64(pokestop.CellId.ValueOrZero())) updateQuestStats(pokestop, haveAr, areas) return fmt.Sprintf("%s %s %s", quest.FortId, haveArStr, questTitle) diff --git a/decoder/pokestop_state.go b/decoder/pokestop_state.go index 40a0f960..1ed79c9c 100644 --- a/decoder/pokestop_state.go +++ b/decoder/pokestop_state.go @@ -213,7 +213,7 @@ func createPokestopFortWebhooks(stop *Pokestop) { func createPokestopWebhooks(stop *Pokestop) { - areas := MatchStatsGeofence(stop.Lat, stop.Lon) + areas := MatchStatsGeofenceWithCell(stop.Lat, stop.Lon, uint64(stop.CellId.ValueOrZero())) pokestopName := "Unknown" if stop.Name.Valid { diff --git a/decoder/station_state.go b/decoder/station_state.go index 267dbdd4..5ac0f06c 100644 --- a/decoder/station_state.go +++ b/decoder/station_state.go @@ -286,7 +286,7 @@ func createStationWebhooks(station *Station) { TotalStationedGmax: station.TotalStationedGmax, Updated: station.Updated, } - areas := MatchStatsGeofence(station.Lat, station.Lon) + areas := MatchStatsGeofenceWithCell(station.Lat, station.Lon, uint64(station.CellId)) webhooksSender.AddMessage(webhooks.MaxBattle, stationHook, areas) statsCollector.UpdateMaxBattleCount(areas, station.BattleLevel.ValueOrZero()) } diff --git a/geo/s2_lookup.go b/geo/s2_lookup.go new file mode 100644 index 00000000..61aeed23 --- /dev/null +++ b/geo/s2_lookup.go @@ -0,0 +1,176 @@ +package geo + +import ( + "runtime" + "sync" + "unsafe" + + "github.com/golang/geo/s2" + "github.com/paulmach/orb" + "github.com/paulmach/orb/geojson" + log "github.com/sirupsen/logrus" +) + +const S2LookupLevel = 15 + +type S2CellLookup struct { + cells map[s2.CellID][]AreaName + edgeCells map[s2.CellID]struct{} +} + +func NewS2CellLookup() *S2CellLookup { + return &S2CellLookup{ + cells: make(map[s2.CellID][]AreaName), + edgeCells: make(map[s2.CellID]struct{}), + } +} + +func (l *S2CellLookup) removeEdgeCells() int { + removed := 0 + for cellID := range l.edgeCells { + if _, exists := l.cells[cellID]; exists { + delete(l.cells, cellID) + removed++ + } + } + l.edgeCells = nil // free memory + return removed +} + +func (l *S2CellLookup) Lookup(cellID s2.CellID) []AreaName { + return l.cells[cellID] +} + +func (l *S2CellLookup) SizeBytes() int64 { + var size int64 + + size += int64(unsafe.Sizeof(l.cells)) + + for cellID, areas := range l.cells { + size += int64(unsafe.Sizeof(cellID)) + size += int64(unsafe.Sizeof(areas)) + for _, area := range areas { + size += int64(unsafe.Sizeof(area)) + size += int64(len(area.Name)) + size += int64(len(area.Parent)) + } + } + + return size +} + +func (l *S2CellLookup) CellCount() int { + return len(l.cells) +} + +type polygonWork struct { + polygon orb.Polygon + area AreaName +} + +func BuildS2LookupFromFeatures(featureCollection *geojson.FeatureCollection) *S2CellLookup { + if featureCollection == nil { + return NewS2CellLookup() + } + + lookup := NewS2CellLookup() + var mu sync.Mutex // Only used during build phase + + // Helper closures for thread-safe writes during build + addArea := func(cellID s2.CellID, area AreaName) { + mu.Lock() + lookup.cells[cellID] = append(lookup.cells[cellID], area) + mu.Unlock() + } + + addEdgeCell := func(cellID s2.CellID) { + mu.Lock() + lookup.edgeCells[cellID] = struct{}{} + mu.Unlock() + } + + numWorkers := max(runtime.NumCPU(), 4) + + workChan := make(chan polygonWork, 100) + var wg sync.WaitGroup + + for range numWorkers { + wg.Go(func() { + for work := range workChan { + processPolygon(work.polygon, work.area, addArea, addEdgeCell) + } + }) + } + + for _, f := range featureCollection.Features { + name := f.Properties.MustString("name", "unknown") + parent := f.Properties.MustString("parent", name) + area := AreaName{Parent: parent, Name: name} + + geoType := f.Geometry.GeoJSONType() + switch geoType { + case "Polygon": + polygon := f.Geometry.(orb.Polygon) + workChan <- polygonWork{polygon: polygon, area: area} + case "MultiPolygon": + multiPolygon := f.Geometry.(orb.MultiPolygon) + for _, polygon := range multiPolygon { + workChan <- polygonWork{polygon: polygon, area: area} + } + } + } + + close(workChan) + wg.Wait() + + removed := lookup.removeEdgeCells() + log.Infof("GEO: Removed %d edge cells from lookup", removed) + + sizeMB := float64(lookup.SizeBytes()) / (1024 * 1024) + log.Infof("GEO: S2 lookup table built with %d cells, size: %.2f MB", lookup.CellCount(), sizeMB) + + return lookup +} + +func processPolygon( + polygon orb.Polygon, + area AreaName, + addArea func(s2.CellID, AreaName), + addEdgeCell func(s2.CellID), +) { + if len(polygon) == 0 || len(polygon[0]) == 0 { + return + } + + // Convert orb.Polygon to s2.Loop for efficient covering + ring := polygon[0] // outer ring + points := make([]s2.Point, len(ring)) + for i, p := range ring { + points[i] = s2.PointFromLatLng(s2.LatLngFromDegrees(p.Lat(), p.Lon())) + } + + loop := s2.LoopFromPoints(points) + s2Polygon := s2.PolygonFromLoops([]*s2.Loop{loop}) + + coverer := s2.RegionCoverer{ + MinLevel: S2LookupLevel, + MaxLevel: S2LookupLevel, + } + covering := coverer.Covering(s2Polygon) + + for _, cellID := range covering { + cell := s2.CellFromCellID(cellID) + allInside := true + for i := range 4 { + if !s2Polygon.ContainsPoint(cell.Vertex(i)) { + allInside = false + break + } + } + if allInside { + addArea(cellID, area) + } else { + addEdgeCell(cellID) + } + } +}