Use BufferedConnection in the http_connection

This commit overall improves the performances of sending journal entries
to the runtime by 4x, by reducing the number of temporary buffer
allocations, promises, and intermediate serialization.

Author: igalshilman

src/connection/buffered_connection.ts

@@ -0,0 +1,47 @@
+import { encodeMessages } from "../io/encoder";
+import { Message } from "../types/types";
+import { Connection } from "./connection";
+
+export class BufferedConnection implements Connection {
+  private queue: Message[] = [];
+  private flushing: Promise<void> = Promise.resolve();
+
+  constructor(private readonly flushFn: (buffer: Buffer) => Promise<void>) {}
+
+  send(msg: Message): Promise<void> {
+    const len = this.queue.push(msg);
+    if (len === 1) {
+      // we are the first in line, therefore we schedule a flush,
+      // BUT we must wait for the previous flush to end.
+      this.flushing = this.flushing.then(() => this.scheduleFlush());
+    }
+    // we don't need to reschedule the `flush` here,
+    // because the flush happens anyway at the end of the current event loop iteration.
+    // tag along to the previously scheduled flush.
+    return this.flushing;
+  }
+
+  end(): Promise<void> {
+    this.flushing = this.flushing.then(() => this.flush());
+    return this.flushing;
+  }
+
+  private scheduleFlush(): Promise<void> {
+    // schedule a flush at the end of the current event loop iteration.
+    return new Promise((resolve, reject) =>
+      setImmediate(() => {
+        this.flush().then(resolve).catch(reject);
+      })
+    );
+  }
+
+  private async flush(): Promise<void> {
+    if (this.queue.length === 0) {
+      return Promise.resolve();
+    }
+    const buffer = encodeMessages(this.queue) as Buffer;
+    this.queue = [];
+
+    return this.flushFn(buffer);
+  }
+}

src/connection/embedded_connection.ts

@@ -10,8 +10,8 @@
  */
 
 import { RemoteContext } from "../generated/proto/services";
-import { encodeMessages } from "../io/encoder";
 import { Message } from "../types/types";
+import { BufferedConnection } from "./buffered_connection";
 import { Connection } from "./connection";
 
 export class FencedOffError extends Error {
@@ -27,47 +27,25 @@ export class InvocationAlreadyCompletedError extends Error {
 }
 
 export class EmbeddedConnection implements Connection {
-  private queue: Message[] = [];
-  private flushing: Promise<void> = Promise.resolve();
+  private buffered: BufferedConnection;
 
   constructor(
     private readonly operationId: string,
     private readonly streamId: string,
     private readonly remote: RemoteContext
-  ) {}
+  ) {
+    this.buffered = new BufferedConnection((buffer) => this.sendBuffer(buffer));
+  }
 
   send(msg: Message): Promise<void> {
-    const len = this.queue.push(msg);
-    if (len === 1) {
-      // we are the first in line, therefore we schedule a flush,
-      // BUT we must wait for the previous flush to end.
-      this.flushing = this.flushing.then(() => this.scheduleFlush());
-    }
-    // tag along to the previously scheduled flush.
-    return this.flushing;
+    return this.buffered.send(msg);
   }
 
   end(): Promise<void> {
-    this.flushing = this.flushing.then(() => this.flush());
-    return this.flushing;
-  }
-
-  private scheduleFlush(): Promise<void> {
-    // schedule a flush at the end of the current event loop iteration.
-    return new Promise((resolve, reject) =>
-      setImmediate(() => {
-        this.flush().then(resolve).catch(reject);
-      })
-    );
+    return this.buffered.end();
   }
 
-  private async flush(): Promise<void> {
-    if (this.queue.length === 0) {
-      return Promise.resolve();
-    }
-    const buffer = encodeMessages(this.queue) as Buffer;
-    this.queue = [];
-
+  private async sendBuffer(buffer: Buffer): Promise<void> {
     const res = await this.remote.send({
       operationId: this.operationId,
       streamId: this.streamId,

src/connection/http_connection.ts

@@ -10,12 +10,12 @@
  */
 
 import stream from "stream";
-import { encodeMessage } from "../io/encoder";
 import { streamDecoder } from "../io/decoder";
 import { Connection, RestateStreamConsumer } from "./connection";
 import { Message } from "../types/types";
 import { rlog } from "../utils/logger";
 import { finished } from "stream/promises";
+import { BufferedConnection } from "./buffered_connection";
 
 // utility promise, for cases where we want to save allocation of an extra promise
 const RESOLVED: Promise<void> = Promise.resolve();
@@ -55,19 +55,25 @@ export class RestateHttp2Connection implements Connection {
   // input as decoded messages
   private readonly sdkInput: stream.Readable;
 
-  // output as encoded bytes. we convert manually, not as transforms,
-  // to skip a layer of stream indirection
-  private readonly sdkOutput: stream.Writable;
-
   // consumer handling
   private currentConsumer: RestateStreamConsumer | null = null;
   private inputBuffer: Message[] = [];
   private consumerError?: Error;
   private consumerInputClosed = false;
 
+  private outputBuffer: BufferedConnection;
+
   constructor(private readonly rawStream: stream.Duplex) {
     this.sdkInput = rawStream.pipe(streamDecoder());
-    this.sdkOutput = rawStream;
+
+    this.outputBuffer = new BufferedConnection((buffer) => {
+      const hasMoreCapacity = rawStream.write(buffer);
+      if (hasMoreCapacity) {
+        return RESOLVED;
+      } else {
+        return new Promise((resolve) => rawStream.once("drain", resolve));
+      }
+    });
 
     // remember and forward messages
     this.sdkInput.on("data", (m: Message) => {
@@ -201,23 +207,16 @@ export class RestateHttp2Connection implements Connection {
    * capacity again, so that at least the operations that await results will respect backpressure.
    */
   public send(msg: Message): Promise<void> {
-    const encodedMessage: Uint8Array = encodeMessage(msg);
-
-    const hasMoreCapacity = this.sdkOutput.write(encodedMessage);
-    if (hasMoreCapacity) {
-      return RESOLVED;
-    }
-
-    return new Promise((resolve) => {
-      this.sdkOutput.once("drain", resolve);
-    });
+    return this.outputBuffer.send(msg);
   }
 
   /**
    * Ends the stream, awaiting pending writes.
    */
   public async end(): Promise<void> {
-    this.sdkOutput.end();
+    await this.outputBuffer.end();
+
+    this.rawStream.end();
 
     const options = {
       error: true,

src/io/encoder.ts

@@ -12,6 +12,7 @@
 import stream from "stream";
 import { PROTOBUF_MESSAGE_BY_TYPE } from "../types/protocol";
 import { Header, Message } from "../types/types";
+import _m0 from "protobufjs/minimal";
 
 export function streamEncoder(): stream.Transform {
   return new stream.Transform({
@@ -32,7 +33,11 @@ export function encodeMessage(msg: Message): Uint8Array {
 }
 
 export function encodeMessages(messages: Message[]): Uint8Array {
-  const chunks: Buffer[] = [];
+  const offsets = [];
+  const headers = [];
+  let off = 0;
+
+  const writer = _m0.Writer.create();
   for (const message of messages) {
     const pbType = PROTOBUF_MESSAGE_BY_TYPE.get(BigInt(message.messageType));
     if (pbType === undefined) {
@@ -41,19 +46,27 @@ export function encodeMessages(messages: Message[]): Uint8Array {
           message.messageType
       );
     }
-    const bodyBuf = pbType.encode(message.message).finish();
+    offsets.push(off);
+    writer.fixed64(0);
+    pbType.encode(message.message, writer);
+    const messageLen = writer.len - 8 - off;
+    off = writer.len;
+
     const header = new Header(
       BigInt(message.messageType),
-      bodyBuf.length,
+      messageLen,
       message.completed,
       message.protocolVersion, // only set for incoming start message
       message.requiresAck
     );
-    const encoded = header.toU64be();
-    const headerBuf = Buffer.alloc(8);
-    headerBuf.writeBigUInt64BE(encoded);
-    chunks.push(headerBuf);
-    chunks.push(bodyBuf);
+    const header64 = header.toU64be();
+    headers.push(header64);
+  }
+  const buffer = writer.finish() as Buffer;
+  for (let i = 0; i < offsets.length; i++) {
+    const offset = offsets[i];
+    const header = headers[i];
+    buffer.writeBigUInt64BE(header, offset);
   }
-  return Buffer.concat(chunks);
+  return buffer;
 }

src/journal.ts

@@ -39,6 +39,8 @@ import { SideEffectEntryMessage } from "./generated/proto/javascript";
 import { Invocation } from "./invocation";
 import { failureToError, RetryableError } from "./types/errors";
 
+const RESOLVED = Promise.resolve(undefined);
+
 export class Journal<I, O> {
   private state = NewExecutionState.REPLAYING;
 
@@ -113,7 +115,7 @@ export class Journal<I, O> {
           case p.COMPLETE_AWAKEABLE_ENTRY_MESSAGE_TYPE:
           case p.BACKGROUND_INVOKE_ENTRY_MESSAGE_TYPE: {
             // Do not need completion
-            return Promise.resolve(undefined);
+            return RESOLVED;
           }
           case p.GET_STATE_ENTRY_MESSAGE_TYPE: {
             const getStateMsg = message as GetStateEntryMessage;

src/restate_context_impl.ts

@@ -444,7 +444,10 @@ async function executeWithRetries<T>(
       //  - they are not retried within the service, because they will never succeed within this service,
       //    but can only succeed within a new invocation going to service with fixed code
       //  we hence break the retries here similar to terminal errors
-      if (e instanceof RestateError && e.code == RestateErrorCodes.JOURNAL_MISMATCH) {
+      if (
+        e instanceof RestateError &&
+        e.code == RestateErrorCodes.JOURNAL_MISMATCH
+      ) {
         throw e;
       }
 

src/types/errors.ts

@@ -225,7 +225,9 @@ export class RetryableError extends RestateError {
         - In the replayed messages: type: ${
           replayMessage.messageType
         }, message: ${printMessageAsJson(replayMessage.message)}`;
-    return new RetryableError(msg, { errorCode: RestateErrorCodes.JOURNAL_MISMATCH });
+    return new RetryableError(msg, {
+      errorCode: RestateErrorCodes.JOURNAL_MISMATCH,
+    });
   }
 
   public static protocolViolation(message: string) {

test/protoutils.ts

@@ -55,7 +55,7 @@ import {
 import { expect } from "@jest/globals";
 import { jsonSerialize, printMessageAsJson } from "../src/utils/utils";
 import { rlog } from "../src/utils/logger";
-import {ErrorCodes, RestateErrorCodes} from "../src/types/errors";
+import { ErrorCodes, RestateErrorCodes } from "../src/types/errors";
 
 export function startMessage(
   knownEntries?: number,