wait for goroutines to finish (#423)

Fixes #422

Author: mangalaman93

cache.go

@@ -72,6 +72,7 @@ type Cache[K Key, V any] struct {
 	keyToHash func(K) (uint64, uint64)
 	// stop is used to stop the processItems goroutine.
 	stop chan struct{}
+	done chan struct{}
 	// indicates whether cache is closed.
 	isClosed atomic.Bool
 	// cost calculates cost from a value.
@@ -227,6 +228,7 @@ func NewCache[K Key, V any](config *Config[K, V]) (*Cache[K, V], error) {
 		setBuf:             make(chan *Item[V], setBufSize),
 		keyToHash:          config.KeyToHash,
 		stop:               make(chan struct{}),
+		done:               make(chan struct{}),
 		cost:               config.Cost,
 		ignoreInternalCost: config.IgnoreInternalCost,
 		cleanupTicker:      time.NewTicker(time.Duration(config.TtlTickerDurationInSec) * time.Second / 2),
@@ -422,7 +424,9 @@ func (c *Cache[K, V]) Close() {
 
 	// Block until processItems goroutine is returned.
 	c.stop <- struct{}{}
+	<-c.done
 	close(c.stop)
+	close(c.done)
 	close(c.setBuf)
 	c.cachePolicy.Close()
 	c.cleanupTicker.Stop()
@@ -438,6 +442,7 @@ func (c *Cache[K, V]) Clear() {
 	}
 	// Block until processItems goroutine is returned.
 	c.stop <- struct{}{}
+	<-c.done
 
 	// Clear out the setBuf channel.
 loop:
@@ -556,6 +561,7 @@ func (c *Cache[K, V]) processItems() {
 		case <-c.cleanupTicker.C:
 			c.storedItems.Cleanup(c.cachePolicy, onEvict)
 		case <-c.stop:
+			c.done <- struct{}{}
 			return
 		}
 	}

cache_test.go

@@ -385,6 +385,7 @@ func TestCacheSet(t *testing.T) {
 	require.Equal(t, 2, val)
 
 	c.stop <- struct{}{}
+	<-c.done
 	for i := 0; i < setBufSize; i++ {
 		key, conflict := z.KeyToHash(1)
 		c.setBuf <- &Item[int]{
@@ -399,6 +400,7 @@ func TestCacheSet(t *testing.T) {
 	require.Equal(t, uint64(1), c.Metrics.SetsDropped())
 	close(c.setBuf)
 	close(c.stop)
+	close(c.done)
 
 	c = nil
 	require.False(t, c.Set(1, 1, 1))

policy.go

@@ -40,6 +40,7 @@ type defaultPolicy[V any] struct {
 	evict    *sampledLFU
 	itemsCh  chan []uint64
 	stop     chan struct{}
+	done     chan struct{}
 	isClosed bool
 	metrics  *Metrics
 }
@@ -50,6 +51,7 @@ func newDefaultPolicy[V any](numCounters, maxCost int64) *defaultPolicy[V] {
 		evict:   newSampledLFU(maxCost),
 		itemsCh: make(chan []uint64, 3),
 		stop:    make(chan struct{}),
+		done:    make(chan struct{}),
 	}
 	go p.processItems()
 	return p
@@ -73,6 +75,7 @@ func (p *defaultPolicy[V]) processItems() {
 			p.admit.Push(items)
 			p.Unlock()
 		case <-p.stop:
+			p.done <- struct{}{}
 			return
 		}
 	}
@@ -226,7 +229,9 @@ func (p *defaultPolicy[V]) Close() {
 
 	// Block until the p.processItems goroutine returns.
 	p.stop <- struct{}{}
+	<-p.done
 	close(p.stop)
+	close(p.done)
 	close(p.itemsCh)
 	p.isClosed = true
 }

policy_test.go

@@ -31,6 +31,7 @@ func TestPolicyProcessItems(t *testing.T) {
 	p.Unlock()
 
 	p.stop <- struct{}{}
+	<-p.done
 	p.itemsCh <- []uint64{3, 3, 3}
 	time.Sleep(wait)
 	p.Lock()