feat(createTimedSpan): Pass along tags from afterCompletion

src/createTimedSpan.test.ts

@@ -66,7 +66,7 @@ describe('timedSpan', () => {
     expect(duration).toBeGreaterThan(0);
   });
 
-  it('should pass along tags', async () => {
+  it('should pass along tags from `timedSpan`', async () => {
     const mockIncrement = jest.spyOn(metricsClient, 'increment');
     const mockTiming = jest.spyOn(metricsClient, 'timing');
 
@@ -90,6 +90,31 @@ describe('timedSpan', () => {
     );
   });
 
+  it('should pass along tags from `afterCompletion`', async () => {
+    const mockIncrement = jest.spyOn(metricsClient, 'increment');
+    const mockTiming = jest.spyOn(metricsClient, 'timing');
+
+    await timedSpan(
+      'test',
+      // This is false but we still successfully resolved
+      () => Promise.resolve(false),
+      () => ({ tags: ['tags-from-after-completion'] }),
+      ['new-tags'],
+    );
+
+    expect(mockIncrement).toHaveBeenCalledWith('test.count', [
+      'success',
+      'new-tags',
+      'tags-from-after-completion',
+    ]);
+
+    expect(mockTiming).toHaveBeenCalledWith(
+      'test.latency',
+      expect.any(Number),
+      ['new-tags', 'tags-from-after-completion'],
+    );
+  });
+
   it('should handle failure', async () => {
     const mockIncrement = jest.spyOn(metricsClient, 'increment');
     const mockTiming = jest.spyOn(metricsClient, 'timing');

src/createTimedSpan.ts

@@ -2,6 +2,10 @@ import type { MetricsClient } from './MetricsClient';
 
 type TimingMetricsClient = Pick<MetricsClient, 'increment' | 'timing'>;
 
+interface AfterCompletion {
+  tags: string[];
+}
+
 /**
  * Sends timing related metrics for an asynchronous operation
  *
@@ -18,28 +22,36 @@ export const createTimedSpan =
   async <T>(
     name: string,
     block: () => PromiseLike<T>,
-    afterCompletion?: (duration: number, success: boolean) => void,
+    afterCompletion?: (
+      duration: number,
+      success: boolean,
+      result: T | undefined,
+    ) => AfterCompletion | void,
     tags?: string[],
   ): Promise<T> => {
     const startTime = process.hrtime.bigint();
 
-    const handleCompletion = (success: boolean) => {
+    const handleCompletion = (success: boolean, result: T | undefined) => {
       const durationNanos = process.hrtime.bigint() - startTime;
       const successTag = success ? 'success' : 'failure';
       const durationMilliseconds = Number(durationNanos) / 1e6;
 
+      const complete = afterCompletion?.(durationMilliseconds, success, result);
+
+      if (complete?.tags) {
+        (tags ?? []).push(...complete.tags);
+      }
+
       metricsClient.timing(`${name}.latency`, durationMilliseconds, tags);
       metricsClient.increment(`${name}.count`, [successTag, ...(tags ?? [])]);
-
-      afterCompletion?.(durationMilliseconds, success);
     };
 
     try {
       const result = await block();
-      handleCompletion(true);
+      handleCompletion(true, result);
       return result;
     } catch (e) {
-      handleCompletion(false);
+      handleCompletion(false, undefined);
       throw e;
     }
   };