Introduce appender for backfilling to avoid 1 appender per sample

Signed-off-by: Ganesh Vernekar <[email protected]>

Author: codesome

pkg/ingester/ingester_v2.go

@@ -20,6 +20,7 @@ import (
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
+	tsdb_errors "github.com/prometheus/prometheus/tsdb/errors"
 	"github.com/thanos-io/thanos/pkg/block/metadata"
 	"github.com/thanos-io/thanos/pkg/objstore"
 	"github.com/thanos-io/thanos/pkg/shipper"
@@ -158,6 +159,7 @@ func newTSDBState(bucketClient objstore.Bucket, registerer prometheus.Registerer
 		tsdbMetrics:         newTSDBMetrics(registerer),
 		forceCompactTrigger: make(chan chan<- struct{}),
 		shipTrigger:         make(chan chan<- struct{}),
+		backfillDBs:         newBackfillTSDBs(),
 
 		compactionsTriggered: promauto.With(registerer).NewCounter(prometheus.CounterOpts{
 			Name: "cortex_ingester_tsdb_compactions_triggered_total",
@@ -426,6 +428,7 @@ func (i *Ingester) v2Push(ctx context.Context, req *client.WriteRequest) (*clien
 
 	// Walk the samples, appending them to the users database
 	app := db.Appender()
+	var backfillApp *backfillAppender
 	for _, ts := range req.Timeseries {
 		// Check if we already have a cached reference for this series. Be aware
 		// that even if we have a reference it's not guaranteed to be still valid.
@@ -470,7 +473,10 @@ func (i *Ingester) v2Push(ctx context.Context, req *client.WriteRequest) (*clien
 			if cause == storage.ErrOutOfBounds &&
 				i.cfg.TSDBConfig.BackfillLimit > 0 &&
 				s.TimestampMs > db.Head().MaxTime()-i.cfg.TSDBConfig.BackfillLimit.Milliseconds() {
-				err := i.v2BackfillPush(userID, ts.Labels, s)
+				if backfillApp == nil {
+					backfillApp = i.newBackfillAppender(userID)
+				}
+				err := backfillApp.add(ts.Labels, s)
 				if err == nil {
 					succeededSamplesCount++
 					continue
@@ -512,8 +518,13 @@ func (i *Ingester) v2Push(ctx context.Context, req *client.WriteRequest) (*clien
 			}
 
 			// The error looks an issue on our side, so we should rollback
-			if rollbackErr := app.Rollback(); rollbackErr != nil {
-				level.Warn(util.Logger).Log("msg", "failed to rollback on error", "user", userID, "err", rollbackErr)
+			var merr tsdb_errors.MultiError
+			merr.Add(errors.Wrap(app.Rollback(), "main appender"))
+			if backfillApp != nil {
+				merr.Add(errors.Wrap(backfillApp.rollback(), "backfill appender"))
+			}
+			if merr.Err() != nil {
+				level.Warn(util.Logger).Log("msg", "failed to rollback on error", "user", userID, "err", merr.Err())
 			}
 
 			return nil, wrapWithUser(err, userID)
@@ -524,8 +535,13 @@ func (i *Ingester) v2Push(ctx context.Context, req *client.WriteRequest) (*clien
 	i.TSDBState.appenderAddDuration.Observe(time.Since(startAppend).Seconds())
 
 	startCommit := time.Now()
-	if err := app.Commit(); err != nil {
-		return nil, wrapWithUser(err, userID)
+	var merr tsdb_errors.MultiError
+	merr.Add(errors.Wrap(app.Commit(), "main appender"))
+	if backfillApp != nil {
+		merr.Add(errors.Wrap(backfillApp.commit(), "backfill appender"))
+	}
+	if merr.Err() != nil {
+		return nil, wrapWithUser(merr.Err(), userID)
 	}
 	i.TSDBState.appenderCommitDuration.Observe(time.Since(startCommit).Seconds())
 
@@ -989,7 +1005,7 @@ func (i *Ingester) createNewTSDB(userID, dbDir string, minBlockDuration, maxBloc
 		WALCompression:    i.cfg.TSDBConfig.WALCompressionEnabled,
 	})
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to open TSDB: %s", udir)
+		return nil, errors.Wrapf(err, "failed to open TSDB: %s", dbDir)
 	}
 	db.DisableCompactions() // we will compact on our own schedule
 
@@ -999,7 +1015,7 @@ func (i *Ingester) createNewTSDB(userID, dbDir string, minBlockDuration, maxBloc
 	level.Info(userLogger).Log("msg", "Running compaction after WAL replay")
 	err = db.Compact()
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to compact TSDB: %s", udir)
+		return nil, errors.Wrapf(err, "failed to compact TSDB: %s", dbDir)
 	}
 
 	userDB.DB = db

pkg/ingester/ingester_v2_backfill.go

@@ -19,22 +19,53 @@ import (
 	"github.com/prometheus/prometheus/pkg/gate"
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/storage"
+	tsdb_errors "github.com/prometheus/prometheus/tsdb/errors"
 
 	"github.com/cortexproject/cortex/pkg/ingester/client"
 	"github.com/cortexproject/cortex/pkg/util"
 )
 
-func (i *Ingester) v2BackfillPush(userID string, la []client.LabelAdapter, s client.Sample) error {
-	bucket, err := i.getOrCreateBackfillTSDB(userID, s.TimestampMs)
-	if err != nil {
-		return err
+// backfillAppender is an appender to ingest old data.
+// This _does not_ implement storage.Appender interface.
+// The methods of this appender should not be called concurrently.
+type backfillAppender struct {
+	userID    string
+	ingester  *Ingester
+	buckets   []*tsdbBucket
+	appenders map[int]storage.Appender
+}
+
+func (i *Ingester) newBackfillAppender(userID string) *backfillAppender {
+	return &backfillAppender{
+		userID:    userID,
+		ingester:  i,
+		buckets:   i.TSDBState.backfillDBs.getBucketsForUser(userID),
+		appenders: make(map[int]storage.Appender),
+	}
+}
+
+func (a *backfillAppender) add(la []client.LabelAdapter, s client.Sample) (err error) {
+	bucket := getBucketForTimestamp(s.TimestampMs, a.buckets)
+	if bucket == nil {
+		var userBuckets []*tsdbBucket
+		bucket, userBuckets, err = a.ingester.getOrCreateBackfillTSDB(a.userID, s.TimestampMs)
+		if err != nil {
+			return err
+		}
+		a.buckets = userBuckets
 	}
 
-	startAppend := time.Now()
 	db := bucket.db
-	cachedRef, cachedRefExists := db.refCache.Ref(startAppend, client.FromLabelAdaptersToLabels(la))
+	var app storage.Appender
+	if ap, ok := a.appenders[bucket.id]; ok {
+		app = ap
+	} else {
+		app = db.Appender()
+		a.appenders[bucket.id] = app
+	}
 
-	app := db.Appender()
+	startAppend := time.Now()
+	cachedRef, cachedRefExists := db.refCache.Ref(startAppend, client.FromLabelAdaptersToLabels(la))
 	// If the cached reference exists, we try to use it.
 	if cachedRefExists {
 		err = app.AddFast(cachedRef, s.TimestampMs, s.Value)
@@ -46,46 +77,66 @@ func (i *Ingester) v2BackfillPush(userID string, la []client.LabelAdapter, s cli
 
 	// If the cached reference doesn't exist, we (re)try without using the reference.
 	if !cachedRefExists {
-		var ref uint64
-
 		// Copy the label set because both TSDB and the cache may retain it.
 		copiedLabels := client.FromLabelAdaptersToLabelsWithCopy(la)
-
-		if ref, err = app.Add(copiedLabels, s.TimestampMs, s.Value); err == nil {
+		if ref, err := app.Add(copiedLabels, s.TimestampMs, s.Value); err == nil {
 			db.refCache.SetRef(startAppend, copiedLabels, ref)
-			cachedRef = ref
-			cachedRefExists = true
 		}
 	}
 
-	if err == nil {
-		err = app.Commit()
+	return err
+}
+
+func (a *backfillAppender) commit() error {
+	var merr tsdb_errors.MultiError
+	for _, app := range a.appenders {
+		merr.Add(app.Commit())
 	}
+	return merr.Err()
+}
 
-	return err
+func (a *backfillAppender) rollback() error {
+	var merr tsdb_errors.MultiError
+	for _, app := range a.appenders {
+		merr.Add(app.Rollback())
+	}
+	return merr.Err()
 }
 
-func (i *Ingester) getOrCreateBackfillTSDB(userID string, ts int64) (*tsdbBucket, error) {
-	userBuckets := i.TSDBState.backfillDBs.getBucketsForUser(userID)
+func getBucketForTimestamp(ts int64, userBuckets []*tsdbBucket) *tsdbBucket {
+	// As the number of buckets will be small, we are iterating instead of binary search.
+	for _, b := range userBuckets {
+		if ts >= b.bucketStart && ts < b.bucketEnd {
+			return b
+		}
+	}
+	return nil
+}
+
+func (i *Ingester) getOrCreateBackfillTSDB(userID string, ts int64) (*tsdbBucket, []*tsdbBucket, error) {
+	i.TSDBState.backfillDBs.tsdbsMtx.Lock()
+	defer i.TSDBState.backfillDBs.tsdbsMtx.Unlock()
+
+	userBuckets := i.TSDBState.backfillDBs.tsdbs[userID]
 
 	start, end := getBucketRangesForTimestamp(ts, 1)
 
-	var bucket *tsdbBucket
 	insertIdx := len(userBuckets)
 	for idx, b := range userBuckets {
 		if ts >= b.bucketStart && ts < b.bucketEnd {
-			bucket = b
-			break
+			return b, userBuckets, nil
 		}
 
-		// Existing: |-----------|
-		//      New:       |------------|
+		//   Existing: |-----------|
+		//        New:       |------------|
+		// Changed to:             |------| (no overlaps)
 		if b.bucketStart < start && start < b.bucketEnd {
 			start = b.bucketEnd
 		}
 
-		// Existing:         |-----------|
-		//      New:  |------------|
+		//   Existing:         |-----------|
+		//        New:  |------------|
+		// Changed to:  |------| (no overlaps)
 		if end > b.bucketStart && end < b.bucketEnd {
 			end = b.bucketStart
 			insertIdx = idx
@@ -98,26 +149,26 @@ func (i *Ingester) getOrCreateBackfillTSDB(userID string, ts int64) (*tsdbBucket
 		}
 	}
 
-	if bucket == nil {
-		tsdb, err := i.createNewTSDB(
-			userID, filepath.Join(userID, getBucketName(start, end)),
-			(end-start)*2, (end-start)*2, prometheus.NewRegistry(),
-		)
-		if err != nil {
-			return nil, err
-		}
-		bucket = &tsdbBucket{
-			db:          tsdb,
-			bucketStart: start,
-			bucketEnd:   end,
-		}
-		userBuckets = append(userBuckets[:insertIdx], append([]*tsdbBucket{bucket}, userBuckets[insertIdx:]...)...)
-		i.TSDBState.backfillDBs.tsdbsMtx.Lock()
-		i.TSDBState.backfillDBs.tsdbs[userID] = userBuckets
-		i.TSDBState.backfillDBs.tsdbsMtx.Unlock()
+	tsdb, err := i.createNewTSDB(
+		userID, filepath.Join(userID, getBucketName(start, end)),
+		(end-start)*2, (end-start)*2, prometheus.NewRegistry(),
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+	bucket := &tsdbBucket{
+		db:          tsdb,
+		bucketStart: start,
+		bucketEnd:   end,
+	}
+	if len(userBuckets) > 0 {
+		bucket.id = userBuckets[len(userBuckets)-1].id + 1
 	}
+	userBuckets = append(userBuckets[:insertIdx], append([]*tsdbBucket{bucket}, userBuckets[insertIdx:]...)...)
 
-	return bucket, nil
+	i.TSDBState.backfillDBs.tsdbs[userID] = userBuckets
+
+	return bucket, userBuckets, nil
 }
 
 func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
@@ -127,6 +178,9 @@ func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
 
 	users, err := ioutil.ReadDir(i.cfg.TSDBConfig.BackfillDir)
 	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
 		return err
 	}
 
@@ -145,7 +199,7 @@ func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
 			continue
 		}
 
-		for _, bucketName := range bucketNames {
+		for bucketID, bucketName := range bucketNames {
 			if bucketName.IsDir() {
 				continue
 			}
@@ -174,21 +228,26 @@ func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
 			}
 
 			wg.Add(1)
-			go func(userID, bucketName string) {
+			go func(bucketID int, userID, bucketName string) {
 				defer wg.Done()
 				defer openGate.Done()
 				defer func(ts time.Time) {
 					i.TSDBState.walReplayTime.Observe(time.Since(ts).Seconds())
 				}(time.Now())
 
 				start, end, err := getBucketRangesForBucketName(bucketName)
+				if err != nil {
+					level.Error(util.Logger).Log("msg", "unable to get bucket range", "err", err, "user", userID, "bucketName", bucketName)
+					return
+				}
 				db, err := i.createNewTSDB(userID, filepath.Join(userID, bucketName), (end-start)*2, (end-start)*2, prometheus.NewRegistry())
 				if err != nil {
 					level.Error(util.Logger).Log("msg", "unable to open user backfill TSDB", "err", err, "user", userID)
 					return
 				}
 
 				bucket := &tsdbBucket{
+					id:          bucketID,
 					db:          db,
 					bucketStart: start,
 					bucketEnd:   end,
@@ -198,7 +257,7 @@ func (i *Ingester) openExistingBackfillTSDB(ctx context.Context) error {
 				// Append at the end, we will sort it at the end.
 				i.TSDBState.backfillDBs.tsdbs[userID] = append(i.TSDBState.backfillDBs.tsdbs[userID], bucket)
 				i.TSDBState.backfillDBs.tsdbsMtx.Unlock()
-			}(userID, bucketName.Name())
+			}(bucketID, userID, bucketName.Name())
 		}
 
 		if runErr != nil {
@@ -278,6 +337,7 @@ func (i *Ingester) backfillSelect(ctx context.Context, userID string, from, thro
 	select {
 	case err := <-errC:
 		return nil, err
+	default:
 	}
 
 	return result, nil
@@ -286,7 +346,8 @@ func (i *Ingester) backfillSelect(ctx context.Context, userID string, from, thro
 // Assumes 1h bucket range for . TODO(codesome): protect stuff with locks.
 type backfillTSDBs struct {
 	tsdbsMtx sync.RWMutex
-	tsdbs    map[string][]*tsdbBucket
+	// TODO(codesome): have more granular locks.
+	tsdbs map[string][]*tsdbBucket
 }
 
 func newBackfillTSDBs() *backfillTSDBs {
@@ -302,6 +363,7 @@ func (b *backfillTSDBs) getBucketsForUser(userID string) []*tsdbBucket {
 }
 
 type tsdbBucket struct {
+	id                     int // This is any number but should be unique among all buckets of a user.
 	db                     *userTSDB
 	bucketStart, bucketEnd int64
 }