Revert "Merge pull request #971 from grafana/batch-index-lookups"

pkg/chunk/aws/storage_client.go

@@ -29,7 +29,6 @@ import (
 	"github.com/weaveworks/common/instrument"
 	"github.com/weaveworks/common/user"
 	"github.com/weaveworks/cortex/pkg/chunk"
-	chunk_util "github.com/weaveworks/cortex/pkg/chunk/util"
 	"github.com/weaveworks/cortex/pkg/util"
 )
 
@@ -289,11 +288,7 @@ func (a storageClient) BatchWrite(ctx context.Context, input chunk.WriteBatch) e
 	return backoff.Err()
 }
 
-func (a storageClient) QueryPages(ctx context.Context, queries []chunk.IndexQuery, callback func(chunk.IndexQuery, chunk.ReadBatch) bool) error {
-	return chunk_util.DoParallelQueries(ctx, a.query, queries, callback)
-}
-
-func (a storageClient) query(ctx context.Context, query chunk.IndexQuery, callback func(result chunk.ReadBatch) (shouldContinue bool)) error {
+func (a storageClient) QueryPages(ctx context.Context, query chunk.IndexQuery, callback func(result chunk.ReadBatch) (shouldContinue bool)) error {
 	sp, ctx := ot.StartSpanFromContext(ctx, "QueryPages", ot.Tag{Key: "tableName", Value: query.TableName}, ot.Tag{Key: "hashValue", Value: query.HashValue})
 	defer sp.Finish()
 
@@ -363,7 +358,7 @@ func (a storageClient) query(ctx context.Context, query chunk.IndexQuery, callba
 	return nil
 }
 
-func (a storageClient) queryPage(ctx context.Context, input *dynamodb.QueryInput, page dynamoDBRequest) (*dynamoDBReadResponse, error) {
+func (a storageClient) queryPage(ctx context.Context, input *dynamodb.QueryInput, page dynamoDBRequest) (dynamoDBReadResponse, error) {
 	backoff := util.NewBackoff(ctx, a.cfg.backoffConfig)
 	defer func() {
 		dynamoQueryRetryCount.WithLabelValues("queryPage").Observe(float64(backoff.NumRetries()))
@@ -393,10 +388,7 @@ func (a storageClient) queryPage(ctx context.Context, input *dynamodb.QueryInput
 		}
 
 		queryOutput := page.Data().(*dynamodb.QueryOutput)
-		return &dynamoDBReadResponse{
-			i:     -1,
-			items: queryOutput.Items,
-		}, nil
+		return dynamoDBReadResponse(queryOutput.Items), nil
 	}
 	return nil, fmt.Errorf("QueryPage error: %s for table %v, last error %v", backoff.Err(), *input.TableName, err)
 }
@@ -780,22 +772,18 @@ func (a storageClient) putS3Chunk(ctx context.Context, key string, buf []byte) e
 }
 
 // Slice of values returned; map key is attribute name
-type dynamoDBReadResponse struct {
-	i     int
-	items []map[string]*dynamodb.AttributeValue
-}
+type dynamoDBReadResponse []map[string]*dynamodb.AttributeValue
 
-func (b *dynamoDBReadResponse) Next() bool {
-	b.i++
-	return b.i < len(b.items)
+func (b dynamoDBReadResponse) Len() int {
+	return len(b)
 }
 
-func (b *dynamoDBReadResponse) RangeValue() []byte {
-	return b.items[b.i][rangeKey].B
+func (b dynamoDBReadResponse) RangeValue(i int) []byte {
+	return b[i][rangeKey].B
 }
 
-func (b *dynamoDBReadResponse) Value() []byte {
-	chunkValue, ok := b.items[b.i][valueKey]
+func (b dynamoDBReadResponse) Value(i int) []byte {
+	chunkValue, ok := b[i][valueKey]
 	if !ok {
 		return nil
 	}

pkg/chunk/cassandra/storage_client.go

@@ -11,7 +11,6 @@ import (
 	"github.com/pkg/errors"
 
 	"github.com/weaveworks/cortex/pkg/chunk"
-	"github.com/weaveworks/cortex/pkg/chunk/util"
 )
 
 const (
@@ -173,11 +172,7 @@ func (s *storageClient) BatchWrite(ctx context.Context, batch chunk.WriteBatch)
 	return nil
 }
 
-func (s *storageClient) QueryPages(ctx context.Context, queries []chunk.IndexQuery, callback func(chunk.IndexQuery, chunk.ReadBatch) bool) error {
-	return util.DoParallelQueries(ctx, s.query, queries, callback)
-}
-
-func (s *storageClient) query(ctx context.Context, query chunk.IndexQuery, callback func(result chunk.ReadBatch) (shouldContinue bool)) error {
+func (s *storageClient) QueryPages(ctx context.Context, query chunk.IndexQuery, callback func(result chunk.ReadBatch) (shouldContinue bool)) error {
 	var q *gocql.Query
 
 	switch {
@@ -210,7 +205,7 @@ func (s *storageClient) query(ctx context.Context, query chunk.IndexQuery, callb
 	defer iter.Close()
 	scanner := iter.Scanner()
 	for scanner.Next() {
-		b := &readBatch{}
+		var b readBatch
 		if err := scanner.Scan(&b.rangeValue, &b.value); err != nil {
 			return errors.WithStack(err)
 		}
@@ -223,26 +218,27 @@ func (s *storageClient) query(ctx context.Context, query chunk.IndexQuery, callb
 
 // readBatch represents a batch of rows read from Cassandra.
 type readBatch struct {
-	consumed   bool
 	rangeValue []byte
 	value      []byte
 }
 
 // Len implements chunk.ReadBatch; in Cassandra we 'stream' results back
 // one-by-one, so this always returns 1.
-func (b *readBatch) Next() bool {
-	if b.consumed {
-		return false
-	}
-	b.consumed = true
-	return true
+func (readBatch) Len() int {
+	return 1
 }
 
-func (b *readBatch) RangeValue() []byte {
+func (b readBatch) RangeValue(index int) []byte {
+	if index != 0 {
+		panic("index != 0")
+	}
 	return b.rangeValue
 }
 
-func (b *readBatch) Value() []byte {
+func (b readBatch) Value(index int) []byte {
+	if index != 0 {
+		panic("index != 0")
+	}
 	return b.value
 }
 

pkg/chunk/chunk_store.go

@@ -347,22 +347,53 @@ func (c *store) lookupChunksByMetricName(ctx context.Context, from, through mode
 }
 
 func (c *store) lookupEntriesByQueries(ctx context.Context, queries []IndexQuery) ([]IndexEntry, error) {
+	incomingEntries := make(chan []IndexEntry)
+	incomingErrors := make(chan error)
+	for _, query := range queries {
+		go func(query IndexQuery) {
+			entries, err := c.lookupEntriesByQuery(ctx, query)
+			if err != nil {
+				incomingErrors <- err
+			} else {
+				incomingEntries <- entries
+			}
+		}(query)
+	}
+
+	// Combine the results into one slice
+	var entries []IndexEntry
+	var lastErr error
+	for i := 0; i < len(queries); i++ {
+		select {
+		case incoming := <-incomingEntries:
+			entries = append(entries, incoming...)
+		case err := <-incomingErrors:
+			lastErr = err
+		}
+	}
+
+	return entries, lastErr
+}
+
+func (c *store) lookupEntriesByQuery(ctx context.Context, query IndexQuery) ([]IndexEntry, error) {
 	var entries []IndexEntry
-	err := c.storage.QueryPages(ctx, queries, func(query IndexQuery, resp ReadBatch) bool {
-		for resp.Next() {
+
+	if err := c.storage.QueryPages(ctx, query, func(resp ReadBatch) (shouldContinue bool) {
+		for i := 0; i < resp.Len(); i++ {
 			entries = append(entries, IndexEntry{
 				TableName:  query.TableName,
 				HashValue:  query.HashValue,
-				RangeValue: resp.RangeValue(),
-				Value:      resp.Value(),
+				RangeValue: resp.RangeValue(i),
+				Value:      resp.Value(i),
 			})
 		}
 		return true
-	})
-	if err != nil {
+	}); err != nil {
 		level.Error(util.WithContext(ctx, util.Logger)).Log("msg", "error querying storage", "err", err)
+		return nil, err
 	}
-	return entries, err
+
+	return entries, nil
 }
 
 func (c *store) parseIndexEntries(ctx context.Context, entries []IndexEntry, matcher *labels.Matcher) ([]string, error) {