fix: cancel per-attempt gRPC contexts promptly

sdk/executable.go

@@ -269,31 +269,15 @@ func _Execute(client *Client, e Executable) (interface{}, error) {
 
 		method := e.getMethod(channel)
 
-		var resp interface{}
-
 		var grpcDeadline time.Duration
 		if e.GetGrpcDeadline() != nil {
 			grpcDeadline = *e.GetGrpcDeadline()
 		} else {
 			grpcDeadline = client.GetGrpcDeadline()
 		}
-		ctx, cancel := context.WithTimeout(context.Background(), grpcDeadline)
-		defer cancel()
-
 		txLogger.Trace("executing gRPC call", "requestId", e.getLogID(e))
 
-		var marshaledResponse []byte
-		if method.query != nil {
-			resp, err = method.query(ctx, protoRequest.(*services.Query))
-			if err == nil {
-				marshaledResponse, _ = protobuf.Marshal(resp.(*services.Response))
-			}
-		} else {
-			resp, err = method.transaction(ctx, protoRequest.(*services.Transaction))
-			if err == nil {
-				marshaledResponse, _ = protobuf.Marshal(resp.(*services.TransactionResponse))
-			}
-		}
+		resp, marshaledResponse, err := _ExecuteRequest(method, grpcDeadline, protoRequest)
 
 		if err != nil {
 			e.advanceRequest()
@@ -380,6 +364,26 @@ func _Execute(client *Client, e Executable) (interface{}, error) {
 	return &services.Response{}, errPersistent
 }
 
+// _ExecuteRequest executes a single gRPC attempt within its own timeout context.
+func _ExecuteRequest(method _Method, grpcDeadline time.Duration, protoRequest interface{}) (interface{}, []byte, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), grpcDeadline)
+	defer cancel()
+
+	var resp protobuf.Message
+	var err error
+	if method.query != nil {
+		resp, err = method.query(ctx, protoRequest.(*services.Query))
+	} else {
+		resp, err = method.transaction(ctx, protoRequest.(*services.Transaction))
+	}
+	if err != nil {
+		return resp, nil, err
+	}
+
+	marshaledResponse, _ := protobuf.Marshal(resp)
+	return resp, marshaledResponse, nil
+}
+
 func _DelayForAttempt(logID string, backoff time.Duration, attempt int64, logger Logger, err error) {
 	logger.Trace("retrying request attempt", "requestId", logID, "delay", backoff, "attempt", attempt+1, "error", err)
 

sdk/mock_test.go

@@ -69,6 +69,50 @@ func TestUnitMockQuery(t *testing.T) {
 	require.NoError(t, err)
 }
 
+type contextTrackingFileCreateTransaction struct {
+	*FileCreateTransaction
+	method _Method
+}
+
+func (tx contextTrackingFileCreateTransaction) getMethod(*_Channel) _Method {
+	return tx.method
+}
+
+func TestUnitMockExecuteCancelsContextBetweenTransactionAttempts(t *testing.T) {
+	responses := [][]interface{}{{}}
+	client, server := NewMockClientAndServer(responses)
+	defer server.Close()
+
+	transaction := NewFileCreateTransaction().
+		SetNodeAccountIDs([]AccountID{{Account: 3}}).
+		SetContents([]byte("hello"))
+	_, err := transaction.FreezeWith(client)
+	require.NoError(t, err)
+
+	var firstAttemptContext context.Context
+	attempt := 0
+	executable := contextTrackingFileCreateTransaction{
+		FileCreateTransaction: transaction,
+		method: _Method{
+			transaction: func(ctx context.Context, _ *services.Transaction, _ ...grpc.CallOption) (*services.TransactionResponse, error) {
+				attempt++
+				if attempt == 1 {
+					firstAttemptContext = ctx
+					return &services.TransactionResponse{NodeTransactionPrecheckCode: services.ResponseCodeEnum_BUSY}, nil
+				}
+
+				require.NotNil(t, firstAttemptContext)
+				require.ErrorIs(t, firstAttemptContext.Err(), context.Canceled)
+				return &services.TransactionResponse{NodeTransactionPrecheckCode: services.ResponseCodeEnum_OK}, nil
+			},
+		},
+	}
+
+	_, err = _Execute(client, executable)
+	require.NoError(t, err)
+	require.Equal(t, 2, attempt)
+}
+
 func DisabledTestUnitMockBackoff(t *testing.T) {
 	responses := [][]interface{}{{
 		&services.TransactionResponse{