ingester_v2: open TSDB on startup, limit the number of concurrently

opening

Signed-off-by: Thor <[email protected]>

Author: Thor

CHANGELOG.md

@@ -23,6 +23,8 @@
 * [FEATURE] EXPERIMENTAL: Added `/series` API endpoint support with TSDB blocks storage. #1830
 * [FEATURE] Added "multi" KV store that can interact with two other KV stores, primary one for all reads and writes, and secondary one, which only receives writes. Primary/secondary store can be modified in runtime via runtime-config mechanism (previously "overrides"). #1749
 * [ENHANCEMENT] Added `password` and `enable_tls` options to redis cache configuration. Enables usage of Microsoft Azure Cache for Redis service.
+* [ENHANCEMENT] Experimental TSDB: Open existing TSDB on startup to prevent ingester from becoming ready before it can accept writes. #1917
+  * `--experimental.tsdb.max-tsdb-opening-concurrency-on-startup`
 * [BUGFIX] Fixed unnecessary CAS operations done by the HA tracker when the jitter is enabled. #1861
 * [BUGFIX] Fixed #1904 ingesters getting stuck in a LEAVING state after coming up from an ungraceful exit. #1921
 * [BUGFIX] TSDB: Fixed handling of out of order/bound samples in ingesters with the experimental TSDB blocks storage. #1864

pkg/ingester/ingester_v2.go

@@ -2,7 +2,10 @@ package ingester
 
 import (
 	"fmt"
+	"io"
 	"net/http"
+	"os"
+	"path/filepath"
 	"sync"
 	"time"
 
@@ -14,6 +17,7 @@ import (
 	"github.com/go-kit/kit/log/level"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/prometheus/pkg/gate"
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/tsdb"
 	"github.com/thanos-io/thanos/pkg/block/metadata"
@@ -73,6 +77,11 @@ func NewV2(cfg Config, clientConfig client.Config, limits *validation.Overrides,
 	i.limiter = NewSeriesLimiter(limits, i.lifecycler, cfg.LifecyclerConfig.RingConfig.ReplicationFactor, cfg.ShardByAllLabels)
 	i.userStates = newUserStates(i.limiter, cfg)
 
+	// Scan and open TSDB's that already exist on disk
+	if err := i.openExistingTSDB(context.Background()); err != nil {
+		return nil, err
+	}
+
 	// Now that user states have been created, we can start the lifecycler
 	i.lifecycler.Start()
 
@@ -382,11 +391,23 @@ func (i *Ingester) getOrCreateTSDB(userID string, force bool) (*tsdb.DB, error)
 		return nil, fmt.Errorf(errTSDBCreateIncompatibleState, ingesterState)
 	}
 
+	// Create the database and a shipper for a user
+	db, err := i.createTSDB(userID)
+	if err != nil {
+		return nil, err
+	}
+
+	// Add the db to list of user databases
+	i.TSDBState.dbs[userID] = db
+	return db, nil
+}
+
+// createTSDB creates a TSDB for a given userID, and returns the created db.
+func (i *Ingester) createTSDB(userID string) (*tsdb.DB, error) {
 	udir := i.cfg.TSDBConfig.BlocksDir(userID)
 
 	// Create a new user database
-	var err error
-	db, err = tsdb.Open(udir, util.Logger, nil, &tsdb.Options{
+	db, err := tsdb.Open(udir, util.Logger, nil, &tsdb.Options{
 		RetentionDuration: uint64(i.cfg.TSDBConfig.Retention / time.Millisecond),
 		BlockRanges:       i.cfg.TSDBConfig.BlockRanges.ToMillisecondRanges(),
 		NoLockfile:        true,
@@ -422,8 +443,6 @@ func (i *Ingester) getOrCreateTSDB(userID string, force bool) (*tsdb.DB, error)
 		}()
 	}
 
-	i.TSDBState.dbs[userID] = db
-
 	return db, nil
 }
 
@@ -459,3 +478,73 @@ func (i *Ingester) closeAllTSDB() {
 	i.userStatesMtx.Unlock()
 	wg.Wait()
 }
+
+// openExistingTSDB walks the user tsdb dir, and opens a tsdb for each user. This may start a WAL replay, so we limit the number of
+// concurrently opening TSDB.
+func (i *Ingester) openExistingTSDB(ctx context.Context) error {
+	level.Info(util.Logger).Log("msg", "opening existing TSDBs")
+	wg := &sync.WaitGroup{}
+	openGate := gate.New(i.cfg.TSDBConfig.MaxTSDBOpeningConcurrencyOnStartup)
+
+	err := filepath.Walk(i.cfg.TSDBConfig.Dir, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return filepath.SkipDir
+		}
+
+		// Skip root dir and all other files
+		if path == i.cfg.TSDBConfig.Dir || !info.IsDir() {
+			return nil
+		}
+
+		// Top level directories are assumed to be user TSDBs
+		userID := info.Name()
+		f, err := os.Open(path)
+		if err != nil {
+			level.Error(util.Logger).Log("msg", "unable to open user TSDB dir", "err", err, "user", userID, "path", path)
+			return filepath.SkipDir
+		}
+		defer f.Close()
+
+		// If the dir is empty skip it
+		if _, err := f.Readdirnames(1); err != nil {
+			if err != io.EOF {
+				level.Error(util.Logger).Log("msg", "unable to read TSDB dir", "err", err, "user", userID, "path", path)
+			}
+
+			return filepath.SkipDir
+		}
+
+		// Limit the number of TSDB's opening concurrently. Start blocks until there's a free spot available or the context is cancelled.
+		if err := openGate.Start(ctx); err != nil {
+			return err
+		}
+
+		wg.Add(1)
+		go func(userID string) {
+			defer wg.Done()
+			defer openGate.Done()
+			db, err := i.createTSDB(userID)
+			if err != nil {
+				level.Error(util.Logger).Log("msg", "unable to open user TSDB", "err", err, "user", userID)
+				return
+			}
+
+			// Add the database to the map of user databases
+			i.userStatesMtx.Lock()
+			i.TSDBState.dbs[userID] = db
+			i.userStatesMtx.Unlock()
+
+		}(userID)
+
+		return filepath.SkipDir // Don't descend into directories
+	})
+
+	// Wait for all opening routines to finish
+	wg.Wait()
+	if err != nil {
+		level.Error(util.Logger).Log("msg", "error while opening existing TSDBs")
+	} else {
+		level.Info(util.Logger).Log("msg", "successfully opened existing TSDBs")
+	}
+	return err
+}

pkg/ingester/ingester_v2_test.go

@@ -6,6 +6,7 @@ import (
 	"math"
 	"net/http"
 	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 	"time"
@@ -745,3 +746,80 @@ func newIngesterMockWithTSDBStorage(ingesterCfg Config, registerer prometheus.Re
 
 	return ingester, cleanup, nil
 }
+
+func TestIngester_v2LoadTSDBOnStartup(t *testing.T) {
+	t.Parallel()
+
+	tests := map[string]struct {
+		setup func(*testing.T, string)
+		check func(*testing.T, *Ingester)
+	}{
+		"empty user dir": {
+			setup: func(t *testing.T, dir string) {
+				require.NoError(t, os.Mkdir(filepath.Join(dir, "user0"), 0700))
+			},
+			check: func(t *testing.T, i *Ingester) {
+				require.Empty(t, i.getTSDB("user0"), "tsdb created for empty user dir")
+			},
+		},
+		"empty tsdbs": {
+			setup: func(t *testing.T, dir string) {},
+			check: func(t *testing.T, i *Ingester) {
+				require.Zero(t, len(i.TSDBState.dbs), "user tsdb's were created on empty dir")
+			},
+		},
+		"missing tsdb dir": {
+			setup: func(t *testing.T, dir string) {
+				require.NoError(t, os.Remove(dir))
+			},
+			check: func(t *testing.T, i *Ingester) {
+				require.Zero(t, len(i.TSDBState.dbs), "user tsdb's were created on missing dir")
+			},
+		},
+		"populated user dirs with unpopulated": {
+			setup: func(t *testing.T, dir string) {
+				require.NoError(t, os.MkdirAll(filepath.Join(dir, "user0", "dummy"), 0700))
+				require.NoError(t, os.MkdirAll(filepath.Join(dir, "user1", "dummy"), 0700))
+				require.NoError(t, os.Mkdir(filepath.Join(dir, "user2"), 0700))
+			},
+			check: func(t *testing.T, i *Ingester) {
+				require.NotNil(t, i.getTSDB("user0"), "tsdb not created for non-empty user dir")
+				require.NotNil(t, i.getTSDB("user1"), "tsdb not created for non-empty user dir")
+				require.Empty(t, i.getTSDB("user2"), "tsdb created for empty user dir")
+			},
+		},
+	}
+
+	for name, test := range tests {
+		testName := name
+		testData := test
+		t.Run(testName, func(t *testing.T) {
+			clientCfg := defaultClientTestConfig()
+			limits := defaultLimitsTestConfig()
+
+			overrides, err := validation.NewOverrides(limits, nil)
+			require.NoError(t, err)
+
+			// Create a temporary directory for TSDB
+			tempDir, err := ioutil.TempDir("", "tsdb")
+			require.NoError(t, err)
+			defer os.RemoveAll(tempDir)
+
+			ingesterCfg := defaultIngesterTestConfig()
+			ingesterCfg.TSDBEnabled = true
+			ingesterCfg.TSDBConfig.Dir = tempDir
+			ingesterCfg.TSDBConfig.Backend = "s3"
+			ingesterCfg.TSDBConfig.S3.Endpoint = "localhost"
+
+			// setup the tsdbs dir
+			testData.setup(t, tempDir)
+
+			ingester, err := NewV2(ingesterCfg, clientCfg, overrides, nil)
+			require.NoError(t, err)
+
+			defer ingester.Shutdown()
+
+			testData.check(t, ingester)
+		})
+	}
+}

pkg/storage/tsdb/config.go

@@ -37,6 +37,9 @@ type Config struct {
 	Backend      string            `yaml:"backend"`
 	BucketStore  BucketStoreConfig `yaml:"bucket_store"`
 
+	// MaxTSDBOpeningConcurrencyOnStartup limits the number of concurrently opening TSDB's during startup
+	MaxTSDBOpeningConcurrencyOnStartup int `yaml:"max_tsdb_opening_concurrency_on_startup"`
+
 	// Backends
 	S3  s3.Config  `yaml:"s3"`
 	GCS gcs.Config `yaml:"gcs"`
@@ -87,6 +90,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.Retention, "experimental.tsdb.retention-period", 6*time.Hour, "TSDB block retention")
 	f.DurationVar(&cfg.ShipInterval, "experimental.tsdb.ship-interval", 30*time.Second, "the frequency at which tsdb blocks are scanned for shipping. 0 means shipping is disabled.")
 	f.StringVar(&cfg.Backend, "experimental.tsdb.backend", "s3", "TSDB storage backend to use")
+	f.IntVar(&cfg.MaxTSDBOpeningConcurrencyOnStartup, "experimental.tsdb.max-tsdb-opening-concurrency-on-startup", 10, "limit the number of concurrently opening TSDB's on startup")
 }
 
 // Validate the config