Fix race condition when opening and closing tsdb concurrently (#3959)

* Fix race condition when opening and closing tsdb concurrently

Signed-off-by: Alan Protasio <[email protected]>

* Changelog

Signed-off-by: Alan Protasio <[email protected]>

Co-authored-by: Marco Pracucci <[email protected]>

Author: alanprot

CHANGELOG.md

@@ -18,6 +18,7 @@
 * [ENHANCEMENT] Query-frontend/scheduler: added querier forget delay (`-query-frontend.querier-forget-delay` and `-query-scheduler.querier-forget-delay`) to mitigate the blast radius in the event queriers crash because of a repeatedly sent "query of death" when shuffle-sharding is enabled. #3901
 * [ENHANCEMENT] Ingester: reduce CPU and memory when an high number of errors are returned by the ingester on the write path with the blocks storage. #3969 #3971 #3973
 * [BUGFIX] Distributor: reverted changes done to rate limiting in #3825. #3948
+* [BUGFIX] Ingester: Fix race condition when opening and closing tsdb concurrently. #3959
 * [BUGFIX] Querier: streamline tracing spans. #3924
 
 ## 1.8.0 in progress

pkg/ingester/ingester_v2.go

@@ -1924,9 +1924,16 @@ func (i *Ingester) closeAndDeleteUserTSDBIfIdle(userID string) tsdbCloseCheckRes
 	// This will prevent going back to "active" state in deferred statement.
 	userDB.casState(closing, closed)
 
-	i.userStatesMtx.Lock()
-	delete(i.TSDBState.dbs, userID)
-	i.userStatesMtx.Unlock()
+	// Only remove user from TSDBState when everything is cleaned up
+	// This will prevent concurrency problems when cortex are trying to open new TSDB - Ie: New request for a given tenant
+	// came in - while closing the tsdb for the same tenant.
+	// If this happens now, the request will get reject as the push will not be able to acquire the lock as the tsdb will be
+	// in closed state
+	defer func() {
+		i.userStatesMtx.Lock()
+		delete(i.TSDBState.dbs, userID)
+		i.userStatesMtx.Unlock()
+	}()
 
 	i.metrics.memUsers.Dec()
 	i.TSDBState.tsdbMetrics.removeRegistryForUser(userID)

pkg/ingester/ingester_v2_test.go

@@ -2190,6 +2190,57 @@ func TestIngester_closeAndDeleteUserTSDBIfIdle_shouldNotCloseTSDBIfShippingIsInP
 	assert.Equal(t, tsdbNotActive, i.closeAndDeleteUserTSDBIfIdle(userID))
 }
 
+func TestIngester_closingAndOpeningTsdbConcurrently(t *testing.T) {
+	ctx := context.Background()
+	cfg := defaultIngesterTestConfig()
+	cfg.BlocksStorageConfig.TSDB.CloseIdleTSDBTimeout = 0 // Will not run the loop, but will allow us to close any TSDB fast.
+
+	// Create ingester
+	i, err := prepareIngesterWithBlocksStorage(t, cfg, nil)
+	require.NoError(t, err)
+
+	require.NoError(t, services.StartAndAwaitRunning(ctx, i))
+	defer services.StopAndAwaitTerminated(ctx, i) //nolint:errcheck
+
+	// Wait until it's ACTIVE
+	test.Poll(t, 1*time.Second, ring.ACTIVE, func() interface{} {
+		return i.lifecycler.GetState()
+	})
+
+	_, err = i.getOrCreateTSDB(userID, false)
+	require.NoError(t, err)
+
+	iterations := 5000
+	chanErr := make(chan error, 1)
+	quit := make(chan bool)
+
+	go func() {
+		for {
+			select {
+			case <-quit:
+				return
+			default:
+				_, err = i.getOrCreateTSDB(userID, false)
+				if err != nil {
+					chanErr <- err
+				}
+			}
+		}
+	}()
+
+	for k := 0; k < iterations; k++ {
+		i.closeAndDeleteUserTSDBIfIdle(userID)
+	}
+
+	select {
+	case err := <-chanErr:
+		assert.Fail(t, err.Error())
+		quit <- true
+	default:
+		quit <- true
+	}
+}
+
 func TestIngester_idleCloseEmptyTSDB(t *testing.T) {
 	ctx := context.Background()
 	cfg := defaultIngesterTestConfig()