[dbnode] Adaptive WriteBatch allocations

src/dbnode/network/server/tchannelthrift/node/service_test.go

@@ -2546,7 +2546,7 @@ func TestServiceWriteBatchRaw(t *testing.T) {
 		{"bar", time.Now().Truncate(time.Second), 42.42},
 	}
 
-	writeBatch := writes.NewWriteBatch(len(values), ident.StringID(nsID), nil)
+	writeBatch := writes.NewWriteBatch(0, ident.StringID(nsID), nil)
 	mockDB.EXPECT().
 		BatchWriter(ident.NewIDMatcher(nsID), len(values)).
 		Return(writeBatch, nil)
@@ -2600,7 +2600,7 @@ func TestServiceWriteBatchRawV2SingleNS(t *testing.T) {
 		{"bar", time.Now().Truncate(time.Second), 42.42},
 	}
 
-	writeBatch := writes.NewWriteBatch(len(values), ident.StringID(nsID), nil)
+	writeBatch := writes.NewWriteBatch(0, ident.StringID(nsID), nil)
 	mockDB.EXPECT().
 		BatchWriter(ident.NewIDMatcher(nsID), len(values)).
 		Return(writeBatch, nil)
@@ -2657,8 +2657,8 @@ func TestServiceWriteBatchRawV2MultiNS(t *testing.T) {
 			{"bar", time.Now().Truncate(time.Second), 42.42},
 		}
 
-		writeBatch1 = writes.NewWriteBatch(len(values), ident.StringID(nsID1), nil)
-		writeBatch2 = writes.NewWriteBatch(len(values), ident.StringID(nsID2), nil)
+		writeBatch1 = writes.NewWriteBatch(0, ident.StringID(nsID1), nil)
+		writeBatch2 = writes.NewWriteBatch(0, ident.StringID(nsID2), nil)
 	)
 
 	mockDB.EXPECT().
@@ -2751,7 +2751,7 @@ func TestServiceWriteBatchRawOverMaxOutstandingRequests(t *testing.T) {
 		testIsComplete       = make(chan struct{}, 0)
 		requestIsOutstanding = make(chan struct{}, 0)
 	)
-	writeBatch := writes.NewWriteBatch(len(values), ident.StringID(nsID), nil)
+	writeBatch := writes.NewWriteBatch(0, ident.StringID(nsID), nil)
 	mockDB.EXPECT().
 		BatchWriter(ident.NewIDMatcher(nsID), len(values)).
 		Do(func(nsID ident.ID, numValues int) {

src/dbnode/persist/fs/commitlog/commit_log_test.go

@@ -161,7 +161,7 @@ func testSeries(
 		UniqueIndex: uniqueIndex,
 		Namespace:   ident.StringID("testNS"),
 		ID:          ident.StringID(id),
-		EncodedTags: ts.EncodedTags(encodedTagsChecked.Bytes()),
+		EncodedTags: encodedTagsChecked.Bytes(),
 		Shard:       shard,
 	}
 }
@@ -1064,7 +1064,7 @@ func TestCommitLogBatchWriteDoesNotAddErroredOrSkippedSeries(t *testing.T) {
 		finalized++
 	}
 
-	writes := writes.NewWriteBatch(4, ident.StringID("ns"), finalizeFn)
+	writes := writes.NewWriteBatch(0, ident.StringID("ns"), finalizeFn)
 
 	testSeriesWrites := []ts.Series{
 		testSeries(t, opts, 0, "foo.bar", testTags0, 42),

src/dbnode/server/server.go

@@ -1598,14 +1598,10 @@ func withEncodingAndPoolingOptions(
 			policy.IteratorPool,
 			scope.SubScope("multi-iterator-pool")))
 
-	var writeBatchPoolInitialBatchSize *int
+	writeBatchPoolInitialBatchSize := 0
 	if policy.WriteBatchPool.InitialBatchSize != nil {
 		// Use config value if available.
-		writeBatchPoolInitialBatchSize = policy.WriteBatchPool.InitialBatchSize
-	} else {
-		// Otherwise use the default batch size that the client will use.
-		clientDefaultSize := client.DefaultWriteBatchSize
-		writeBatchPoolInitialBatchSize = &clientDefaultSize
+		writeBatchPoolInitialBatchSize = *policy.WriteBatchPool.InitialBatchSize
 	}
 
 	var writeBatchPoolMaxBatchSize *int
@@ -1623,7 +1619,10 @@ func withEncodingAndPoolingOptions(
 		// writes without allocating because these objects are very expensive to
 		// allocate.
 		commitlogQueueSize := opts.CommitLogOptions().BacklogQueueSize()
-		expectedBatchSize := *writeBatchPoolInitialBatchSize
+		expectedBatchSize := writeBatchPoolInitialBatchSize
+		if expectedBatchSize == 0 {
+			expectedBatchSize = client.DefaultWriteBatchSize
+		}
 		writeBatchPoolSize = commitlogQueueSize / expectedBatchSize
 	}
 

src/dbnode/storage/options.go

@@ -193,7 +193,7 @@ func newOptions(poolOpts pool.ObjectPoolOptions) Options {
 	bytesPool.Init()
 	seriesOpts := series.NewOptions()
 
-	writeBatchPool := writes.NewWriteBatchPool(poolOpts, nil, nil)
+	writeBatchPool := writes.NewWriteBatchPool(poolOpts, 0, nil)
 	writeBatchPool.Init()
 
 	segmentReaderPool := xio.NewSegmentReaderPool(poolOpts)

src/dbnode/ts/writes/write_batch.go

@@ -33,6 +33,13 @@ var (
 	errTagsAndEncodedTagsRequired = errors.New("tags iterator and encoded tags must be provided")
 )
 
+const (
+	// preallocateBatchCoeff is used for allocating write batches of slightly bigger
+	// capacity than needed for the current request, in order to reduce allocations on
+	// subsequent reuse of pooled write batch.
+	preallocateBatchCoeff = 1.2
+)
+
 type writeBatch struct {
 	writes       []BatchWrite
 	pendingIndex []PendingIndexInsert
@@ -50,13 +57,25 @@ type writeBatch struct {
 
 // NewWriteBatch creates a new WriteBatch.
 func NewWriteBatch(
-	batchSize int,
+	initialBatchSize int,
 	ns ident.ID,
 	finalizeFn func(WriteBatch),
 ) WriteBatch {
+	var (
+		writes       []BatchWrite
+		pendingIndex []PendingIndexInsert
+	)
+
+	if initialBatchSize > 0 {
+		writes = make([]BatchWrite, 0, initialBatchSize)
+		pendingIndex = make([]PendingIndexInsert, 0, initialBatchSize)
+		// Leaving nil slices if initialBatchSize == 0,
+		// they will be allocated when needed, based on the actual batch size.
+	}
+
 	return &writeBatch{
-		writes:       make([]BatchWrite, 0, batchSize),
-		pendingIndex: make([]PendingIndexInsert, 0, batchSize),
+		writes:       writes,
+		pendingIndex: pendingIndex,
 		ns:           ns,
 		finalizeFn:   finalizeFn,
 	}
@@ -102,14 +121,29 @@ func (b *writeBatch) Reset(
 	batchSize int,
 	ns ident.ID,
 ) {
-	var writes []BatchWrite
+	// Preallocate slightly more when not using initialBatchSize.
+	preallocateBatchCap := int(float32(batchSize) * preallocateBatchCoeff)
+
 	if batchSize > cap(b.writes) {
-		writes = make([]BatchWrite, 0, batchSize)
+		batchCap := batchSize
+		if cap(b.writes) == 0 {
+			batchCap = preallocateBatchCap
+		}
+		b.writes = make([]BatchWrite, 0, batchCap)
+	} else {
+		b.writes = b.writes[:0]
+	}
+
+	if batchSize > cap(b.pendingIndex) {
+		batchCap := batchSize
+		if cap(b.pendingIndex) == 0 {
+			batchCap = preallocateBatchCap
+		}
+		b.pendingIndex = make([]PendingIndexInsert, 0, batchCap)
 	} else {
-		writes = b.writes[:0]
+		b.pendingIndex = b.pendingIndex[:0]
 	}
 
-	b.writes = writes
 	b.ns = ns
 	b.finalizeEncodedTagsFn = nil
 	b.finalizeAnnotationFn = nil
@@ -144,14 +178,14 @@ func (b *writeBatch) PendingIndex() []PendingIndexInsert {
 	return b.pendingIndex
 }
 
-// Set the function that will be called to finalize annotations when a WriteBatch
-// is finalized, allowing the caller to pool them.
+// SetFinalizeEncodedTagsFn sets the function that will be called to finalize encodedTags
+// when a WriteBatch is finalized, allowing the caller to pool them.
 func (b *writeBatch) SetFinalizeEncodedTagsFn(f FinalizeEncodedTagsFn) {
 	b.finalizeEncodedTagsFn = f
 }
 
-// Set the function that will be called to finalize annotations when a WriteBatch
-// is finalized, allowing the caller to pool them.
+// SetFinalizeAnnotationFn sets the function that will be called to finalize annotations
+// when a WriteBatch is finalized, allowing the caller to pool them.
 func (b *writeBatch) SetFinalizeAnnotationFn(f FinalizeAnnotationFn) {
 	b.finalizeAnnotationFn = f
 }

src/dbnode/ts/writes/write_batch_pool.go

@@ -25,13 +25,10 @@ import (
 )
 
 const (
-	// defaultInitialiBatchSize determines the initial batch size that will be used when filling up the
-	// pool.
-	defaultInitialBatchSize = 1024
 	// defaultWritePoolMaxBatchSize is the default maximum size for a writeBatch that the pool
 	// will allow to remain in the pool. Any batches larger than that will be discarded to prevent
 	// excessive memory use forever in the case of an exceptionally large batch write.
-	defaultMaxBatchSize = 100000
+	defaultMaxBatchSize = 10240
 )
 
 // WriteBatchPool is a pool of WriteBatch.
@@ -44,14 +41,9 @@ type WriteBatchPool struct {
 // NewWriteBatchPool constructs a new WriteBatchPool.
 func NewWriteBatchPool(
 	opts pool.ObjectPoolOptions,
-	initialBatchSizeOverride,
+	initialBatchSize int,
 	maxBatchSizeOverride *int,
 ) *WriteBatchPool {
-	initialBatchSize := defaultInitialBatchSize
-	if initialBatchSizeOverride != nil {
-		initialBatchSize = *initialBatchSizeOverride
-	}
-
 	maxBatchSize := defaultMaxBatchSize
 	if maxBatchSizeOverride != nil {
 		maxBatchSize = *maxBatchSizeOverride