fixing bug in chuck cache that can cause panic during shutdown

CHANGELOG.md

@@ -45,6 +45,7 @@
 * [BUGFIX] Ruler: fixed counting of PromQL evaluation errors as user-errors when updating `cortex_ruler_queries_failed_total`. #4335
 * [BUGFIX] Ingester: When using block storage, prevent any reads or writes while the ingester is stopping. This will prevent accessing TSDB blocks once they have been already closed. #4304
 * [BUGFIX] Ingester: fixed ingester stuck on start up (LEAVING ring state) when `-ingester.heartbeat-period=0` and `-ingester.unregister-on-shutdown=false`. #4366
+* [BUGFIX] Ingester: panic during shutdown while fetching batches from cache. #4397
 * [BUGFIX] Querier: After query-frontend restart, querier may have lower than configured concurrency. #4417
 
 ## 1.10.0 / 2021-08-03

pkg/chunk/cache/cache_test.go

@@ -125,6 +125,25 @@ func testChunkFetcher(t *testing.T, c cache.Cache, keys []string, chunks []chunk
 	require.Equal(t, chunks, found)
 }
 
+// testChunkFetcherStop checks that stopping the fetcher while fetching chunks don't result an error
+func testChunkFetcherStop(t *testing.T, c cache.Cache, keys []string, chunks []chunk.Chunk) {
+	fetcher, err := chunk.NewChunkFetcher(c, false, chunk.NewMockStorage())
+	require.NoError(t, err)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		if _, err := fetcher.FetchChunks(context.Background(), chunks, keys); err != nil {
+			// Since we stop fetcher while FetchChunks is running, we may not get everything back
+			// which requires the fetcher to fetch keys from storage, which is missing the keys
+			// so errors here is expected. Need to check the error because of the lint check.
+			require.NotNil(t, err)
+		}
+	}()
+	fetcher.Stop()
+	<-done
+}
+
 type byExternalKey []chunk.Chunk
 
 func (a byExternalKey) Len() int           { return len(a) }
@@ -155,6 +174,11 @@ func testCache(t *testing.T, cache cache.Cache) {
 	t.Run("Fetcher", func(t *testing.T) {
 		testChunkFetcher(t, cache, keys, chunks)
 	})
+	t.Run("FetcherStop", func(t *testing.T) {
+		// Refill the cache to avoid nil pointer error during fetch for getting missing keys from storage
+		keys, chunks = fillCache(t, cache)
+		testChunkFetcherStop(t, cache, keys, chunks)
+	})
 }
 
 func TestMemcache(t *testing.T) {

pkg/chunk/cache/memcached.go

@@ -55,6 +55,7 @@ type Memcached struct {
 
 	wg      sync.WaitGroup
 	inputCh chan *work
+	quit    chan struct{}
 
 	logger log.Logger
 }
@@ -83,19 +84,24 @@ func NewMemcached(cfg MemcachedConfig, client MemcachedClient, name string, reg
 	}
 
 	c.inputCh = make(chan *work)
+	c.quit = make(chan struct{})
 	c.wg.Add(cfg.Parallelism)
 
 	for i := 0; i < cfg.Parallelism; i++ {
 		go func() {
-			for input := range c.inputCh {
-				res := &result{
-					batchID: input.batchID,
+			defer c.wg.Done()
+			for {
+				select {
+				case <-c.quit:
+					return
+				case input := <-c.inputCh:
+					res := &result{
+						batchID: input.batchID,
+					}
+					res.found, res.bufs, res.missed = c.fetch(input.ctx, input.keys)
+					input.resultCh <- res
 				}
-				res.found, res.bufs, res.missed = c.fetch(input.ctx, input.keys)
-				input.resultCh <- res
 			}
-
-			c.wg.Done()
 		}()
 	}
 
@@ -187,11 +193,15 @@ func (c *Memcached) fetchKeysBatched(ctx context.Context, keys []string) (found
 	go func() {
 		for i, j := 0, 0; i < len(keys); i += batchSize {
 			batchKeys := keys[i:math.Min(i+batchSize, len(keys))]
-			c.inputCh <- &work{
+			select {
+			case <-c.quit:
+				return
+			case c.inputCh <- &work{
 				keys:     batchKeys,
 				ctx:      ctx,
 				resultCh: resultsCh,
 				batchID:  j,
+			}:
 			}
 			j++
 		}
@@ -205,13 +215,21 @@ func (c *Memcached) fetchKeysBatched(ctx context.Context, keys []string) (found
 
 	// We need to order found by the input keys order.
 	results := make([]*result, numResults)
+loopResults:
 	for i := 0; i < numResults; i++ {
-		result := <-resultsCh
-		results[result.batchID] = result
+		select {
+		case <-c.quit:
+			break loopResults
+		case result := <-resultsCh:
+			results[result.batchID] = result
+		}
 	}
 	close(resultsCh)
 
 	for _, result := range results {
+		if result == nil {
+			continue
+		}
 		found = append(found, result.found...)
 		bufs = append(bufs, result.bufs...)
 		missed = append(missed, result.missed...)
@@ -239,11 +257,15 @@ func (c *Memcached) Store(ctx context.Context, keys []string, bufs [][]byte) {
 
 // Stop does nothing.
 func (c *Memcached) Stop() {
-	if c.inputCh == nil {
+	if c.quit == nil {
 		return
 	}
 
-	close(c.inputCh)
+	select {
+	case <-c.quit:
+	default:
+		close(c.quit)
+	}
 	c.wg.Wait()
 }
 

pkg/chunk/cache/memcached_test.go

@@ -156,3 +156,39 @@ func testMemcacheFailing(t *testing.T, memcache *cache.Memcached) {
 		}
 	}
 }
+
+func TestMemcacheStop(t *testing.T) {
+	t.Run("unbatched", func(t *testing.T) {
+		client := newMockMemcacheFailing()
+		memcache := cache.NewMemcached(cache.MemcachedConfig{}, client,
+			"test", nil, log.NewNopLogger())
+
+		testMemcachedStopping(t, memcache)
+	})
+
+	t.Run("batched", func(t *testing.T) {
+		client := newMockMemcacheFailing()
+		memcache := cache.NewMemcached(cache.MemcachedConfig{
+			BatchSize:   10,
+			Parallelism: 5,
+		}, client, "test", nil, log.NewNopLogger())
+
+		testMemcachedStopping(t, memcache)
+	})
+}
+
+func testMemcachedStopping(t *testing.T, memcache *cache.Memcached) {
+	numKeys := 1000
+	ctx := context.Background()
+	keys := make([]string, 0, numKeys)
+	bufs := make([][]byte, 0, numKeys)
+	for i := 0; i < numKeys; i++ {
+		keys = append(keys, fmt.Sprint(i))
+		bufs = append(bufs, []byte(fmt.Sprint(i)))
+	}
+
+	memcache.Store(ctx, keys, bufs)
+
+	go memcache.Fetch(ctx, keys)
+	memcache.Stop()
+}

pkg/chunk/chunk_store_utils.go

@@ -86,6 +86,7 @@ type Fetcher struct {
 
 	wait           sync.WaitGroup
 	decodeRequests chan decodeRequest
+	quit           chan struct{}
 }
 
 type decodeRequest struct {
@@ -105,6 +106,7 @@ func NewChunkFetcher(cacher cache.Cache, cacheStubs bool, storage Client) (*Fetc
 		cache:          cacher,
 		cacheStubs:     cacheStubs,
 		decodeRequests: make(chan decodeRequest),
+		quit:           make(chan struct{}),
 	}
 
 	c.wait.Add(chunkDecodeParallelism)
@@ -117,22 +119,32 @@ func NewChunkFetcher(cacher cache.Cache, cacheStubs bool, storage Client) (*Fetc
 
 // Stop the ChunkFetcher.
 func (c *Fetcher) Stop() {
-	close(c.decodeRequests)
+	select {
+	case <-c.quit:
+	default:
+		close(c.quit)
+	}
+
 	c.wait.Wait()
 	c.cache.Stop()
 }
 
 func (c *Fetcher) worker() {
 	defer c.wait.Done()
 	decodeContext := NewDecodeContext()
-	for req := range c.decodeRequests {
-		err := req.chunk.Decode(decodeContext, req.buf)
-		if err != nil {
-			cacheCorrupt.Inc()
-		}
-		req.responses <- decodeResponse{
-			chunk: req.chunk,
-			err:   err,
+	for {
+		select {
+		case <-c.quit:
+			return
+		case req := <-c.decodeRequests:
+			err := req.chunk.Decode(decodeContext, req.buf)
+			if err != nil {
+				cacheCorrupt.Inc()
+			}
+			req.responses <- decodeResponse{
+				chunk: req.chunk,
+				err:   err,
+			}
 		}
 	}
 }
@@ -230,22 +242,31 @@ func (c *Fetcher) processCacheResponse(ctx context.Context, chunks []Chunk, keys
 
 	go func() {
 		for _, request := range requests {
-			c.decodeRequests <- request
+			select {
+			case <-c.quit:
+				return
+			case c.decodeRequests <- request:
+			}
 		}
 	}()
 
 	var (
 		err   error
 		found []Chunk
 	)
-	for i := 0; i < len(requests); i++ {
-		response := <-responses
 
-		// Don't exit early, as we don't want to block the workers.
-		if response.err != nil {
-			err = response.err
-		} else {
-			found = append(found, response.chunk)
+loopResponses:
+	for i := 0; i < len(requests); i++ {
+		select {
+		case <-c.quit:
+			break loopResponses
+		case response := <-responses:
+			// Don't exit early, as we don't want to block the workers.
+			if response.err != nil {
+				err = response.err
+			} else {
+				found = append(found, response.chunk)
+			}
 		}
 	}
 	return found, missing, err