Use a pool of buffers when publishing

carlhoerberg · carlhoerberg · commit 3ad05ade3dfd · 2022-08-04T12:46:55.000+02:00
Before when a single buffer was shared multiple simulations calls to AMQPChannel#basicPublish could corrupt the buffer. Fixes #36 Downside with this fix is that if many microtasks run simulantious many buffers will be created, but never cleaned up, so there's a risk of memory leak. Upside compared to allocating a new buffer for each call to basicPublish is that this way is a lot faster (about 50%).
diff --git a/src/amqp-base-client.ts b/src/amqp-base-client.ts
@@ -24,6 +24,10 @@ export abstract class AMQPBaseClient {
   frameMax: number
   heartbeat: number
   onerror: (error: AMQPError) => void
+  /** Used for string -> arraybuffer when publishing */
+  readonly textEncoder = new TextEncoder()
+  // Buffer pool for publishes, let multiple microtasks publish at the same time but save on allocations
+  readonly bufferPool: AMQPView[] = []
 
   /**
    * @param name - name of the connection, set in client properties
diff --git a/src/amqp-channel.ts b/src/amqp-channel.ts
@@ -16,10 +16,6 @@ export class AMQPChannel {
   readonly promises: [(value?: any) => void, (err?: Error) => void][] = []
   private readonly unconfirmedPublishes: [number, (confirmId: number) => void, (err?: Error) => void][] = []
   closed = false
-  /** Used for string -> arraybuffer when publishing */
-  private static textEncoder = new TextEncoder()
-  /** Frame buffer, reuse when publishes to avoid repated allocations */
-  private readonly buffer: AMQPView
   confirmId = 0
   delivery?: AMQPMessage
   getMessage?: AMQPMessage
@@ -31,7 +27,6 @@ export class AMQPChannel {
   constructor(connection: AMQPBaseClient, id: number) {
     this.connection = connection
     this.id = id
-    this.buffer = new AMQPView(new ArrayBuffer(connection.frameMax))
   }
 
   /**
@@ -288,13 +283,14 @@ export class AMQPChannel {
     } else if (data === null) {
       body = new Uint8Array(0)
     } else if (typeof data === "string") {
-      body = AMQPChannel.textEncoder.encode(data)
+      body = this.connection.textEncoder.encode(data)
     } else {
       throw new TypeError(`Invalid type ${typeof data} for parameter data`)
     }
 
     let j = 0
-    const buffer = this.buffer
+    // get a buffer from the pool or create a new, it will later be returned to the pool for reuse
+    const buffer = this.connection.bufferPool.pop() || new AMQPView(new ArrayBuffer(this.connection.frameMax))
     buffer.setUint8(j, 1); j += 1 // type: method
     buffer.setUint16(j, this.id); j += 2 // channel
     j += 4 // frame size, update later
@@ -346,6 +342,7 @@ export class AMQPChannel {
       bodyPos += frameSize
       j = 0
     }
+    this.connection.bufferPool.push(buffer) // return buffer to buffer pool for later reuse
     // if publish confirm is enabled, put a promise on a queue if the sends were ok
     // the promise on the queue will be fullfilled by the read loop when an ack/nack
     // comes from the server
diff --git a/test/test.ts b/test/test.ts
@@ -613,3 +613,54 @@ test("onerror is not called when conn is closed by client", async t => {
   await new Promise((resolv) => setTimeout(resolv, 10))
   t.pass()
 })
+
+test("will throw on too large headers", async t => {
+  const amqp = new AMQPClient("amqp://127.0.0.1?frameMax=4096")
+  const conn = await amqp.connect()
+  const ch = await conn.channel()
+  await t.throwsAsync(() => ch.basicPublish("", "x".repeat(255), null, {"headers": {a: Array(4000).fill(1)}}),
+                      { instanceOf: RangeError })
+  await t.throwsAsync(() => ch.basicPublish("", "", null, {"headers": {a: "x".repeat(5000)}}),
+                      { instanceOf: RangeError })
+})
+
+test("will split body over multiple frames", async t => {
+  const amqp = new AMQPClient("amqp://127.0.0.1?frameMax=4096")
+  const conn = await amqp.connect()
+  const ch = await conn.channel()
+  const q = await ch.queue("")
+  await ch.confirmSelect()
+  await q.publish("x".repeat(5000))
+  const msg = await q.get()
+  if (msg)
+    if (msg.body)
+      t.is(msg.body.length, 5000)
+    else
+      t.fail("no body")
+  else
+    t.fail("no msg")
+})
+
+test("can republish in consume block without race condition", async t => {
+  const amqp = new AMQPClient("amqp://127.0.0.1")
+  const conn = await amqp.connect()
+  const ch = await conn.channel()
+  await ch.prefetch(0)
+  const q = await ch.queue("")
+  await ch.confirmSelect()
+  await q.publish("x".repeat(500))
+  const consumer = await q.subscribe({noAck: false}, async (msg) => {
+    if (msg.deliveryTag < 10000) {
+      await Promise.all([
+        q.publish(msg.body),
+        q.publish(msg.body),
+        msg.ack()
+      ])
+    } else if (msg.deliveryTag === 10000) {
+      await consumer.cancel()
+    }
+  })
+  await t.notThrowsAsync(() => consumer.wait())
+  await t.notThrowsAsync(() => conn.close())
+  console.log(conn.bufferPool.length)
+})
