Fix creation of user states in WAL recover

Signed-off-by: Ganesh Vernekar <cs15btech11018@iith.ac.in>
cortexproject · Jul 11, 2019 · f023ddb · f023ddb
1 parent bc25534
commit f023ddb
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 43 deletions.
diff --git a/pkg/ingester/user_state.go b/pkg/ingester/user_state.go
@@ -120,6 +120,43 @@ func (us *userStates) get(userID string) (*userState, bool) {
 	return state.(*userState), ok
 }
 
+func (us *userStates) getOrCreate(userID string) *userState {
+	rawState, ok := us.states.Load(userID)
+	if ok {
+		return rawState.(*userState)
+	}
+
+	seriesInMetric := make([]metricCounterShard, 0, metricCounterShards)
+	for i := 0; i < metricCounterShards; i++ {
+		seriesInMetric = append(seriesInMetric, metricCounterShard{
+			m: map[string]int{},
+		})
+	}
+
+	// Speculatively create a userState object and try to store it
+	// in the map.  Another goroutine may have got there before
+	// us, in which case this userState will be discarded
+	state := &userState{
+		userID:              userID,
+		limits:              us.limits,
+		fpToSeries:          newSeriesMap(),
+		fpLocker:            newFingerprintLocker(16 * 1024),
+		index:               index.New(),
+		ingestedAPISamples:  newEWMARate(0.2, us.cfg.RateUpdatePeriod),
+		ingestedRuleSamples: newEWMARate(0.2, us.cfg.RateUpdatePeriod),
+		seriesInMetric:      seriesInMetric,
+
+		memSeriesCreatedTotal: memSeriesCreatedTotal.WithLabelValues(userID),
+		memSeriesRemovedTotal: memSeriesRemovedTotal.WithLabelValues(userID),
+	}
+	state.mapper = newFPMapper(state.fpToSeries)
+	stored, ok := us.states.LoadOrStore(userID, state)
+	if !ok {
+		memUsers.Inc()
+	}
+	return stored.(*userState)
+}
+
 func (us *userStates) getViaContext(ctx context.Context) (*userState, bool, error) {
 	userID, err := user.ExtractOrgID(ctx)
 	if err != nil {
@@ -135,39 +172,7 @@ func (us *userStates) getOrCreateSeries(ctx context.Context, labels []client.Lab
 		return nil, 0, nil, fmt.Errorf("no user id")
 	}
 
-	state, ok := us.get(userID)
-	if !ok {
-
-		seriesInMetric := make([]metricCounterShard, 0, metricCounterShards)
-		for i := 0; i < metricCounterShards; i++ {
-			seriesInMetric = append(seriesInMetric, metricCounterShard{
-				m: map[string]int{},
-			})
-		}
-
-		// Speculatively create a userState object and try to store it
-		// in the map.  Another goroutine may have got there before
-		// us, in which case this userState will be discarded
-		state = &userState{
-			userID:              userID,
-			limits:              us.limits,
-			fpToSeries:          newSeriesMap(),
-			fpLocker:            newFingerprintLocker(16 * 1024),
-			index:               index.New(),
-			ingestedAPISamples:  newEWMARate(0.2, us.cfg.RateUpdatePeriod),
-			ingestedRuleSamples: newEWMARate(0.2, us.cfg.RateUpdatePeriod),
-			seriesInMetric:      seriesInMetric,
-
-			memSeriesCreatedTotal: memSeriesCreatedTotal.WithLabelValues(userID),
-			memSeriesRemovedTotal: memSeriesRemovedTotal.WithLabelValues(userID),
-		}
-		state.mapper = newFPMapper(state.fpToSeries)
-		stored, ok := us.states.LoadOrStore(userID, state)
-		if !ok {
-			memUsers.Inc()
-		}
-		state = stored.(*userState)
-	}
+	state := us.getOrCreate(userID)
 
 	fp, series, err := state.getSeries(labels, record)
 	return state, fp, series, err

diff --git a/pkg/ingester/wal.go b/pkg/ingester/wal.go
@@ -2,7 +2,6 @@ package ingester
 
 import (
 	"flag"
-	"fmt"
 	"path"
 	"sync"
 	"time"
@@ -220,7 +219,7 @@ func (w *wrapper) truncateSamples() error {
 	return nil
 }
 
-func (w *wrapper) recover(ctx context.Context) error {
+func (w *wrapper) recover(ctx context.Context) (err error) {
 	// Use a local userStates, so we don't need to worry about locking.
 	userStates := newUserStates(w.ingester.limits, w.ingester.cfg)
 
@@ -233,10 +232,7 @@ func (w *wrapper) recover(ctx context.Context) error {
 			return err
 		}
 
-		state, ok := userStates.get(walSeries.UserId)
-		if !ok {
-			return fmt.Errorf("user state not found for userid=%s", walSeries.UserId)
-		}
+		state := userStates.getOrCreate(walSeries.UserId)
 
 		la = la[:0]
 		for _, l := range walSeries.Labels {
@@ -258,10 +254,7 @@ func (w *wrapper) recover(ctx context.Context) error {
 	if err := w.recoverRecords("samples", &Record{}, func(msg proto.Message) error {
 		record := msg.(*Record)
 
-		state, ok := userStates.get(record.UserId)
-		if !ok {
-			return fmt.Errorf("user state not found for userid=%s", record.UserId)
-		}
+		state := userStates.getOrCreate(record.UserId)
 
 		for _, labels := range record.Labels {
 			_, ok := state.fpToSeries.get(model.Fingerprint(labels.Fingerprint))