fix(Instrumentation/Messenger): remove TraceableStamp usage from TraceableMiddleware and TraceableStack (#165)

* fix(Instrumentation/Messenger): remove TraceableStamp usage from TraceableMiddleware and TraceableStack

* chore: format

* test(Instrumentation): remove messenger doctrine transport test

Author: gaelreyrol

src/Instrumentation/Symfony/Messenger/TraceableMessengerMiddleware.php

@@ -2,10 +2,7 @@
 
 namespace FriendsOfOpenTelemetry\OpenTelemetryBundle\Instrumentation\Symfony\Messenger;
 
-use OpenTelemetry\API\Trace\SpanKind;
-use OpenTelemetry\API\Trace\StatusCode;
 use OpenTelemetry\API\Trace\TracerInterface;
-use OpenTelemetry\Context\Context;
 use Psr\Log\LoggerInterface;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Middleware\MiddlewareInterface;
@@ -23,51 +20,18 @@ public function __construct(
 
     public function handle(Envelope $envelope, StackInterface $stack): Envelope
     {
-        $scope = Context::storage()->scope();
-        if (null !== $scope) {
-            $this->logger?->debug(sprintf('Using scope "%s"', spl_object_id($scope)));
-        } else {
-            $this->logger?->debug('No active scope');
-        }
-
-        $traceableStamp = $this->getTraceableStamp($envelope);
-        if (null !== $traceableStamp && $traceableStamp->getSpan()->isRecording()) {
-            $span = $traceableStamp->getSpan();
-            $span->setStatus(StatusCode::STATUS_OK);
-            $this->logger?->debug(sprintf('Ending span "%s"', $span->getContext()->getSpanId()));
-            $span->end();
-        }
-
-        $spanBuilder = $this->tracer
-            ->spanBuilder('messenger.middleware')
-            ->setSpanKind(SpanKind::KIND_INTERNAL)
-            ->setParent($scope?->context())
-            ->setAttribute('event.category', $this->eventCategory)
-            ->setAttribute('bus.name', $this->busName)
-        ;
-
-        $span = $spanBuilder->startSpan();
-
-        $this->logger?->debug(sprintf('Starting span "%s"', $span->getContext()->getSpanId()));
-
         $stack = new TraceableMessengerStack(
-            $span,
+            $this->tracer,
             $stack,
             $this->busName,
+            $this->eventCategory,
             $this->logger,
         );
 
-        $envelope = $envelope->with(new TraceableStamp($span));
-
         try {
             return $stack->next()->handle($envelope, $stack);
         } finally {
             $stack->stop();
         }
     }
-
-    private function getTraceableStamp(Envelope $envelope): ?TraceableStamp
-    {
-        return $envelope->last(TraceableStamp::class);
-    }
 }

src/Instrumentation/Symfony/Messenger/TraceableMessengerStack.php

@@ -2,8 +2,11 @@
 
 namespace FriendsOfOpenTelemetry\OpenTelemetryBundle\Instrumentation\Symfony\Messenger;
 
-use OpenTelemetry\API\Trace\SpanInterface;
+use OpenTelemetry\API\Trace\SpanKind;
 use OpenTelemetry\API\Trace\StatusCode;
+use OpenTelemetry\API\Trace\TracerInterface;
+use OpenTelemetry\Context\Context;
+use OpenTelemetry\SDK\Trace\Span;
 use Psr\Log\LoggerInterface;
 use Symfony\Component\Messenger\Middleware\MiddlewareInterface;
 use Symfony\Component\Messenger\Middleware\StackInterface;
@@ -13,32 +16,77 @@ class TraceableMessengerStack implements StackInterface
     private ?string $currentEvent = null;
 
     public function __construct(
-        private SpanInterface $span,
+        private TracerInterface $tracer,
         private StackInterface $stack,
         private string $busName,
+        private string $eventCategory,
         private ?LoggerInterface $logger = null,
     ) {
     }
 
     public function next(): MiddlewareInterface
     {
+        $scope = Context::storage()->scope();
+        if (null !== $scope) {
+            $this->logger?->debug(sprintf('Using scope "%s"', spl_object_id($scope)));
+        } else {
+            $this->logger?->debug('No active scope');
+        }
+
+        /*        if (null !== $scope) {
+                    $span = Span::fromContext($scope->context());
+
+                    if ($span->isRecording()) {
+                        $scope->detach();
+
+                        $span->setStatus(StatusCode::STATUS_OK);
+                        $this->logger?->debug(sprintf('Ending span "%s"', $span->getContext()->getSpanId()));
+                        $span->end();
+                    }
+                }*/
+
+        $spanBuilder = $this->tracer
+            ->spanBuilder('messenger.middleware')
+            ->setSpanKind(SpanKind::KIND_INTERNAL)
+            ->setParent($scope?->context())
+            ->setAttribute('event.category', $this->eventCategory)
+            ->setAttribute('bus.name', $this->busName)
+        ;
+
+        $parent = Context::getCurrent();
+
+        $span = $spanBuilder->setParent($parent)->startSpan();
+
+        $this->logger?->debug(sprintf('Starting span "%s"', $span->getContext()->getSpanId()));
+
         if ($this->stack === $nextMiddleware = $this->stack->next()) {
             $this->currentEvent = 'Tail';
         } else {
             $this->currentEvent = sprintf('"%s"', get_debug_type($nextMiddleware));
         }
         $this->currentEvent .= sprintf(' on "%s"', $this->busName);
 
-        $this->span->setAttribute('event.current', $this->currentEvent);
+        $span->setAttribute('event.current', $this->currentEvent);
+
+        $context = $span->storeInContext($parent);
+        Context::storage()->attach($context);
 
         return $nextMiddleware;
     }
 
     public function stop(): void
     {
-        $this->span->setStatus(StatusCode::STATUS_OK);
-        $this->logger?->debug(sprintf('Ending span "%s"', $this->span->getContext()->getSpanId()));
-        $this->span->end();
+        $scope = Context::storage()->scope();
+        if (null === $scope) {
+            return;
+        }
+
+        $scope->detach();
+
+        $span = Span::fromContext($scope->context());
+        $span->setStatus(StatusCode::STATUS_OK);
+        $this->logger?->debug(sprintf('Ending span "%s"', $span->getContext()->getSpanId()));
+        $span->end();
         $this->currentEvent = null;
     }