Send halfClose immediately after messages to prevent late halfClose issues with Envoy

Author: serkanerip

packages/grpc-js/src/retrying-call.ts

@@ -123,6 +123,15 @@ interface UnderlyingCall {
   state: UnderlyingCallState;
   call: LoadBalancingCall;
   nextMessageToSend: number;
+  /**
+   * Tracks the highest message index that has been sent to the underlying call.
+   * This is different from nextMessageToSend which tracks completion/acknowledgment.
+   */
+  highestSentMessageIndex: number;
+  /**
+   * Tracks whether halfClose has been sent to this child call.
+   */
+  halfCloseSent: boolean;
   startTime: Date;
 }
 
@@ -695,6 +704,8 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
       state: 'ACTIVE',
       call: child,
       nextMessageToSend: 0,
+      highestSentMessageIndex: -1,
+      halfCloseSent: false,
       startTime: new Date(),
     });
     const previousAttempts = this.attempts - 1;
@@ -778,6 +789,7 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
       const bufferEntry = this.getBufferEntry(childCall.nextMessageToSend);
       switch (bufferEntry.entryType) {
         case 'MESSAGE':
+          childCall.highestSentMessageIndex = childCall.nextMessageToSend;
           childCall.call.sendMessageWithContext(
             {
               callback: error => {
@@ -787,10 +799,26 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
             },
             bufferEntry.message!.message
           );
+          // Optimization: if the next entry is HALF_CLOSE, send it immediately
+          // without waiting for the message callback. This is safe because the message
+          // has already been passed to the underlying transport.
+          const nextEntry = this.getBufferEntry(childCall.nextMessageToSend + 1);
+          if (nextEntry.entryType === 'HALF_CLOSE' && !childCall.halfCloseSent) {
+            this.trace(
+              'Sending halfClose immediately after message to child [' +
+                childCall.call.getCallNumber() +
+                '] - optimizing for unary/final message'
+            );
+            childCall.halfCloseSent = true;
+            childCall.call.halfClose();
+          }
           break;
         case 'HALF_CLOSE':
-          childCall.nextMessageToSend += 1;
-          childCall.call.halfClose();
+          if (!childCall.halfCloseSent) {
+            childCall.nextMessageToSend += 1;
+            childCall.halfCloseSent = true;
+            childCall.call.halfClose();
+          }
           break;
         case 'FREED':
           // Should not be possible
@@ -819,6 +847,7 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
           call.state === 'ACTIVE' &&
           call.nextMessageToSend === messageIndex
         ) {
+          call.highestSentMessageIndex = messageIndex;
           call.call.sendMessageWithContext(
             {
               callback: error => {
@@ -839,6 +868,7 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
       const call = this.underlyingCalls[this.committedCallIndex];
       bufferEntry.callback = context.callback;
       if (call.state === 'ACTIVE' && call.nextMessageToSend === messageIndex) {
+        call.highestSentMessageIndex = messageIndex;
         call.call.sendMessageWithContext(
           {
             callback: error => {
@@ -868,12 +898,22 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
       allocated: false,
     });
     for (const call of this.underlyingCalls) {
-      if (
-        call?.state === 'ACTIVE' &&
-        call.nextMessageToSend === halfCloseIndex
-      ) {
-        call.nextMessageToSend += 1;
-        call.call.halfClose();
+      if (call?.state === 'ACTIVE' && !call.halfCloseSent) {
+        // Send halfClose immediately if all messages have been sent to this call
+        // We check highestSentMessageIndex >= halfCloseIndex - 1 because:
+        // - If halfCloseIndex is 0, there are no messages, so send immediately
+        // - If halfCloseIndex is N, the last message is at index N-1
+        // - If highestSentMessageIndex >= N-1, all messages have been sent
+        if (halfCloseIndex === 0 || call.highestSentMessageIndex >= halfCloseIndex - 1) {
+          this.trace(
+            'Sending halfClose immediately to child [' +
+              call.call.getCallNumber() +
+              '] - all messages already sent'
+          );
+          call.halfCloseSent = true;
+          call.call.halfClose();
+        }
+        // Otherwise, halfClose will be sent by sendNextChildMessage when messages complete
       }
     }
   }