refactor: Pull out withActiveSpan to ensure we always end the span regardless of exception

Signed-off-by: Mike Bryant <mike@mikebryant.me.uk>

plugins/catalog-backend/src/processing/DefaultCatalogProcessingEngine.ts

@@ -23,7 +23,7 @@ import { assertError, serializeError, stringifyError } from '@backstage/errors';
 import { Hash } from 'crypto';
 import stableStringify from 'fast-json-stable-stringify';
 import { Logger } from 'winston';
-import { metrics, SpanStatusCode, trace } from '@opentelemetry/api';
+import { metrics, trace } from '@opentelemetry/api';
 import { ProcessingDatabase, RefreshStateItem } from '../database/types';
 import { createCounterMetric, createSummaryMetric } from '../util/metrics';
 import {
@@ -35,7 +35,11 @@ import { Stitcher } from '../stitching/Stitcher';
 import { startTaskPipeline } from './TaskPipeline';
 import { PluginTaskScheduler } from '@backstage/backend-tasks';
 import { Config } from '@backstage/config';
-import { addEntityAttributes, TRACER_ID } from '../util/opentelemetry';
+import {
+  addEntityAttributes,
+  TRACER_ID,
+  withActiveSpan,
+} from '../util/opentelemetry';
 
 const CACHE_TTL = 5;
 
@@ -134,7 +138,7 @@ export class DefaultCatalogProcessingEngine implements CatalogProcessingEngine {
         }
       },
       processTask: async item => {
-        await tracer.startActiveSpan('ProcessingRun', async span => {
+        await withActiveSpan(tracer, 'ProcessingRun', async span => {
           const track = this.tracker.processStart(item, this.logger);
           addEntityAttributes(span, item.entityRef);
 
@@ -157,7 +161,8 @@ export class DefaultCatalogProcessingEngine implements CatalogProcessingEngine {
             if (result.ok) {
               const { ttl: _, ...stateWithoutTtl } = state ?? {};
               if (
-                stableStringify(stateWithoutTtl) !== stableStringify(result.state)
+                stableStringify(stateWithoutTtl) !==
+                stableStringify(result.state)
               ) {
                 await this.processingDatabase.transaction(async tx => {
                   await this.processingDatabase.updateEntityCache(tx, {
@@ -216,7 +221,6 @@ export class DefaultCatalogProcessingEngine implements CatalogProcessingEngine {
               // significant effect on our surroundings; therefore, we just abort
               // without any updates / stitching.
               track.markSuccessfulWithNoChanges();
-              span.end();
               return;
             }
 
@@ -255,8 +259,6 @@ export class DefaultCatalogProcessingEngine implements CatalogProcessingEngine {
                 new Set([stringifyEntityRef(unprocessedEntity)]),
               );
               track.markSuccessfulWithErrors();
-              span.setStatus({ code: SpanStatusCode.ERROR });
-              span.end();
               return;
             }
 
@@ -309,10 +311,7 @@ export class DefaultCatalogProcessingEngine implements CatalogProcessingEngine {
           } catch (error) {
             assertError(error);
             track.markFailed(error);
-            span.recordException(error);
-            span.setStatus({ code: SpanStatusCode.ERROR });
           }
-          span.end();
         });
       },
     });

plugins/catalog-backend/src/processing/DefaultCatalogProcessingOrchestrator.ts

@@ -56,7 +56,11 @@ import {
 } from './util';
 import { CatalogRulesEnforcer } from '../ingestion/CatalogRules';
 import { ProcessorCacheManager } from './ProcessorCacheManager';
-import { addEntityAttributes, TRACER_ID } from '../util/opentelemetry';
+import {
+  addEntityAttributes,
+  TRACER_ID,
+  withActiveSpan,
+} from '../util/opentelemetry';
 
 const tracer = trace.getTracer(TRACER_ID);
 
@@ -200,7 +204,7 @@ export class DefaultCatalogProcessingOrchestrator
     for (const processor of this.options.processors) {
       if (processor.preProcessEntity) {
         let innerRes = res;
-        res = await tracer.startActiveSpan('ProcessingStep', async span => {
+        res = await withActiveSpan(tracer, 'ProcessingStep', async span => {
           addEntityAttributes(span, context.entityRef);
           addProcessorAttributes(span, 'preProcessEntity', processor);
           try {
@@ -212,15 +216,11 @@ export class DefaultCatalogProcessingOrchestrator
               context.cache.forProcessor(processor),
             );
           } catch (e) {
-            span.recordException(e);
-            span.setStatus({ code: SpanStatusCode.ERROR });
-            span.end();
             throw new InputError(
               `Processor ${processor.constructor.name} threw an error while preprocessing`,
               e,
             );
           }
-          span.end();
           return innerRes;
         });
       }
@@ -388,7 +388,7 @@ export class DefaultCatalogProcessingOrchestrator
     for (const processor of this.options.processors) {
       if (processor.postProcessEntity) {
         let innerRes = res;
-        res = await tracer.startActiveSpan('ProcessingStep', async span => {
+        res = await withActiveSpan(tracer, 'ProcessingStep', async span => {
           addEntityAttributes(span, context.entityRef);
           addProcessorAttributes(span, 'postProcessEntity', processor);
           try {
@@ -399,15 +399,11 @@ export class DefaultCatalogProcessingOrchestrator
               context.cache.forProcessor(processor),
             );
           } catch (e) {
-            span.recordException(e);
-            span.setStatus({ code: SpanStatusCode.ERROR });
-            span.end();
             throw new InputError(
               `Processor ${processor.constructor.name} threw an error while postprocessing`,
               e,
             );
           }
-          span.end();
           return innerRes;
         });
       }

plugins/catalog-backend/src/util/opentelemetry.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { Span, SpanStatusCode } from '@opentelemetry/api';
+import { Span, SpanOptions, SpanStatusCode, Tracer } from '@opentelemetry/api';
 import { parseEntityRef } from '@backstage/catalog-model';
 
 export const TRACER_ID = 'backstage-plugin-catalog-backend';
@@ -30,3 +30,48 @@ export function addEntityAttributes(span: Span, entityRef: string) {
     span.setStatus({ code: SpanStatusCode.ERROR });
   }
 }
+
+// Adapted from https://github.com/open-telemetry/opentelemetry-js/blob/359fbcc40a859057a02b14e84599eac399b8dba7/api/src/trace/SugaredTracer.ts
+// While waiting for something like https://github.com/open-telemetry/opentelemetry-js/pull/3317 to land upstream
+
+const onException = (e: Error, span: Span) => {
+  span.recordException(e);
+  span.setStatus({
+    code: SpanStatusCode.ERROR,
+  });
+};
+
+function handleFn<F extends (span: Span) => ReturnType<F>>(
+  span: Span,
+  fn: F,
+): ReturnType<F> {
+  try {
+    const ret = fn(span) as Promise<ReturnType<F>>;
+    // if fn is an async function attach a recordException and spanEnd callback to the promise
+    if (typeof ret.then === 'function' && typeof ret.catch === 'function') {
+      return ret
+        .catch((e: Error) => {
+          onException(e, span);
+          throw e;
+        })
+        .finally(() => span.end()) as ReturnType<F>;
+    }
+    span.end();
+    return ret as ReturnType<F>;
+  } catch (e) {
+    onException(e, span);
+    span.end();
+    throw e;
+  }
+}
+
+export function withActiveSpan<F extends (span: Span) => ReturnType<F>>(
+  tracer: Tracer,
+  name: string,
+  fn: F,
+  spanOptions: SpanOptions = {},
+): ReturnType<F> {
+  return tracer.startActiveSpan(name, spanOptions, (span: Span) => {
+    return handleFn(span, fn);
+  });
+}