Merge pull request #1389 from cortexproject/combine-dynamo-writes

Send chunks writes to DynamoDB in the same request as index entries

Author: bboreham

pkg/chunk/aws/dynamodb_storage_client.go

@@ -616,28 +616,47 @@ func processChunkResponse(response *dynamodb.BatchGetItemOutput, chunksByKey map
 	return result, nil
 }
 
+// PutChunkAndIndex implements chunk.ObjectAndIndexClient
+// Combine both sets of writes before sending to DynamoDB, for performance
+func (a dynamoDBStorageClient) PutChunkAndIndex(ctx context.Context, c chunk.Chunk, index chunk.WriteBatch) error {
+	dynamoDBWrites, err := a.writesForChunks([]chunk.Chunk{c})
+	if err != nil {
+		return err
+	}
+	dynamoDBWrites.TakeReqs(index.(dynamoDBWriteBatch), 0)
+	return a.BatchWrite(ctx, dynamoDBWrites)
+}
+
 // PutChunks implements chunk.ObjectClient.
 func (a dynamoDBStorageClient) PutChunks(ctx context.Context, chunks []chunk.Chunk) error {
+	dynamoDBWrites, err := a.writesForChunks(chunks)
+	if err != nil {
+		return err
+	}
+	return a.BatchWrite(ctx, dynamoDBWrites)
+}
+
+func (a dynamoDBStorageClient) writesForChunks(chunks []chunk.Chunk) (dynamoDBWriteBatch, error) {
 	var (
 		dynamoDBWrites = dynamoDBWriteBatch{}
 	)
 
 	for i := range chunks {
 		buf, err := chunks[i].Encoded()
 		if err != nil {
-			return err
+			return nil, err
 		}
 		key := chunks[i].ExternalKey()
 
 		table, err := a.schemaCfg.ChunkTableFor(chunks[i].From)
 		if err != nil {
-			return err
+			return nil, err
 		}
 
 		dynamoDBWrites.Add(table, key, placeholder, buf)
 	}
 
-	return a.BatchWrite(ctx, dynamoDBWrites)
+	return dynamoDBWrites, nil
 }
 
 // Slice of values returned; map key is attribute name

pkg/chunk/series_store.go

@@ -328,21 +328,25 @@ func (c *seriesStore) Put(ctx context.Context, chunks []Chunk) error {
 func (c *seriesStore) PutOne(ctx context.Context, from, through model.Time, chunk Chunk) error {
 	chunks := []Chunk{chunk}
 
-	err := c.storage.PutChunks(ctx, chunks)
-	if err != nil {
-		return err
-	}
-
-	c.writeBackCache(ctx, chunks)
-
 	writeReqs, keysToCache, err := c.calculateIndexEntries(from, through, chunk)
 	if err != nil {
 		return err
 	}
 
-	if err := c.index.BatchWrite(ctx, writeReqs); err != nil {
-		return err
+	if oic, ok := c.storage.(ObjectAndIndexClient); ok {
+		if err = oic.PutChunkAndIndex(ctx, chunk, writeReqs); err != nil {
+			return err
+		}
+	} else {
+		err := c.storage.PutChunks(ctx, chunks)
+		if err != nil {
+			return err
+		}
+		if err := c.index.BatchWrite(ctx, writeReqs); err != nil {
+			return err
+		}
 	}
+	c.writeBackCache(ctx, chunks)
 
 	bufs := make([][]byte, len(keysToCache))
 	c.writeDedupeCache.Store(ctx, keysToCache, bufs)

pkg/chunk/storage_client.go

@@ -22,6 +22,11 @@ type ObjectClient interface {
 	GetChunks(ctx context.Context, chunks []Chunk) ([]Chunk, error)
 }
 
+// ObjectAndIndexClient allows optimisations where the same client handles both
+type ObjectAndIndexClient interface {
+	PutChunkAndIndex(ctx context.Context, c Chunk, index WriteBatch) error
+}
+
 // WriteBatch represents a batch of writes.
 type WriteBatch interface {
 	Add(tableName, hashValue string, rangeValue []byte, value []byte)