From db2cac67450cd75c95ce1ed761b2c6f0993f7bde Mon Sep 17 00:00:00 2001
From: Mike Bryant <mike@mikebryant.me.uk>
Date: Thu, 18 May 2023 21:06:51 +0100
Subject: [PATCH] feat: Add ProcessingStage span

Signed-off-by: Mike Bryant <mike@mikebryant.me.uk>
---
 .../DefaultCatalogProcessingOrchestrator.ts   | 112 ++++++++++--------
 1 file changed, 63 insertions(+), 49 deletions(-)

diff --git a/plugins/catalog-backend/src/processing/DefaultCatalogProcessingOrchestrator.ts b/plugins/catalog-backend/src/processing/DefaultCatalogProcessingOrchestrator.ts
index 8091c6f120..8b4f4127f3 100644
--- a/plugins/catalog-backend/src/processing/DefaultCatalogProcessingOrchestrator.ts
+++ b/plugins/catalog-backend/src/processing/DefaultCatalogProcessingOrchestrator.ts
@@ -199,34 +199,41 @@ export class DefaultCatalogProcessingOrchestrator
     entity: Entity,
     context: Context,
   ): Promise<Entity> {
-    let res = entity;
+    return await withActiveSpan(tracer, 'ProcessingStage', async stageSpan => {
+      addEntityAttributes(stageSpan, context.entityRef);
+      stageSpan.setAttribute(
+        'backstage.catalog.processor.stage',
+        'preProcessEntity',
+      );
+      let res = entity;
 
-    for (const processor of this.options.processors) {
-      if (processor.preProcessEntity) {
-        let innerRes = res;
-        res = await withActiveSpan(tracer, 'ProcessingStep', async span => {
-          addEntityAttributes(span, context.entityRef);
-          addProcessorAttributes(span, 'preProcessEntity', processor);
-          try {
-            innerRes = await processor.preProcessEntity!(
-              innerRes,
-              context.location,
-              context.collector.forProcessor(processor),
-              context.originLocation,
-              context.cache.forProcessor(processor),
-            );
-          } catch (e) {
-            throw new InputError(
-              `Processor ${processor.constructor.name} threw an error while preprocessing`,
-              e,
-            );
-          }
-          return innerRes;
-        });
+      for (const processor of this.options.processors) {
+        if (processor.preProcessEntity) {
+          let innerRes = res;
+          res = await withActiveSpan(tracer, 'ProcessingStep', async span => {
+            addEntityAttributes(span, context.entityRef);
+            addProcessorAttributes(span, 'preProcessEntity', processor);
+            try {
+              innerRes = await processor.preProcessEntity!(
+                innerRes,
+                context.location,
+                context.collector.forProcessor(processor),
+                context.originLocation,
+                context.cache.forProcessor(processor),
+              );
+            } catch (e) {
+              throw new InputError(
+                `Processor ${processor.constructor.name} threw an error while preprocessing`,
+                e,
+              );
+            }
+            return innerRes;
+          });
+        }
       }
-    }
 
-    return res;
+      return res;
+    });
   }
 
   /**
@@ -383,32 +390,39 @@ export class DefaultCatalogProcessingOrchestrator
     entity: Entity,
     context: Context,
   ): Promise<Entity> {
-    let res = entity;
+    return await withActiveSpan(tracer, 'ProcessingStage', async stageSpan => {
+      addEntityAttributes(stageSpan, context.entityRef);
+      stageSpan.setAttribute(
+        'backstage.catalog.processor.stage',
+        'postProcessEntity',
+      );
+      let res = entity;
 
-    for (const processor of this.options.processors) {
-      if (processor.postProcessEntity) {
-        let innerRes = res;
-        res = await withActiveSpan(tracer, 'ProcessingStep', async span => {
-          addEntityAttributes(span, context.entityRef);
-          addProcessorAttributes(span, 'postProcessEntity', processor);
-          try {
-            innerRes = await processor.postProcessEntity!(
-              innerRes,
-              context.location,
-              context.collector.forProcessor(processor),
-              context.cache.forProcessor(processor),
-            );
-          } catch (e) {
-            throw new InputError(
-              `Processor ${processor.constructor.name} threw an error while postprocessing`,
-              e,
-            );
-          }
-          return innerRes;
-        });
+      for (const processor of this.options.processors) {
+        if (processor.postProcessEntity) {
+          let innerRes = res;
+          res = await withActiveSpan(tracer, 'ProcessingStep', async span => {
+            addEntityAttributes(span, context.entityRef);
+            addProcessorAttributes(span, 'postProcessEntity', processor);
+            try {
+              innerRes = await processor.postProcessEntity!(
+                innerRes,
+                context.location,
+                context.collector.forProcessor(processor),
+                context.cache.forProcessor(processor),
+              );
+            } catch (e) {
+              throw new InputError(
+                `Processor ${processor.constructor.name} threw an error while postprocessing`,
+                e,
+              );
+            }
+            return innerRes;
+          });
+        }
       }
-    }
 
-    return res;
+      return res;
+    });
   }
 }