From e6a286fb3eec7f50e1992ea1468d3adecebdb7a7 Mon Sep 17 00:00:00 2001
From: Weizhen Wang <wangweizhen@pingcap.com>
Date: Tue, 5 Apr 2022 14:33:02 +0800
Subject: [PATCH] fix: use atomic in the cache

Signed-off-by: Weizhen Wang <wangweizhen@pingcap.com>
---
 cache.go | 17 +++++++++--------
 go.mod   |  1 +
 go.sum   |  3 +++
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/cache.go b/cache.go
index af9a67e..18e3647 100644
--- a/cache.go
+++ b/cache.go
@@ -26,6 +26,7 @@ import (
 	"unsafe"
 
 	"github.com/outcaste-io/ristretto/z"
+	"go.uber.org/atomic"
 )
 
 var (
@@ -64,7 +65,7 @@ type Cache struct {
 	// stop is used to stop the processItems goroutine.
 	stop chan struct{}
 	// indicates whether cache is closed.
-	isClosed bool
+	isClosed atomic.Bool
 	// cost calculates cost from a value.
 	cost func(value interface{}) int64
 	// ignoreInternalCost dictates whether to ignore the cost of internally storing
@@ -214,7 +215,7 @@ func NewCache(config *Config) (*Cache, error) {
 }
 
 func (c *Cache) Wait() {
-	if c == nil || c.isClosed {
+	if c == nil || c.isClosed.Load() {
 		return
 	}
 	wg := &sync.WaitGroup{}
@@ -227,7 +228,7 @@ func (c *Cache) Wait() {
 // value was found or not. The value can be nil and the boolean can be true at
 // the same time.
 func (c *Cache) Get(key interface{}) (interface{}, bool) {
-	if c == nil || c.isClosed || key == nil {
+	if c == nil || c.isClosed.Load() || key == nil {
 		return nil, false
 	}
 	keyHash, conflictHash := c.keyToHash(key)
@@ -270,7 +271,7 @@ func (c *Cache) SetIfPresent(key, value interface{}, cost int64) bool {
 
 func (c *Cache) setInternal(key, value interface{},
 	cost int64, ttl time.Duration, onlyUpdate bool) bool {
-	if c == nil || c.isClosed || key == nil {
+	if c == nil || c.isClosed.Load() || key == nil {
 		return false
 	}
 
@@ -326,7 +327,7 @@ func (c *Cache) setInternal(key, value interface{},
 
 // Del deletes the key-value item from the cache if it exists.
 func (c *Cache) Del(key interface{}) {
-	if c == nil || c.isClosed || key == nil {
+	if c == nil || c.isClosed.Load() || key == nil {
 		return
 	}
 	keyHash, conflictHash := c.keyToHash(key)
@@ -373,7 +374,7 @@ func (c *Cache) GetTTL(key interface{}) (time.Duration, bool) {
 
 // Close stops all goroutines and closes all channels.
 func (c *Cache) Close() {
-	if c == nil || c.isClosed {
+	if c == nil || c.isClosed.Load() {
 		return
 	}
 	c.Clear()
@@ -383,14 +384,14 @@ func (c *Cache) Close() {
 	close(c.stop)
 	close(c.setBuf)
 	c.policy.Close()
-	c.isClosed = true
+	c.isClosed.Store(true)
 }
 
 // Clear empties the hashmap and zeroes all policy counters. Note that this is
 // not an atomic operation (but that shouldn't be a problem as it's assumed that
 // Set/Get calls won't be occurring until after this).
 func (c *Cache) Clear() {
-	if c == nil || c.isClosed {
+	if c == nil || c.isClosed.Load() {
 		return
 	}
 	// Block until processItems goroutine is returned.
diff --git a/go.mod b/go.mod
index 2721911..7c53b7e 100644
--- a/go.mod
+++ b/go.mod
@@ -10,5 +10,6 @@ require (
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.7.0
+	go.uber.org/atomic v1.9.0 // indirect
 	golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f
 )
diff --git a/go.sum b/go.sum
index 7ea0736..7327953 100644
--- a/go.sum
+++ b/go.sum
@@ -14,8 +14,11 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
+go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f h1:+Nyd8tzPX9R7BWHguqsrbFdRx3WQ/1ib8I44HXV5yTA=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=