Use nextMessageToSend for early half-close

Author: serkanerip

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

@@ -123,11 +123,6 @@ 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;
   startTime: Date;
 }
 
@@ -700,7 +695,6 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
       state: 'ACTIVE',
       call: child,
       nextMessageToSend: 0,
-      highestSentMessageIndex: -1,
       startTime: new Date(),
     });
     const previousAttempts = this.attempts - 1;
@@ -793,7 +787,6 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
             },
             bufferEntry.message!.message
           );
-          childCall.highestSentMessageIndex = messageIndex;
           // 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.
@@ -833,7 +826,11 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
     };
     this.writeBuffer.push(bufferEntry);
     if (bufferEntry.allocated) {
-      context.callback?.();
+      // Run this in next tick to avoid suspending the current execution context
+      // otherwise it might cause half closing the call before sending message
+      process.nextTick(() => {
+        context.callback?.();
+      });
       for (const [callIndex, call] of this.underlyingCalls.entries()) {
         if (
           call.state === 'ACTIVE' &&
@@ -848,7 +845,6 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
             },
             message
           );
-          call.highestSentMessageIndex = messageIndex;
         }
       }
     } else {
@@ -860,7 +856,6 @@ 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 => {
@@ -891,11 +886,11 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
     });
     for (const call of this.underlyingCalls) {
       if (call?.state === 'ACTIVE') {
-        // Send halfClose immediately if all messages have been sent to this call
-        // We check highestSentMessageIndex >= halfCloseIndex - 1 because:
-        // - If halfCloseIndex is N, the last message is at index N-1
-        // - If highestSentMessageIndex >= N-1, all messages have been sent
-        if (call.highestSentMessageIndex >= halfCloseIndex - 1) {
+        // Send halfClose to call when either:
+        // - nextMessageToSend === halfCloseIndex - 1: last message sent, callback pending (optimization)
+        // - nextMessageToSend === halfCloseIndex: all messages sent and acknowledged
+        if (call.nextMessageToSend === halfCloseIndex 
+          || call.nextMessageToSend === halfCloseIndex - 1) {
           this.trace(
             'Sending halfClose immediately to child [' +
               call.call.getCallNumber() +
@@ -904,7 +899,7 @@ export class RetryingCall implements Call, DeadlineInfoProvider {
           call.nextMessageToSend += 1;
           call.call.halfClose();
         }
-        // Otherwise, halfClose will be sent by sendNextChildMessage when messages complete
+        // Otherwise, halfClose will be sent by sendNextChildMessage when message callbacks complete
       }
     }
   }

packages/grpc-js/test/test-end-to-end.ts

@@ -18,7 +18,7 @@
 import * as assert from 'assert';
 import * as path from 'path';
 import { loadProtoFile } from './common';
-import { Metadata, Server, ServerDuplexStream, ServerUnaryCall, ServiceClientConstructor, ServiceError, experimental, sendUnaryData } from '../src';
+import { Metadata, Server, ServerCredentials, ServerDuplexStream, ServerReadableStream, ServerUnaryCall, ServiceClientConstructor, ServiceError, credentials, experimental, sendUnaryData } from '../src';
 import { ServiceClient } from '../src/make-client';
 
 const protoFile = path.join(__dirname, 'fixtures', 'echo_service.proto');
@@ -36,6 +36,15 @@ const echoServiceImplementation = {
       call.end();
     });
   },
+  echoClientStream(call: ServerReadableStream<any, any>, callback: sendUnaryData<any>) {
+    const messages: any[] = [];
+    call.on('data', (message: any) => {
+      messages.push(message);
+    });
+    call.on('end', () => {
+      callback(null, { value: messages.map(m => m.value).join(','), value2: messages.length });
+    });
+  },
 };
 
 describe('Client should successfully communicate with server', () => {
@@ -77,4 +86,20 @@ describe('Client should successfully communicate with server', () => {
       });
     });
   }).timeout(5000);
+
+  it('Client streaming with one message should work', done => {
+    server = new Server();
+    server.addService(EchoService.service, echoServiceImplementation);
+    server.bindAsync('localhost:0', ServerCredentials.createInsecure(), (error, port) => {
+      assert.ifError(error);
+      client = new EchoService(`localhost:${port}`, credentials.createInsecure());
+      const call = client.echoClientStream((error: ServiceError, response: any) => {
+        assert.ifError(error);
+        assert.deepStrictEqual(response, { value: 'test value', value2: 1 });
+        done();
+      });
+      call.write({ value: 'test value', value2: 42 });
+      call.end();
+    });
+  });
 });