improve performance by using splitting steps to different transforms

Author: rluvaton

packages/pg-protocol/src/index.ts

@@ -1,10 +1,24 @@
 import { BackendMessage, DatabaseError } from './messages'
 import { serialize } from './serializer'
 import { Parser, MessageCallback } from './parser'
+import {pipeline} from 'stream';
 
 export function parse(stream: NodeJS.ReadableStream, callback: MessageCallback): Promise<void> {
   const parser = new Parser()
-  stream.on('data', (buffer: Buffer) => parser.parse(buffer, callback))
+  stream.on('data', (buffer: Buffer) => parser.parse(buffer, callback));
+
+  pipeline(
+    stream,
+    parser.splitMessagesTransform.bind(parser),
+    parser.convertToMessageTransform.bind(parser),
+    async function* (stream) {
+      for await(const message of stream) {
+        callback(message)
+      }
+    },
+    (err) => err ? reject(err) : resolve()
+  );
+
   return new Promise((resolve) => stream.on('end', () => resolve()))
 }
 

packages/pg-protocol/src/parser.ts

@@ -119,6 +119,48 @@ export class Parser {
     }
   }
 
+  async* splitMessagesTransform(stream: NodeJS.ReadableStream) {
+    for await(const buffer of stream) {
+      this.mergeBuffer(buffer as Buffer)
+      const bufferFullLength = this.bufferOffset + this.bufferLength
+      let offset = this.bufferOffset
+      while (offset + HEADER_LENGTH <= bufferFullLength) {
+        // code is 1 byte long - it identifies the message type
+        const code = this.buffer[offset]
+        // length is 1 Uint32BE - it is the length of the message EXCLUDING the code
+        const length = this.buffer.readUInt32BE(offset + CODE_LENGTH)
+        const fullMessageLength = CODE_LENGTH + length
+        if (fullMessageLength + offset <= bufferFullLength) {
+          yield {
+            offset: offset + HEADER_LENGTH,
+            code,
+            length,
+            bytes: this.buffer
+          };
+          offset += fullMessageLength
+        } else {
+          break
+        }
+      }
+      if (offset === bufferFullLength) {
+        // No more use for the buffer
+        this.buffer = emptyBuffer
+        this.bufferLength = 0
+        this.bufferOffset = 0
+      } else {
+        // Adjust the cursors of remainingBuffer
+        this.bufferLength = bufferFullLength - offset
+        this.bufferOffset = offset
+      }
+    }
+  }
+
+  async* convertToMessageTransform(stream: NodeJS.ReadableStream) {
+    for await(const rawMessage of stream) {
+      yield this.handlePacket((rawMessage as any).offset, (rawMessage as any).code, (rawMessage as any).length, (rawMessage as any).bytes);
+    }
+  }
+
   private mergeBuffer(buffer: Buffer): void {
     if (this.bufferLength > 0) {
       const newLength = this.bufferLength + buffer.byteLength