fix(server-utils): Record amqplib consumer info before the first delivery and instrument the callback API

dev-packages/node-integration-tests/suites/tracing/amqplib/scenario-callback-api.mjs

@@ -0,0 +1,51 @@
+import * as Sentry from '@sentry/node';
+import amqp from 'amqplib/callback_api.js';
+
+const queueName = 'callback-queue';
+const amqpUsername = 'sentry';
+const amqpPassword = 'sentry';
+
+const AMQP_URL = `amqp://${amqpUsername}:${amqpPassword}@localhost:5672/`;
+
+(async () => {
+  const connection = await connectToRabbitMQ();
+  const channel = await toPromise(cb => connection.createChannel(cb));
+  await toPromise(cb => channel.assertQueue(queueName, { durable: false, exclusive: true }, cb));
+
+  // Without the info from `consume`, a `noAck` span stays open until the channel closes and ends as an error.
+  const received = new Promise((resolve, reject) => {
+    channel.consume(
+      queueName,
+      message => (message ? resolve() : reject(new Error('No message received'))),
+      { noAck: true },
+      err => err && reject(err),
+    );
+  });
+
+  await Sentry.startSpan({ name: 'root span' }, async () => {
+    channel.sendToQueue(queueName, Buffer.from(JSON.stringify({ foo: 'bar01' })));
+  });
+
+  await received;
+
+  await toPromise(cb => channel.close(cb));
+  await toPromise(cb => connection.close(cb));
+})();
+
+async function connectToRabbitMQ() {
+  // The broker accepts TCP before AMQP handshakes succeed, so early connects can fail while it boots.
+  let lastError;
+  for (let attempt = 0; attempt < 5; attempt++) {
+    try {
+      return await toPromise(cb => amqp.connect(AMQP_URL, cb));
+    } catch (err) {
+      lastError = err;
+      await new Promise(resolve => setTimeout(resolve, 1000));
+    }
+  }
+  throw lastError;
+}
+
+function toPromise(fn) {
+  return new Promise((resolve, reject) => fn((err, result) => (err ? reject(err) : resolve(result))));
+}

dev-packages/node-integration-tests/suites/tracing/amqplib/scenario-topic-noack.mjs

@@ -0,0 +1,63 @@
+import * as Sentry from '@sentry/node';
+import amqp from 'amqplib';
+
+// A topic exchange, so the routing key and the queue name differ.
+const exchangeName = 'orders';
+const queueName = 'orders-worker';
+const routingKey = 'order.created.12345';
+const amqpUsername = 'sentry';
+const amqpPassword = 'sentry';
+
+const AMQP_URL = `amqp://${amqpUsername}:${amqpPassword}@localhost:5672/`;
+
+// A `noAck` consumer never settles a message, so its span has to end when dispatch returns.
+const NO_ACKNOWLEDGEMENT = { noAck: true };
+
+(async () => {
+  const { connection, channel } = await connectToRabbitMQ();
+
+  await channel.assertExchange(exchangeName, 'topic', { durable: false, autoDelete: true });
+  await channel.assertQueue(queueName, { durable: false, exclusive: true });
+  await channel.bindQueue(queueName, exchangeName, 'order.created.#');
+
+  // Publish before consuming, so the broker sends `BasicConsumeOk` and the delivery back to back.
+  await Sentry.startSpan({ name: 'root span' }, async () => {
+    channel.publish(exchangeName, routingKey, Buffer.from(JSON.stringify({ foo: 'bar01' })));
+  });
+
+  await consumeMessageFromQueue(queueName, channel);
+
+  await channel.close();
+  await connection.close();
+})();
+
+async function connectToRabbitMQ() {
+  // The broker accepts TCP before AMQP handshakes succeed, so early connects can reject while it boots.
+  let lastError;
+  for (let attempt = 0; attempt < 5; attempt++) {
+    try {
+      const connection = await amqp.connect(AMQP_URL);
+      const channel = await connection.createChannel();
+      return { connection, channel };
+    } catch (err) {
+      lastError = err;
+      await new Promise(resolve => setTimeout(resolve, 1000));
+    }
+  }
+  throw lastError;
+}
+
+async function consumeMessageFromQueue(queueName, channel) {
+  return new Promise((resolve, reject) => {
+    const onMessage = message => (message ? resolve() : reject(new Error('No message received')));
+
+    // Invalid `arguments` throw before the RPC, so the channel stays open for the next `consume`.
+    try {
+      channel.consume('orders-archive', onMessage, { arguments: 'invalid' });
+    } catch {
+      // Expected.
+    }
+
+    channel.consume(queueName, onMessage, NO_ACKNOWLEDGEMENT).catch(reject);
+  });
+}

dev-packages/node-integration-tests/suites/tracing/amqplib/test.ts

@@ -168,5 +168,56 @@ describeWithDockerCompose('amqplib auto-instrumentation', { workingDirectory: [_
       },
       { additionalDependencies },
     );
+
+    createEsmAndCjsTests(
+      __dirname,
+      'scenario-topic-noack.mjs',
+      'instrument.mjs',
+      (createTestRunner, test) => {
+        test('ends a noAck consumer span as ok', { timeout: 60_000 }, async () => {
+          await createTestRunner()
+            .expect({
+              span: container => {
+                const consumer = container.items.find(
+                  span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.amqplib.consumer',
+                );
+
+                expect(consumer).toBeDefined();
+                expect(consumer!.name).toBe('process orders');
+                expect(consumer!.status).toBe('ok');
+                expect(consumer!.attributes[MESSAGING_DESTINATION_NAME]?.value).toBe('orders');
+                expect(consumer!.attributes[MESSAGING_RABBITMQ_DESTINATION_ROUTING_KEY]?.value).toBe(
+                  'order.created.12345',
+                );
+              },
+            })
+            .start()
+            .completed();
+        });
+      },
+      { additionalDependencies },
+    );
+  });
+
+  createEsmAndCjsTests(__dirname, 'scenario-callback-api.mjs', 'instrument.mjs', (createTestRunner, test) => {
+    test('instruments publish and noAck consume on the callback API', { timeout: 60_000 }, async () => {
+      await createTestRunner()
+        .expect({
+          span: container => {
+            const producerSpan = container.items.find(
+              span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.amqplib.publisher',
+            );
+            expect(producerSpan?.name).toBe('send callback-queue');
+
+            const consumerSpan = container.items.find(
+              span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.amqplib.consumer',
+            );
+            expect(consumerSpan?.name).toBe('process callback-queue');
+            expect(consumerSpan?.status).toBe('ok');
+          },
+        })
+        .start()
+        .completed();
+    });
   });
 });

packages/server-utils/src/integrations/amqplib.ts

@@ -68,6 +68,7 @@ const MESSAGE_STORED_SPAN: unique symbol = Symbol('sentry.amqplib.message.stored
 const CHANNEL_SPANS_NOT_ENDED: unique symbol = Symbol('sentry.amqplib.channel.spans-not-ended');
 const CHANNEL_CONSUME_TIMEOUT_TIMER: unique symbol = Symbol('sentry.amqplib.channel.consume-timeout-timer');
 const CHANNEL_CONSUMER_INFO: unique symbol = Symbol('sentry.amqplib.channel.consumer-info');
+const CHANNEL_PENDING_CONSUMERS: unique symbol = Symbol('sentry.amqplib.channel.pending-consumers');
 const CHANNEL_IS_CONFIRM_PUBLISHING: unique symbol = Symbol('sentry.amqplib.channel.is-confirm-publishing');
 const CONNECTION_ATTRIBUTES: unique symbol = Symbol('sentry.amqplib.connection.attributes');
 
@@ -108,12 +109,19 @@ interface ConsumerInfo {
   queue: string;
 }
 
+/** A `consume` call waiting for the broker to assign it a consumer tag. */
+interface PendingConsumer {
+  callback: unknown;
+  info: ConsumerInfo;
+}
+
 interface ChannelLike {
   connection?: ConnectionLike;
   on?: (event: string, listener: (...args: unknown[]) => void) => unknown;
   [CHANNEL_SPANS_NOT_ENDED]?: { msg: ConsumeMessage; timeOfConsume: number }[];
   [CHANNEL_CONSUME_TIMEOUT_TIMER]?: ReturnType<typeof setInterval>;
   [CHANNEL_CONSUMER_INFO]?: Map<string, ConsumerInfo>;
+  [CHANNEL_PENDING_CONSUMERS]?: PendingConsumer[];
   [CHANNEL_IS_CONFIRM_PUBLISHING]?: boolean;
 }
 
@@ -135,6 +143,11 @@ interface AmqpDispatchContext extends AmqpChannelContext {
   _sentryNoAck?: boolean;
 }
 
+interface AmqpConsumeContext extends AmqpChannelContext {
+  // The entry this call queued, so the error hook can drop it.
+  _sentryPendingConsumer?: PendingConsumer;
+}
+
 interface AmqpConnectContext {
   arguments?: unknown[];
   result?: unknown;
@@ -156,6 +169,7 @@ function instrumentAmqplib(): void {
   subscribePublish();
   subscribeConfirmPublish();
   subscribeConsume();
+  subscribeRegisterConsumer();
   subscribeDispatch();
   subscribeSettle();
 }
@@ -199,32 +213,68 @@ function subscribeConfirmPublish(): void {
   });
 }
 
-/**
- * Records `consumerTag -> { noAck, queue }` when a consumer is registered, so the per-message
- * dispatch hook can name the span after the queue and know when to end it.
- */
+/** Stashes the `{ noAck, queue }` each `consume` asks for until the broker assigns it a tag. */
 function subscribeConsume(): void {
-  const channel = diagnosticsChannel.tracingChannel<AmqpChannelContext>(CHANNELS.AMQPLIB_CONSUME);
+  const channel = diagnosticsChannel.tracingChannel<AmqpConsumeContext>(CHANNELS.AMQPLIB_CONSUME);
 
-  // A `start` subscriber is required for orchestrion to wrap `consume` at all.
-  channel.start.subscribe(NOOP);
-  channel.asyncEnd.subscribe(message => {
-    const data = message as AmqpChannelContext;
+  channel.start.subscribe(message => {
+    const data = message as AmqpConsumeContext;
     const consumerChannel = data.self;
-    const result = data.result as { consumerTag?: string } | undefined;
-    const consumerTag = result?.consumerTag;
-    if (!consumerChannel || !consumerTag) {
+    if (!consumerChannel) {
       return;
     }
 
     ensureChannelState(consumerChannel);
     const queueArg = data.arguments[0];
     const queue = typeof queueArg === 'string' ? queueArg : '<unknown>';
+    const callback = data.arguments[1];
     const options = data.arguments[2] as { noAck?: boolean } | undefined;
-    consumerChannel[CHANNEL_CONSUMER_INFO]?.set(consumerTag, { noAck: !!options?.noAck, queue });
+    const entry = { callback, info: { noAck: !!options?.noAck, queue } };
+    consumerChannel[CHANNEL_PENDING_CONSUMERS]?.push(entry);
+    data._sentryPendingConsumer = entry;
+  });
+
+  // A synchronous throw, such as for invalid `arguments`, leaves the channel open. RPC failures close it.
+  // A stale entry would label the next consumer that shares this callback.
+  channel.error.subscribe(message => {
+    const data = message as AmqpConsumeContext;
+    if (data.self && data._sentryPendingConsumer) {
+      removePendingConsumer(data.self, data._sentryPendingConsumer);
+    }
+  });
+}
+
+/**
+ * Records `consumerTag -> { noAck, queue }` for the dispatch hook. `consume`'s promise resolves too late,
+ * because amqplib can decode `BasicConsumeOk` and the first delivery in one socket read. amqplib always
+ * registers a consumer before it dispatches to it.
+ */
+function subscribeRegisterConsumer(): void {
+  const channel = diagnosticsChannel.tracingChannel<AmqpChannelContext>(CHANNELS.AMQPLIB_REGISTER_CONSUMER);
+
+  channel.start.subscribe(message => {
+    const data = message as AmqpChannelContext;
+    const consumerChannel = data.self;
+    const consumerTag = data.arguments[0];
+    const callback = data.arguments[1];
+    if (!consumerChannel || typeof consumerTag !== 'string') {
+      return;
+    }
+
+    const entry = consumerChannel[CHANNEL_PENDING_CONSUMERS]?.find(pending => pending.callback === callback);
+    if (!entry) {
+      return;
+    }
+
+    removePendingConsumer(consumerChannel, entry);
+    consumerChannel[CHANNEL_CONSUMER_INFO]?.set(consumerTag, entry.info);
   });
 }
 
+function removePendingConsumer(channel: ChannelLike, entry: PendingConsumer): void {
+  channel[CHANNEL_PENDING_CONSUMERS] = channel[CHANNEL_PENDING_CONSUMERS]?.filter(pending => pending !== entry);
+}
+
 /**
  * Per delivered message (`BaseChannel.prototype.dispatchMessage`): continues the producer's trace,
  * opens a CONSUMER span, and runs the user callback under it. Manual-ack consumers keep the span open
@@ -356,6 +406,7 @@ function ensureChannelState(channel: ChannelLike): void {
 
   channel[CHANNEL_SPANS_NOT_ENDED] = [];
   channel[CHANNEL_CONSUMER_INFO] = new Map();
+  channel[CHANNEL_PENDING_CONSUMERS] = [];
 
   const timer = setInterval(() => checkConsumeTimeoutOnChannel(channel), CONSUME_TIMEOUT_MS);
   timer.unref?.();