feat(sentry): support nested spans for console commands

src/sentry/src/Tracing/Aspect/CoroutineAspect.php

@@ -18,11 +18,11 @@
 use Hyperf\Di\Aop\ProceedingJoinPoint;
 use Hyperf\Engine\Coroutine as Co;
 use Sentry\SentrySdk;
+use Sentry\State\Scope;
 use Sentry\Tracing\SpanContext;
-use Sentry\Tracing\SpanStatus;
-use Throwable;
 
 use function FriendsOfHyperf\Sentry\startTransaction;
+use function FriendsOfHyperf\Sentry\trace;
 use function Hyperf\Coroutine\defer;
 use function Sentry\continueTrace;
 
@@ -52,76 +52,55 @@ public function process(ProceedingJoinPoint $proceedingJoinPoint)
 
         $callingOnFunction = CoroutineBacktraceHelper::foundCallingOnFunction();
 
-        // Only trace the top-level coroutine creation.
-        if (! $callingOnFunction) {
-            return $proceedingJoinPoint->process();
-        }
-
-        // Get the current transaction from the current scope.
-        $transaction = SentrySdk::getCurrentHub()->getTransaction();
-
-        // If there's no active transaction, skip tracing.
-        if (! $transaction?->getSampled()) {
-            return $proceedingJoinPoint->process();
-        }
+        return trace(
+            function (Scope $scope) use ($proceedingJoinPoint, $callingOnFunction) {
+                if ($callingOnFunction && $span = $scope->getSpan()) {
+                    $cid = Co::id();
+                    $keys = $this->keys;
+                    $callable = $proceedingJoinPoint->arguments['keys']['callable'];
+
+                    // Transfer the Sentry context to the new coroutine.
+                    $proceedingJoinPoint->arguments['keys']['callable'] = function () use ($callable, $span, $callingOnFunction, $cid, $keys) {
+                        SentrySdk::init(); // Ensure Sentry is initialized in the new coroutine.
+
+                        $from = Co::getContextFor($cid);
+                        $current = Co::getContextFor();
+
+                        foreach ($keys as $key) {
+                            if (isset($from[$key]) && ! isset($current[$key])) {
+                                $current[$key] = $from[$key];
+                            }
+                        }
+
+                        $transaction = startTransaction(
+                            continueTrace($span->toTraceparent(), $span->toBaggage())
+                                ->setName('coroutine')
+                                ->setOp('coroutine.prepare')
+                                ->setDescription($callingOnFunction)
+                                ->setOrigin('auto.coroutine')
+                        );
+
+                        defer(function () use ($transaction) {
+                            $transaction->finish();
+                            Integration::flushEvents();
+                        });
+
+                        return trace(
+                            fn () => $callable(),
+                            SpanContext::make()
+                                ->setOp('coroutine.execute')
+                                ->setDescription($callingOnFunction)
+                                ->setOrigin('auto.coroutine')
+                        );
+                    };
+                }
 
-        // Start a span for the coroutine creation.
-        $parent = $transaction->startChild(
+                return $proceedingJoinPoint->process();
+            },
             SpanContext::make()
                 ->setOp('coroutine.create')
                 ->setDescription($callingOnFunction)
                 ->setOrigin('auto.coroutine')
-                ->setData(['coroutine.id' => Co::id()])
         );
-        SentrySdk::getCurrentHub()->setSpan($parent);
-
-        $cid = Co::id();
-        $keys = $this->keys;
-        $callable = $proceedingJoinPoint->arguments['keys']['callable'];
-
-        // Transfer the Sentry context to the new coroutine.
-        $proceedingJoinPoint->arguments['keys']['callable'] = function () use ($callable, $parent, $callingOnFunction, $cid, $keys) {
-            $from = Co::getContextFor($cid);
-            $current = Co::getContextFor();
-
-            foreach ($keys as $key) {
-                if (isset($from[$key]) && ! isset($current[$key])) {
-                    $current[$key] = $from[$key];
-                }
-            }
-
-            $coTransaction = startTransaction(
-                continueTrace($parent->toTraceparent(), $parent->toBaggage())
-                    ->setName('coroutine')
-                    ->setOp('coroutine.execute')
-                    ->setDescription($callingOnFunction)
-                    ->setOrigin('auto.coroutine')
-            );
-
-            defer(function () use ($coTransaction) {
-                // Set the transaction on the current scope to ensure it's the active one.
-                SentrySdk::getCurrentHub()->setSpan($coTransaction);
-
-                // Finish the transaction when the coroutine ends.
-                $coTransaction->finish();
-
-                // Flush events
-                Integration::flushEvents();
-            });
-
-            try {
-                $callable();
-            } catch (Throwable $exception) {
-                $coTransaction->setStatus(SpanStatus::internalError());
-
-                throw $exception;
-            }
-        };
-
-        try {
-            return $proceedingJoinPoint->process();
-        } finally {
-            $parent->finish();
-        }
     }
 }

src/sentry/src/Tracing/Listener/EventHandleListener.php

@@ -16,6 +16,7 @@
 use FriendsOfHyperf\Sentry\Feature;
 use FriendsOfHyperf\Sentry\Integration;
 use FriendsOfHyperf\Sentry\Util\Carrier;
+use FriendsOfHyperf\Sentry\Util\CoContainer;
 use FriendsOfHyperf\Sentry\Util\SqlParser;
 use FriendsOfHyperf\Support\RedisCommand;
 use Hyperf\Amqp\Event as AmqpEvent;
@@ -43,6 +44,7 @@
 use Psr\Container\ContainerInterface;
 use Sentry\SentrySdk;
 use Sentry\State\Scope;
+use Sentry\Tracing\Span;
 use Sentry\Tracing\SpanContext;
 use Sentry\Tracing\SpanStatus;
 use Sentry\Tracing\TransactionContext;
@@ -376,61 +378,85 @@ protected function handleCommandStarting(CommandEvent\BeforeHandle $event): void
 
         $command = $event->getCommand();
 
-        $transaction = startTransaction(
-            TransactionContext::make()
-                ->setName($command->getName() ?: '<unnamed command>')
-                ->setOp('console.command')
-                ->setDescription($command->getDescription())
-                ->setOrigin('auto.console')
+        if (! $parentSpan = SentrySdk::getCurrentHub()->getSpan()) {
+            $parentSpan = startTransaction(
+                TransactionContext::make()
+                    ->setName($command->getName() ?: '<unnamed command>')
+                    ->setOp('console.command')
+                    ->setDescription($command->getDescription())
+                    ->setOrigin('auto.console')
+                    ->setData([
+                        'command.arguments' => (fn () => $this->input->getArguments())->call($command),
+                        'command.options' => (fn () => $this->input->getOptions())->call($command),
+                    ])
+                    ->setTags([
+                        'command.name' => $command->getName(),
+                    ])
+                    ->setSource(TransactionSource::task())
+            );
+
+            CoContainer::set($command, $parentSpan);
+        }
+
+        if (! $parentSpan->getSampled()) {
+            CoContainer::del($command);
+            return;
+        }
+
+        $scope = SentrySdk::getCurrentHub()->pushScope();
+        $span = $parentSpan->startChild(
+            SpanContext::make()
+                ->setOp('console.command.execute')
+                ->setDescription($command->getName() ?: $command->getDescription())
                 ->setData([
                     'command.arguments' => (fn () => $this->input->getArguments())->call($command),
                     'command.options' => (fn () => $this->input->getOptions())->call($command),
                 ])
                 ->setTags([
                     'command.name' => $command->getName(),
                 ])
-                ->setSource(TransactionSource::task())
         );
-
-        Coroutine::inCoroutine() && defer(function () use ($transaction) {
-            // Make sure the transaction is finished after the command is executed
-            SentrySdk::getCurrentHub()->setSpan($transaction);
-
-            // Finish transaction
-            $transaction->finish();
-
-            // Flush events
-            Integration::flushEvents();
-        });
+        $scope->setSpan($span);
     }
 
     protected function handleCommandFinished(CommandEvent\AfterExecute $event): void
     {
-        $transaction = SentrySdk::getCurrentHub()->getTransaction();
+        $span = SentrySdk::getCurrentHub()->getSpan();
 
-        if (! $transaction?->getSampled()) {
+        if (! $span instanceof Span) {
             return;
         }
 
+        $command = $event->getCommand();
+        $sampled = $span->getSampled();
+
         try {
-            $command = $event->getCommand();
+            if (! $sampled) {
+                return;
+            }
+
             /** @var int $exitCode */
             $exitCode = (fn () => $this->exitCode ?? SymfonyCommand::SUCCESS)->call($command);
 
-            $transaction->setTags([
+            $span->setTags([
                 'command.exit_code' => (string) $exitCode,
-            ]);
-
-            $transaction->setStatus(
+            ])->setStatus(
                 $event->getThrowable() || $exitCode !== SymfonyCommand::SUCCESS
                 ? SpanStatus::internalError()
                 : SpanStatus::ok()
             );
         } finally {
-            if (! Coroutine::inCoroutine()) {
-                SentrySdk::getCurrentHub()->setSpan($transaction);
+            if ($sampled) {
+                $span->finish();
+                Integration::flushEvents();
+            }
+
+            $parentSpan = CoContainer::pull($command);
 
-                $transaction->finish();
+            if ($parentSpan instanceof Span) {
+                $parentSpan->finish();
+                SentrySdk::getCurrentHub()->setSpan($parentSpan);
+                SentrySdk::getCurrentHub()->popScope();
             }
         }
     }

src/sentry/src/Util/CoContainer.php

@@ -0,0 +1,75 @@
+<?php
+
+declare(strict_types=1);
+/**
+ * This file is part of friendsofhyperf/components.
+ *
+ * @link     https://github.com/friendsofhyperf/components
+ * @document https://github.com/friendsofhyperf/components/blob/main/README.md
+ * @contact  huangdijia@gmail.com
+ */
+
+namespace FriendsOfHyperf\Sentry\Util;
+
+use Hyperf\Context\Context;
+use WeakMap;
+
+/**
+ * @template TKey of object
+ * @template TValue of object
+ */
+class CoContainer
+{
+    public const CONTEXT_KEY = 'sentry.context.container';
+
+    /**
+     * @param TKey $key
+     * @param TValue $value
+     * @return TValue
+     */
+    public static function set(object $key, object $value): object
+    {
+        self::getContainer()[$key] = $value;
+
+        return $value;
+    }
+
+    /**
+     * @param TKey $key
+     * @return null|TValue
+     */
+    public static function get(object $key): ?object
+    {
+        $container = self::getContainer();
+
+        return $container[$key] ?? null;
+    }
+
+    /**
+     * @param TKey $key
+     * @return null|TValue
+     */
+    public static function pull(object $key): ?object
+    {
+        $container = self::getContainer();
+
+        if (! isset($container[$key])) {
+            return null;
+        }
+
+        $value = $container[$key];
+        unset($container[$key]);
+
+        return $value;
+    }
+
+    public static function del(object $key): void
+    {
+        unset(self::getContainer()[$key]);
+    }
+
+    private static function getContainer(): WeakMap
+    {
+        return Context::getOrSet(self::CONTEXT_KEY, fn () => new WeakMap());
+    }
+}