From 5a1ae974a763a4e471648c7d7855bdcdb515fe3f Mon Sep 17 00:00:00 2001 From: James Berry Date: Tue, 1 Sep 2026 22:34:32 +0100 Subject: [PATCH] fix: make masterfile reload safe under concurrent queries Reloading the masterfile while queries run crashed consumers with "fatal error: concurrent map read and map write" (hit in production by Golbat, see UnownHash/Golbat#403): - LoadPokemonData unmarshaled into the live PokemonData struct; encoding/json reuses non-nil maps, writing into the very maps QueryPvPRank readers walk with no lock. - WatchPokemonData published a new PokemonData via an unsynchronized multi-word struct assignment (torn reads, no happens-before). - ClearCache overwrote the live sync.Map while queries were mid Load/Store on it. The masterfile snapshot and the rank cache computed from it now live in one immutable bundle behind an atomic.Pointer. Reloads build a fresh bundle (fresh maps, empty cache) and publish it with a single atomic store; readers pick up the bundle once per call. This also closes the window where a reload served new masterfile data with ranks cached from the old one (data was stored before the cache was cleared), and stops calculateAllRanksCompact from lazily writing RankingComparator on the query path. The exported PokemonData field is kept for backward compatibility: it still seeds the first bundle when populated directly, and reloads mirror into it for legacy readers, but internal code no longer reads it after publication. Regression test exercises LoadPokemonData and ClearCache concurrently with QueryPvPRank and CalculateCp under -race. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01BMJMXHUiq6Pnce9KshgFTC --- ohbem.go | 108 ++++++++++++++++++++++++++++++++------------- ohbem_race_test.go | 62 ++++++++++++++++++++++++++ ohbem_test.go | 6 +-- structs.go | 17 ++++++- utils.go | 13 +++--- 5 files changed, 165 insertions(+), 41 deletions(-) create mode 100644 ohbem_race_test.go diff --git a/ohbem.go b/ohbem.go index dfef4c7..206052c 100644 --- a/ohbem.go +++ b/ohbem.go @@ -7,7 +7,6 @@ import ( "os" "reflect" "sort" - "sync" "time" ) @@ -19,13 +18,11 @@ const VERSION = "0.12.0" // FetchPokemonData Fetch remote MasterFile and keep it in memory. func (o *Ohbem) FetchPokemonData() error { - var err error - - o.PokemonData, err = fetchMasterFile() + pokemonData, err := fetchMasterFile() if err != nil { return err } - o.ClearCache() + o.publishPokemonData(pokemonData) return nil } @@ -35,17 +32,22 @@ func (o *Ohbem) LoadPokemonData(filePath string) error { if err != nil { return ErrMasterFileOpen } - if err := json.Unmarshal(data, &o.PokemonData); err != nil { + var pokemonData PokemonData + if err := json.Unmarshal(data, &pokemonData); err != nil { return ErrMasterFileUnmarshall } - o.PokemonData.Initialized = true - o.ClearCache() + pokemonData.Initialized = true + o.publishPokemonData(pokemonData) return nil } // SavePokemonData Save MasterFile from memory to provided location. func (o *Ohbem) SavePokemonData(filePath string) error { - data, err := json.Marshal(o.PokemonData) + var pokemonData PokemonData + if b := o.currentBundle(); b != nil { + pokemonData = b.data + } + data, err := json.Marshal(pokemonData) if err != nil { return ErrMasterFileMarshall } @@ -88,13 +90,15 @@ func (o *Ohbem) WatchPokemonData() error { o.log("Remote MasterFile fetch failed") continue } - if reflect.DeepEqual(o.PokemonData, pokemonData) { + var current PokemonData + if b := o.currentBundle(); b != nil { + current = b.data + } + if reflect.DeepEqual(current, pokemonData) { continue } else { o.log("New MasterFile found! Updating PokemonData") - o.PokemonData = pokemonData // overwrite PokemonData using new MasterFile - o.PokemonData.Initialized = true - o.ClearCache() // clean compactRankCache cache + o.publishPokemonData(pokemonData) // fresh snapshot with an empty rank cache // when provided store latest version of MasterFile under provided path if o.MasterFileCachePath != "" { err = o.SavePokemonData(o.MasterFileCachePath) @@ -120,24 +124,59 @@ func (o *Ohbem) StopWatchingPokemonData() error { return nil } +// publishPokemonData atomically swaps in a fresh MasterFile snapshot with an +// empty rank cache. Readers already holding the previous snapshot finish on +// it undisturbed. The exported PokemonData field is updated as a best-effort +// mirror for legacy callers; see its comment for the concurrency caveat. +func (o *Ohbem) publishPokemonData(data PokemonData) { + o.PokemonData = data + o.bundle.Store(&pokemonBundle{data: data}) +} + +// currentBundle returns the active MasterFile snapshot, or nil when no data +// is loaded. If PokemonData was populated directly instead of via +// Load/FetchPokemonData, it is promoted to a snapshot on first use. +func (o *Ohbem) currentBundle() *pokemonBundle { + if b := o.bundle.Load(); b != nil { + return b + } + if !o.PokemonData.Initialized { + return nil + } + b := &pokemonBundle{data: o.PokemonData} + if o.bundle.CompareAndSwap(nil, b) { + return b + } + return o.bundle.Load() +} + func (o *Ohbem) ClearCache() { if !o.DisableCache { - o.compactRankCache = sync.Map{} + for { + b := o.bundle.Load() + if b == nil { + break // nothing loaded, nothing cached + } + if o.bundle.CompareAndSwap(b, &pokemonBundle{data: b.data}) { + break + } + } o.log("Cache cleaned") } } // calculateAllRanksCompact Calculate all PvP ranks for a specific base stats with the specified CP cap. Compact version intended to be used with cache. -func (o *Ohbem) calculateAllRanksCompact(stats *PokemonStats, cpCap int) (map[int]compactCacheValue, bool) { +func (o *Ohbem) calculateAllRanksCompact(b *pokemonBundle, stats *PokemonStats, cpCap int) (map[int]compactCacheValue, bool) { cacheKey := int64(cpCap*999*999*999 + stats.Attack*999*999 + stats.Defense*999 + stats.Stamina) if !o.DisableCache { - if obj, ok := o.compactRankCache.Load(cacheKey); ok { + if obj, ok := b.cache.Load(cacheKey); ok { return obj.(map[int]compactCacheValue), true } } - if o.RankingComparator == nil { - o.RankingComparator = RankingComparatorDefault + comparator := o.RankingComparator + if comparator == nil { + comparator = RankingComparatorDefault } filled := false @@ -150,7 +189,7 @@ func (o *Ohbem) calculateAllRanksCompact(stats *PokemonStats, cpCap int) (map[in continue } - combinations, sortedRanks := calculateRanksCompact(stats, cpCap, lvCapFloat, o.RankingComparator, 0) + combinations, sortedRanks := calculateRanksCompact(stats, cpCap, lvCapFloat, comparator, 0) res := compactCacheValue{ Combinations: combinations, TopValue: sortedRanks[0].Value, @@ -163,7 +202,7 @@ func (o *Ohbem) calculateAllRanksCompact(stats *PokemonStats, cpCap int) (map[in } } if filled && !maxed { - combinations, sortedRanks := calculateRanksCompact(stats, cpCap, MaxLevel, o.RankingComparator, 0) + combinations, sortedRanks := calculateRanksCompact(stats, cpCap, MaxLevel, comparator, 0) res := compactCacheValue{ Combinations: combinations, @@ -172,7 +211,7 @@ func (o *Ohbem) calculateAllRanksCompact(stats *PokemonStats, cpCap int) (map[in result[MaxLevel] = res } if !o.DisableCache && filled { - o.compactRankCache.Store(cacheKey, result) + b.cache.Store(cacheKey, result) } return result, filled } @@ -355,7 +394,11 @@ func (o *Ohbem) CalculateTopRanks(maxRank int16, pokemonId int, form int, evolut // CalculateCp calculates CP for your pokemon. Errors if pokemon cannot be found in master. func (o *Ohbem) CalculateCp(pokemonId, form, evolution, attack, defense, stamina int, level float64) (int, error) { - masterPokemon, ok := o.PokemonData.Pokemon[pokemonId] + var pokemonData PokemonData + if b := o.currentBundle(); b != nil { + pokemonData = b.data + } + masterPokemon, ok := pokemonData.Pokemon[pokemonId] if !ok { return 0, ErrMissingPokemon } @@ -393,7 +436,8 @@ func (o *Ohbem) CalculateCp(pokemonId, form, evolution, attack, defense, stamina func (o *Ohbem) QueryPvPRank(pokemonId int, form int, costume int, gender int, attack int, defense int, stamina int, level float64) (map[string][]PokemonEntry, error) { result := make(map[string][]PokemonEntry) - if err := safetyCheck(o); err != nil { + b, err := safetyCheck(o) + if err != nil { return result, err } @@ -405,8 +449,8 @@ func (o *Ohbem) QueryPvPRank(pokemonId int, form int, costume int, gender int, a var masterPokemon Pokemon var baseEntry = PokemonEntry{Pokemon: pokemonId} - if _, ok := o.PokemonData.Pokemon[pokemonId]; ok { - masterPokemon = o.PokemonData.Pokemon[pokemonId] + if _, ok := b.data.Pokemon[pokemonId]; ok { + masterPokemon = b.data.Pokemon[pokemonId] } else { return result, ErrMissingPokemon } @@ -434,7 +478,7 @@ func (o *Ohbem) QueryPvPRank(pokemonId int, form int, costume int, gender int, a if leagueOptions.LittleCupRules && !(masterForm.Little || masterPokemon.Little) { continue } - combinationIndex, filled := o.calculateAllRanksCompact(stats, leagueOptions.Cap) + combinationIndex, filled := o.calculateAllRanksCompact(b, stats, leagueOptions.Cap) if !filled { continue } @@ -529,7 +573,7 @@ func (o *Ohbem) QueryPvPRank(pokemonId int, form int, costume int, gender int, a canEvolve := true if costume != 0 { - canEvolve = !o.PokemonData.Costumes[costume] || containsInt(masterForm.CostumeOverrideEvolutions, costume) + canEvolve = !b.data.Costumes[costume] || containsInt(masterForm.CostumeOverrideEvolutions, costume) } if canEvolve && len(masterForm.Evolutions) != 0 { for _, evolution := range masterForm.Evolutions { @@ -577,11 +621,12 @@ func (o *Ohbem) QueryPvPRank(pokemonId int, form int, costume int, gender int, a // FindBaseStats Look up base stats of a Pokémon. func (o *Ohbem) FindBaseStats(pokemonId int, form int, evolution int) (PokemonStats, error) { - if err := safetyCheck(o); err != nil { + b, err := safetyCheck(o) + if err != nil { return PokemonStats{}, err } - masterPokemon, ok := o.PokemonData.Pokemon[pokemonId] + masterPokemon, ok := b.data.Pokemon[pokemonId] if !ok { return PokemonStats{}, ErrMissingPokemon } @@ -628,11 +673,12 @@ func (o *Ohbem) FindBaseStats(pokemonId int, form int, evolution int) (PokemonSt // IsMegaUnreleased Check whether the stats for a given mega is speculated. func (o *Ohbem) IsMegaUnreleased(pokemonId int, evolution int) (bool, error) { - if err := safetyCheck(o); err != nil { + b, err := safetyCheck(o) + if err != nil { return false, err } - masterPokemon := o.PokemonData.Pokemon[pokemonId] + masterPokemon := b.data.Pokemon[pokemonId] if masterPokemon.Attack != 0 { evo := masterPokemon.TempEvolutions[evolution] return evo.Unreleased, nil diff --git a/ohbem_race_test.go b/ohbem_race_test.go new file mode 100644 index 0000000..6442914 --- /dev/null +++ b/ohbem_race_test.go @@ -0,0 +1,62 @@ +package gohbem + +import ( + "sync" + "testing" +) + +// TestReloadConcurrentWithQueries pins that reloading the masterfile +// (LoadPokemonData, as WatchPokemonData does on a changed remote file) and +// clearing the rank cache are safe while other goroutines query. Run under +// -race. Before the atomic snapshot change, LoadPokemonData unmarshaled into +// the live PokemonData maps and ClearCache overwrote the live sync.Map, +// crashing production consumers with "concurrent map read and map write" +// (UnownHash/Golbat#403). +func TestReloadConcurrentWithQueries(t *testing.T) { + ohbem := Ohbem{Leagues: leagues, LevelCaps: levelCaps} + if err := ohbem.LoadPokemonData("./test/master-test.json"); err != nil { + t.Fatalf("LoadPokemonData: %v", err) + } + + // Sanity checks: the queries below must reach the masterfile maps and the + // rank cache, otherwise the concurrent phase is vacuous. + if pvp, err := ohbem.QueryPvPRank(661, 0, 0, 1, 15, 15, 14, 1); err != nil || len(pvp) == 0 { + t.Fatalf("QueryPvPRank sanity check failed: pvp=%v err=%v", pvp, err) + } + if _, err := ohbem.CalculateCp(661, 0, 0, 15, 15, 14, 40); err != nil { + t.Fatalf("CalculateCp sanity check failed: %v", err) + } + + stop := make(chan struct{}) + var wg sync.WaitGroup + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + if _, err := ohbem.QueryPvPRank(661, 0, 0, 1, 15, 15, 14, 1); err != nil { + t.Errorf("QueryPvPRank during reload: %v", err) + return + } + if _, err := ohbem.CalculateCp(661, 0, 0, 15, 15, 14, 40); err != nil { + t.Errorf("CalculateCp during reload: %v", err) + return + } + } + }() + } + + for i := 0; i < 5; i++ { + if err := ohbem.LoadPokemonData("./test/master-test.json"); err != nil { + t.Errorf("LoadPokemonData during queries: %v", err) + } + ohbem.ClearCache() + } + close(stop) + wg.Wait() +} diff --git a/ohbem_test.go b/ohbem_test.go index 926aae8..37c2bac 100644 --- a/ohbem_test.go +++ b/ohbem_test.go @@ -52,7 +52,7 @@ func TestCalculateAllRanksCompact(t *testing.T) { for ix, test := range tests { testName := fmt.Sprintf("%d", ix) t.Run(testName, func(t *testing.T) { - combinations, _ := ohbem.calculateAllRanksCompact(&test.stats, test.cpCap) + combinations, _ := ohbem.calculateAllRanksCompact(ohbem.currentBundle(), &test.stats, test.cpCap) comb := combinations[test.lvCap].Combinations topValue := combinations[test.lvCap].TopValue if comb[test.pos] != test.value || topValue != test.topVal { @@ -68,7 +68,7 @@ func BenchmarkCalculateAllRanksCompact(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { - _, _ = ohbem.calculateAllRanksCompact(&PikachuStats, 5000) + _, _ = ohbem.calculateAllRanksCompact(ohbem.currentBundle(), &PikachuStats, 5000) } } @@ -78,7 +78,7 @@ func BenchmarkCalculateAllRanksCompactCached(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { - _, _ = ohbem.calculateAllRanksCompact(&PikachuStats, 5000) + _, _ = ohbem.calculateAllRanksCompact(ohbem.currentBundle(), &PikachuStats, 5000) } } diff --git a/structs.go b/structs.go index 90c5bba..436363a 100644 --- a/structs.go +++ b/structs.go @@ -2,11 +2,17 @@ package gohbem import ( "sync" + "sync/atomic" "time" ) // Ohbem struct is holding main configuration, cache and channels. type Ohbem struct { + // PokemonData may be populated directly before first use as an + // alternative to Load/FetchPokemonData. After that it is only a + // best-effort mirror kept for backward compatibility: it is not + // synchronized, so do not read it concurrently with a reload — use the + // query methods instead, which read the atomically published snapshot. PokemonData PokemonData LevelCaps []int Leagues map[string]League @@ -15,11 +21,20 @@ type Ohbem struct { RankingComparator RankingComparator IncludeHundosUnderCap bool WatcherInterval time.Duration - compactRankCache sync.Map + bundle atomic.Pointer[pokemonBundle] watcherChan chan bool Logger Logger } +// pokemonBundle pairs one immutable MasterFile snapshot with the rank cache +// computed from it. Reloads publish a fresh bundle through Ohbem.bundle, so +// readers never observe data and cache from different MasterFile versions, +// and never race a reload mutating live maps. +type pokemonBundle struct { + data PokemonData + cache sync.Map +} + // Logger interface // // Example implementation: diff --git a/utils.go b/utils.go index f96890b..1750272 100644 --- a/utils.go +++ b/utils.go @@ -48,17 +48,18 @@ func fetchMasterFile() (PokemonData, error) { return data, nil } -func safetyCheck(o *Ohbem) error { - if !o.PokemonData.Initialized { - return ErrMasterFileUnloaded +func safetyCheck(o *Ohbem) (*pokemonBundle, error) { + b := o.currentBundle() + if b == nil { + return nil, ErrMasterFileUnloaded } if len(o.Leagues) == 0 { - return ErrLeaguesMissing + return nil, ErrLeaguesMissing } if len(o.LevelCaps) == 0 { - return ErrLevelCapsMissing + return nil, ErrLevelCapsMissing } - return nil + return b, nil } // log logs the given message using the provided logger, if available. If no logger is set, the message is ignored.