[P2P-Store] Release batch IDs on all performTransfer exit paths

mooncake-p2p-store/src/p2pstore/core.go

@@ -366,60 +366,110 @@ func (store *P2PStore) performTransfer(ctx context.Context, source uintptr, shar
 	retryCount := 0
 	maxRetryCount := max(3, shard.Count())
 	for retryCount < maxRetryCount {
-		batchID, err := store.transfer.allocateBatchID(1)
-		if err != nil {
-			return err
-		}
-
 		location := shard.GetLocation(retryCount)
 		if location == nil {
 			break
 		}
 
-		targetID, err := store.transfer.openSegment(location.SegmentName, retryCount == 0)
+		status, err := performTransferOnce(ctx, store.transfer, source, shard.Length, location, retryCount == 0)
 		if err != nil {
 			return err
 		}
 
-		request := TransferRequest{
-			Opcode:       OPCODE_READ,
-			Source:       uint64(source),
-			TargetID:     targetID,
-			TargetOffset: location.Offset,
-			Length:       shard.Length,
+		if status == STATUS_COMPLETED {
+			return nil
 		}
 
-		err = store.transfer.submitTransfer(batchID, []TransferRequest{request})
-		if err != nil {
-			return err
-		}
+		retryCount++
+	}
 
-		var status int
-		for status == STATUS_WAITING || status == STATUS_PENDING {
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			default:
-				status, _, err = store.transfer.getTransferStatus(batchID, 0)
-				if err != nil {
-					return err
-				}
+	return ErrTooManyRetries
+}
+
+// batchTransport is the subset of TransferEngine a single transfer attempt
+// needs. It lets the batch lifecycle be exercised without the native engine.
+type batchTransport interface {
+	allocateBatchID(batchSize int) (BatchID, error)
+	openSegment(segmentName string, useCache bool) (int64, error)
+	submitTransfer(batchID BatchID, requests []TransferRequest) error
+	getTransferStatus(batchID BatchID, taskID int) (int, uint64, error)
+	freeBatchID(batchID BatchID) error
+}
+
+func isTransferInFlight(status int) bool {
+	return status == STATUS_WAITING || status == STATUS_PENDING
+}
+
+// performTransferOnce runs a single transfer attempt against the given
+// location. The allocated batch ID is released on every return path.
+//
+// The engine refuses to free a batch while any of its tasks is still in
+// flight (BatchBusy), so on context cancellation the task is drained to a
+// terminal state before the deferred free runs. Cancellation latency is
+// therefore bounded by the engine's own transfer timeout.
+func performTransferOnce(ctx context.Context, transport batchTransport, source uintptr, length uint64, location *Location, useCache bool) (status int, err error) {
+	batchID, err := transport.allocateBatchID(1)
+	if err != nil {
+		return STATUS_FAILED, err
+	}
+	defer func() {
+		freeErr := transport.freeBatchID(batchID)
+		if freeErr != nil {
+			log.Println("cascading error: failed to free batch ID:", freeErr)
+			if err == nil {
+				err = freeErr
 			}
 		}
+	}()
 
-		err = store.transfer.freeBatchID(batchID)
-		if err != nil {
-			return err
-		}
+	targetID, err := transport.openSegment(location.SegmentName, useCache)
+	if err != nil {
+		return STATUS_FAILED, err
+	}
 
-		if status == STATUS_COMPLETED {
-			return nil
-		}
+	request := TransferRequest{
+		Opcode:       OPCODE_READ,
+		Source:       uint64(source),
+		TargetID:     targetID,
+		TargetOffset: location.Offset,
+		Length:       length,
+	}
 
-		retryCount++
+	err = transport.submitTransfer(batchID, []TransferRequest{request})
+	if err != nil {
+		return STATUS_FAILED, err
 	}
 
-	return ErrTooManyRetries
+	for isTransferInFlight(status) {
+		select {
+		case <-ctx.Done():
+			if _, drainErr := drainTransfer(transport, batchID); drainErr != nil {
+				return STATUS_FAILED, drainErr
+			}
+			return STATUS_FAILED, ctx.Err()
+		default:
+			status, _, err = transport.getTransferStatus(batchID, 0)
+			if err != nil {
+				return STATUS_FAILED, err
+			}
+		}
+	}
+
+	return status, nil
+}
+
+// drainTransfer polls the batch until its task reaches a terminal state so
+// the batch can be freed.
+func drainTransfer(transport batchTransport, batchID BatchID) (int, error) {
+	status := STATUS_WAITING
+	for isTransferInFlight(status) {
+		var err error
+		status, _, err = transport.getTransferStatus(batchID, 0)
+		if err != nil {
+			return STATUS_FAILED, err
+		}
+	}
+	return status, nil
 }
 
 func (store *P2PStore) updatePayloadMetadata(ctx context.Context, name string, addrList []uintptr, sizeList []uint64, payload *Payload, revision int64) error {

mooncake-p2p-store/src/p2pstore/perform_transfer_test.go

@@ -0,0 +1,120 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package p2pstore
+
+import (
+	"context"
+	"errors"
+	"testing"
+)
+
+var errBatchBusy = errors.New("BatchID cannot be freed until all tasks are done")
+
+// fakeBatchTransport mirrors the engine's batch lifecycle: a submitted task
+// stays in flight for pollsUntilDone status polls, and freeBatchID refuses
+// with BatchBusy while the task is in flight, as MultiTransport::freeBatchID does.
+type fakeBatchTransport struct {
+	pollsUntilDone int
+	finalStatus    int
+	submitErr      error
+
+	polls          int
+	submitted      bool
+	freed          bool
+	busyRejections int
+}
+
+func (f *fakeBatchTransport) allocateBatchID(int) (BatchID, error) { return BatchID(42), nil }
+
+func (f *fakeBatchTransport) openSegment(string, bool) (int64, error) { return 1, nil }
+
+func (f *fakeBatchTransport) submitTransfer(BatchID, []TransferRequest) error {
+	if f.submitErr != nil {
+		return f.submitErr
+	}
+	f.submitted = true
+	return nil
+}
+
+func (f *fakeBatchTransport) inFlight() bool {
+	return f.submitted && f.polls < f.pollsUntilDone
+}
+
+func (f *fakeBatchTransport) getTransferStatus(BatchID, int) (int, uint64, error) {
+	f.polls++
+	if f.inFlight() {
+		return STATUS_PENDING, 0, nil
+	}
+	return f.finalStatus, 0, nil
+}
+
+func (f *fakeBatchTransport) freeBatchID(BatchID) error {
+	if f.inFlight() {
+		f.busyRejections++
+		return errBatchBusy
+	}
+	f.freed = true
+	return nil
+}
+
+func testLocation() *Location {
+	return &Location{SegmentName: "peer:12345", Offset: 0}
+}
+
+// TestPerformTransferOnceCancellationFreesBatch is the regression for the
+// cancellation leak: with the task still pending, returning on ctx.Done()
+// without draining makes the deferred free hit BatchBusy and leak the batch.
+func TestPerformTransferOnceCancellationFreesBatch(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	fake := &fakeBatchTransport{pollsUntilDone: 3, finalStatus: STATUS_COMPLETED}
+	_, err := performTransferOnce(ctx, fake, 0, 4096, testLocation(), true)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("err = %v, want context.Canceled", err)
+	}
+	if !fake.freed {
+		t.Fatal("batch ID leaked on cancellation")
+	}
+	if fake.busyRejections != 0 {
+		t.Fatalf("free attempted %d time(s) while the task was in flight", fake.busyRejections)
+	}
+}
+
+func TestPerformTransferOnceCompletedFreesBatch(t *testing.T) {
+	fake := &fakeBatchTransport{pollsUntilDone: 2, finalStatus: STATUS_COMPLETED}
+	status, err := performTransferOnce(context.Background(), fake, 0, 4096, testLocation(), true)
+	if err != nil {
+		t.Fatalf("err = %v", err)
+	}
+	if status != STATUS_COMPLETED {
+		t.Fatalf("status = %d, want STATUS_COMPLETED", status)
+	}
+	if !fake.freed {
+		t.Fatal("batch ID leaked on completion")
+	}
+}
+
+func TestPerformTransferOnceSubmitErrorFreesBatch(t *testing.T) {
+	submitErr := errors.New("submit failed")
+	fake := &fakeBatchTransport{submitErr: submitErr}
+	_, err := performTransferOnce(context.Background(), fake, 0, 4096, testLocation(), true)
+	if !errors.Is(err, submitErr) {
+		t.Fatalf("err = %v, want submit error", err)
+	}
+	if !fake.freed {
+		t.Fatal("batch ID leaked on submit failure")
+	}
+}