- Fixed high concurrency panic on slow geoIP resolve mode
- Added test case for concurrent geodb access
This commit is contained in:
Toby Chui 2024-12-01 21:21:53 +08:00
parent 4a4483e09d
commit 6bf944e13c
4 changed files with 55 additions and 15 deletions

View File

@ -46,7 +46,6 @@ import (
"imuslab.com/zoraxy/mod/utils"
)
/* SIGTERM handler, do shutdown sequences before closing */
func SetupCloseHandler() {
c := make(chan os.Signal, 2)

View File

@ -3,6 +3,7 @@ package geodb
import (
_ "embed"
"net/http"
"sync"
"time"
"imuslab.com/zoraxy/mod/database"
@ -21,10 +22,10 @@ type Store struct {
geotrie *trie
geotrieIpv6 *trie
sysdb *database.Database
slowLookupCacheIpv4 map[string]string //Cache for slow lookup
slowLookupCacheIpv6 map[string]string //Cache for slow lookup
cacheClearTicker *time.Ticker //Ticker for clearing cache
cacheClearTickerStopChan chan bool //Stop channel for cache clear ticker
slowLookupCacheIpv4 sync.Map //Cache for slow lookup, ip -> cc
slowLookupCacheIpv6 sync.Map //Cache for slow lookup ipv6, ip -> cc
cacheClearTicker *time.Ticker //Ticker for clearing cache
cacheClearTickerStopChan chan bool //Stop channel for cache clear ticker
option *StoreOptions
}
@ -61,7 +62,7 @@ func NewGeoDb(sysdb *database.Database, option *StoreOptions) (*Store, error) {
}
if option.SlowLookupCacheClearInterval == 0 {
option.SlowLookupCacheClearInterval = 15 * time.Minute
option.SlowLookupCacheClearInterval = 30 * time.Minute
}
//Create a new store
@ -71,8 +72,8 @@ func NewGeoDb(sysdb *database.Database, option *StoreOptions) (*Store, error) {
geodbIpv6: parsedGeoDataIpv6,
geotrieIpv6: ipv6Trie,
sysdb: sysdb,
slowLookupCacheIpv4: make(map[string]string),
slowLookupCacheIpv6: make(map[string]string),
slowLookupCacheIpv4: sync.Map{},
slowLookupCacheIpv6: sync.Map{},
cacheClearTicker: time.NewTicker(option.SlowLookupCacheClearInterval),
cacheClearTickerStopChan: make(chan bool),
option: option,
@ -86,8 +87,8 @@ func NewGeoDb(sysdb *database.Database, option *StoreOptions) (*Store, error) {
case <-store.cacheClearTickerStopChan:
return
case <-thisGeoDBStore.cacheClearTicker.C:
thisGeoDBStore.slowLookupCacheIpv4 = make(map[string]string)
thisGeoDBStore.slowLookupCacheIpv6 = make(map[string]string)
thisGeoDBStore.slowLookupCacheIpv4 = sync.Map{}
thisGeoDBStore.slowLookupCacheIpv6 = sync.Map{}
}
}
}(thisGeoDBStore)

View File

@ -42,7 +42,7 @@ func TestTrieConstruct(t *testing.T) {
func TestResolveCountryCodeFromIP(t *testing.T) {
// Create a new store
store, err := geodb.NewGeoDb(nil, &geodb.StoreOptions{
false,
true,
true,
0,
})
@ -84,4 +84,24 @@ func TestResolveCountryCodeFromIP(t *testing.T) {
if info.CountryIsoCode != expected {
t.Errorf("expected country code %s, but got %s for IP %s", expected, info.CountryIsoCode, ip)
}
// Test for issue #401
// Create 100 concurrent goroutines to resolve country code for random IP addresses in the test cases above
for i := 0; i < 100; i++ {
go func() {
for _, testcase := range knownIpCountryMap {
ip := testcase[0]
expected := testcase[1]
info, err := store.ResolveCountryCodeFromIP(ip)
if err != nil {
t.Errorf("error resolving country code for IP %s: %v", ip, err)
return
}
if info.CountryIsoCode != expected {
t.Errorf("expected country code %s, but got %s for IP %s", expected, info.CountryIsoCode, ip)
}
}
}()
}
}

View File

@ -58,7 +58,8 @@ func (s *Store) slowSearchIpv4(ipAddr string) string {
}
//Check if already in cache
if cc, ok := s.slowLookupCacheIpv4[ipAddr]; ok {
cc := s.GetSlowSearchCachedIpv4(ipAddr)
if cc != "" {
return cc
}
@ -70,7 +71,7 @@ func (s *Store) slowSearchIpv4(ipAddr string) string {
inRange, _ := isIPv4InRange(startIp, endIp, ipAddr)
if inRange {
//Add to cache
s.slowLookupCacheIpv4[ipAddr] = cc
s.slowLookupCacheIpv4.Store(ipAddr, cc)
return cc
}
}
@ -83,7 +84,8 @@ func (s *Store) slowSearchIpv6(ipAddr string) string {
}
//Check if already in cache
if cc, ok := s.slowLookupCacheIpv6[ipAddr]; ok {
cc := s.GetSlowSearchCachedIpv6(ipAddr)
if cc != "" {
return cc
}
@ -95,9 +97,27 @@ func (s *Store) slowSearchIpv6(ipAddr string) string {
inRange, _ := isIPv6InRange(startIp, endIp, ipAddr)
if inRange {
//Add to cache
s.slowLookupCacheIpv6[ipAddr] = cc
s.slowLookupCacheIpv6.Store(ipAddr, cc)
return cc
}
}
return ""
}
// GetSlowSearchCachedIpv4 return the country code for the given ipv4 address, return empty string if not found
func (s *Store) GetSlowSearchCachedIpv4(ipAddr string) string {
cc, ok := s.slowLookupCacheIpv4.Load(ipAddr)
if ok {
return cc.(string)
}
return ""
}
// GetSlowSearchCachedIpv6 return the country code for the given ipv6 address, return empty string if not found
func (s *Store) GetSlowSearchCachedIpv6(ipAddr string) string {
cc, ok := s.slowLookupCacheIpv6.Load(ipAddr)
if ok {
return cc.(string)
}
return ""
}