userTSDB: Avoid race condition when setting db

Signed-off-by: Arve Knudsen <[email protected]>

pkg/ingester/ingester.go

@@ -2720,7 +2720,7 @@ func (i *Ingester) createTSDB(userID string, walReplayConcurrency int) (*userTSD
 		return nil, errors.Wrapf(err, "failed to compact TSDB: %s", udir)
 	}
 
-	userDB.db = db
+	userDB.setDB(db)
 	// We set the limiter here because we don't want to limit
 	// series during WAL replay.
 	userDB.limiter = i.limiter

pkg/ingester/user_tsdb.go

@@ -92,6 +92,7 @@ type ownedSeriesState struct {
 }
 
 type userTSDB struct {
+	hasDB          atomic.Bool
 	db             *tsdb.DB
 	userID         string
 	activeSeries   *activeseries.ActiveSeries
@@ -144,6 +145,12 @@ type userTSDB struct {
 	requiresOwnedSeriesUpdate atomic.String // Non-empty string means that we need to recompute "owned series" for the user. Value will be used in the log message.
 }
 
+// setDB sets u.db.
+func (u *userTSDB) setDB(db *tsdb.DB) {
+	u.db = db
+	u.hasDB.Store(true)
+}
+
 func (u *userTSDB) Appender(ctx context.Context) storage.Appender {
 	return u.db.Appender(ctx)
 }
@@ -361,7 +368,7 @@ func (u *userTSDB) PostDeletion(metrics map[chunks.HeadSeriesRef]labels.Labels)
 
 // blocksToDelete filters the input blocks and returns the blocks which are safe to be deleted from the ingester.
 func (u *userTSDB) blocksToDelete(blocks []*tsdb.Block) map[ulid.ULID]struct{} {
-	if u.db == nil {
+	if !u.hasDB.Load() {
 		return nil
 	}
 

pkg/ingester/user_tsdb_test.go

@@ -5,6 +5,7 @@ package ingester
 import (
 	"context"
 	"math"
+	"sync"
 	"testing"
 	"time"
 
@@ -103,6 +104,40 @@ func TestUserTSDB_acquireAppendLock(t *testing.T) {
 	})
 }
 
+func TestUserTSDB_blocksToDelete(t *testing.T) {
+	t.Run("blocksToDelete is concurrency safe", func(t *testing.T) {
+		var wg sync.WaitGroup
+		wg.Add(2)
+
+		setDBStarted := make(chan struct{})
+		blockDeleteStarted := make(chan struct{})
+		var userDB userTSDB
+
+		go func() {
+			defer wg.Done()
+
+			<-setDBStarted
+
+			for i := 0; i < 1000; i++ {
+				userDB.blocksToDelete(nil)
+				if i == 0 {
+					close(blockDeleteStarted)
+				}
+			}
+		}()
+
+		go func() {
+			defer wg.Done()
+
+			close(setDBStarted)
+			<-blockDeleteStarted
+			userDB.setDB(&tsdb.DB{})
+		}()
+
+		wg.Wait()
+	})
+}
+
 func TestNextForcedHeadCompactionRange(t *testing.T) {
 	const blockDuration = 10