kvstreamer: remove temporary allocations for []Result

pkg/kv/kvclient/kvstreamer/results_buffer.go

@@ -45,6 +45,8 @@ type resultsBuffer interface {
 	// get returns all the Results that the buffer can send to the client at the
 	// moment. The boolean indicates whether all expected Results have been
 	// returned. Must be called without holding the budget's mutex.
+	// TODO(yuzefovich): consider changing the interface to return a single
+	// Result object in order to avoid some allocations.
 	get(context.Context) (_ []Result, allComplete bool, _ error)
 
 	// wait blocks until there is at least one Result available to be returned
@@ -64,10 +66,26 @@ type resultsBuffer interface {
 	//                                                                       //
 	///////////////////////////////////////////////////////////////////////////
 
-	// add adds the provided Results into the buffer. If any Results are
-	// available to be returned to the client and there is a goroutine blocked
-	// in wait(), the goroutine is woken up.
-	add([]Result)
+	// Lock and Unlock expose methods on the mutex of the resultsBuffer. If the
+	// Streamer's mutex needs to be locked, then the Streamer's mutex must be
+	// acquired first.
+	Lock()
+	Unlock()
+
+	// addLocked adds the provided Result into the buffer. Note that if the
+	// Result is available to be returned to the client and there is a goroutine
+	// blocked in wait(), the goroutine is **not** woken up - doneAddingLocked()
+	// has to be called.
+	//
+	// The combination of multiple addLocked() calls followed by a single
+	// doneAddingLocked() call allows us to simulate adding many Results at
+	// once, without having to allocate a slice for that.
+	addLocked(Result)
+	// doneAddingLocked notifies the resultsBuffer that the worker goroutine
+	// added all Results it could, and the resultsBuffer checks whether any
+	// Results are available to be returned to the client. If there is a
+	// goroutine blocked in wait(), the goroutine is woken up.
+	doneAddingLocked()
 
 	///////////////////////////////////////////////////////////////////////////
 	//                                                                       //
@@ -159,11 +177,10 @@ func (b *resultsBufferBase) initLocked(isEmpty bool, numExpectedResponses int) e
 	return nil
 }
 
-func (b *resultsBufferBase) findCompleteResponses(results []Result) {
-	for i := range results {
-		if results[i].GetResp != nil || results[i].scanComplete {
-			b.numCompleteResponses++
-		}
+func (b *resultsBufferBase) checkIfCompleteLocked(r Result) {
+	b.Mutex.AssertHeld()
+	if r.GetResp != nil || r.scanComplete {
+		b.numCompleteResponses++
 	}
 }
 
@@ -248,12 +265,15 @@ func (b *outOfOrderResultsBuffer) init(_ context.Context, numExpectedResponses i
 	return nil
 }
 
-func (b *outOfOrderResultsBuffer) add(results []Result) {
-	b.Lock()
-	defer b.Unlock()
-	b.results = append(b.results, results...)
-	b.findCompleteResponses(results)
-	b.numUnreleasedResults += len(results)
+func (b *outOfOrderResultsBuffer) addLocked(r Result) {
+	b.Mutex.AssertHeld()
+	b.results = append(b.results, r)
+	b.checkIfCompleteLocked(r)
+	b.numUnreleasedResults++
+}
+
+func (b *outOfOrderResultsBuffer) doneAddingLocked() {
+	b.Mutex.AssertHeld()
 	b.signal()
 }
 
@@ -408,35 +428,32 @@ func (b *inOrderResultsBuffer) init(ctx context.Context, numExpectedResponses in
 	return nil
 }
 
-func (b *inOrderResultsBuffer) add(results []Result) {
-	b.Lock()
-	defer b.Unlock()
+func (b *inOrderResultsBuffer) addLocked(r Result) {
+	b.Mutex.AssertHeld()
 	// Note that we don't increase b.numUnreleasedResults because all these
 	// results are "buffered".
-	b.findCompleteResponses(results)
-	foundHeadOfLine := false
-	for _, r := range results {
-		if debug {
-			subRequestIdx := ""
-			if !b.singleRowLookup {
-				subRequestIdx = fmt.Sprintf(" (%d)", r.subRequestIdx)
-			}
-			fmt.Printf("adding a result for position %d%s of size %d\n", r.Position, subRequestIdx, r.memoryTok.toRelease)
-		}
-		// All the Results have already been registered with the budget, so
-		// we're keeping them in-memory.
-		heap.Push(b, inOrderBufferedResult{Result: r, onDisk: false, addEpoch: b.addCounter})
-		if r.Position == b.headOfLinePosition && r.subRequestIdx == b.headOfLineSubRequestIdx {
-			foundHeadOfLine = true
+	b.checkIfCompleteLocked(r)
+	if debug {
+		subRequestIdx := ""
+		if !b.singleRowLookup {
+			subRequestIdx = fmt.Sprintf(" (%d)", r.subRequestIdx)
 		}
+		fmt.Printf("adding a result for position %d%s of size %d\n", r.Position, subRequestIdx, r.memoryTok.toRelease)
 	}
-	if foundHeadOfLine {
+	// All the Results have already been registered with the budget, so we're
+	// keeping them in-memory.
+	heap.Push(b, inOrderBufferedResult{Result: r, onDisk: false, addEpoch: b.addCounter})
+	b.addCounter++
+}
+
+func (b *inOrderResultsBuffer) doneAddingLocked() {
+	b.Mutex.AssertHeld()
+	if len(b.buffered) > 0 && b.buffered[0].Position == b.headOfLinePosition && b.buffered[0].subRequestIdx == b.headOfLineSubRequestIdx {
 		if debug {
 			fmt.Println("found head-of-the-line")
 		}
 		b.signal()
 	}
-	b.addCounter++
 }
 
 func (b *inOrderResultsBuffer) get(ctx context.Context) ([]Result, bool, error) {

pkg/kv/kvclient/kvstreamer/results_buffer_test.go

@@ -113,13 +113,14 @@ func TestInOrderResultsBuffer(t *testing.T) {
 				break
 			}
 
+			b.Lock()
 			numToAdd := rng.Intn(len(addOrder)) + 1
-			toAdd := make([]Result, numToAdd)
-			for i := range toAdd {
-				toAdd[i] = results[addOrder[0]]
+			for i := 0; i < numToAdd; i++ {
+				b.addLocked(results[addOrder[0]])
 				addOrder = addOrder[1:]
 			}
-			b.add(toAdd)
+			b.doneAddingLocked()
+			b.Unlock()
 
 			// With 50% probability, try spilling some of the buffered results
 			// to disk.