diff --git a/main.go b/main.go index 172b085..7a489dd 100644 --- a/main.go +++ b/main.go @@ -19,6 +19,7 @@ import ( "sync/atomic" "syscall" "time" + "unsafe" _ "net/http/pprof" @@ -45,7 +46,7 @@ var ( var ( _BackendAddrCacheMutex sync.Mutex - _BackendAddrCache atomic.Value + _BackendAddrCache unsafe.Pointer ) var ( @@ -60,7 +61,7 @@ func main() { runtime.GOMAXPROCS(runtime.NumCPU()) os.Setenv("GOTRACEBACK", "crash") - _BackendAddrCache.Store(make(backendAddrMap)) + _BackendAddrCache = unsafe.Pointer(&map[uint]uint{}) var lim syscall.Rlimit syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim) @@ -339,7 +340,7 @@ func dialTimeout(network, address string, timeout time.Duration) (conn net.Conn, func backendAddrDecrypt(key []byte) ([]byte, error) { // Try to check cache - m1 := _BackendAddrCache.Load().(backendAddrMap) + m1 := *(*backendAddrMap)(atomic.LoadPointer(&_BackendAddrCache)) k1 := string(key) addr, ok := m1[k1] if ok { @@ -357,25 +358,26 @@ func backendAddrDecrypt(key []byte) ([]byte, error) { } func backendAddrList(key string, val []byte) { - _BackendAddrCacheMutex.Lock() - defer _BackendAddrCacheMutex.Unlock() - - m1 := _BackendAddrCache.Load().(backendAddrMap) - // double check - if _, ok := m1[key]; ok { - return - } + done := false + for !done { + mptr := (atomic.LoadPointer(&_BackendAddrCache)) + m1 := *(*backendAddrMap)(mptr) + // double check + if _, ok := m1[key]; ok { + return + } - m2 := make(backendAddrMap) - // flush cache if there is way too many - if len(m1) < _MaxBackendAddrCacheCount { - // copy-on-write - for k, v := range m1 { - m2[k] = v // copy all data from the current object to the new one + m2 := make(backendAddrMap) + // flush cache if there is way too many + if len(m1) < _MaxBackendAddrCacheCount { + // copy-on-write + for k, v := range m1 { + m2[k] = v // copy all data from the current object to the new one + } } + m2[key] = val + done = atomic.CompareAndSwapPointer(&_BackendAddrCache, mptr, unsafe.Pointer(&m2)) } - m2[key] = val - _BackendAddrCache.Store(m2) // atomically replace the current object with the new one } // Request.RemoteAddress contains port, which we want to remove i.e.: