fix(Aws): fixes wrong aws client name and region binding when appending middlewares

src/Aws/phpunit.xml.dist

@@ -39,6 +39,9 @@
         <testsuite name="unit">
             <directory>tests/Unit</directory>
         </testsuite>
+        <testsuite name="integration">
+            <directory>tests/Integration</directory>
+        </testsuite>
     </testsuites>
 
 </phpunit>

src/Aws/src/AwsSdkInstrumentation.php

@@ -8,12 +8,10 @@
 use Aws\ResultInterface;
 use OpenTelemetry\API\Instrumentation\InstrumentationInterface;
 use OpenTelemetry\API\Instrumentation\InstrumentationTrait;
-use OpenTelemetry\API\Trace\SpanInterface;
 use OpenTelemetry\API\Trace\SpanKind;
 use OpenTelemetry\API\Trace\TracerInterface;
 use OpenTelemetry\API\Trace\TracerProviderInterface;
 use OpenTelemetry\Context\Propagation\TextMapPropagatorInterface;
-use OpenTelemetry\Context\ScopeInterface;
 
 /**
  * @experimental
@@ -25,13 +23,12 @@ class AwsSdkInstrumentation implements InstrumentationInterface
     public const NAME = 'AWS SDK Instrumentation';
     public const VERSION = '0.0.1';
     public const SPAN_KIND = SpanKind::KIND_CLIENT;
-    private TextMapPropagatorInterface $propagator;
-    private TracerProviderInterface $tracerProvider;
-    private $clients = [] ;
-    private string $clientName;
-    private string $region;
-    private SpanInterface $span;
-    private ScopeInterface $scope;
+
+    private array $clients = [];
+
+    private array $instrumentedClients = [];
+
+    private array $spanStorage = [];
 
     public function getName(): string
     {
@@ -79,61 +76,70 @@ public function getTracer(): TracerInterface
     }
 
     /** @psalm-api */
-    public function instrumentClients($clientsArray) : void
+    public function instrumentClients($clientsArray): void
     {
         $this->clients = $clientsArray;
     }
 
-    /** @psalm-suppress ArgumentTypeCoercion */
     public function activate(): bool
     {
         try {
-            $middleware = Middleware::tap(function ($cmd, $_req) {
-                $tracer = $this->getTracer();
-                $propagator = $this->getPropagator();
-
-                $carrier = [];
-                /** @phan-suppress-next-line PhanTypeMismatchArgument */
-                $this->span = $tracer->spanBuilder($this->clientName)->setSpanKind(AwsSdkInstrumentation::SPAN_KIND)->startSpan();
-                $this->scope = $this->span->activate();
-
-                $propagator->inject($carrier);
-
-                /** @psalm-suppress PossiblyInvalidArgument */
-                $this->span->setAttributes([
-                    'rpc.method' => $cmd->getName(),
-                    'rpc.service' => $this->clientName,
-                    'rpc.system' => 'aws-api',
-                    'aws.region' => $this->region,
-                    ]);
-            });
-
-            /** @psalm-suppress PossiblyInvalidArgument */
-            $end_middleware = Middleware::mapResult(function (ResultInterface $result) {
-                /**
-                 * Some AWS SDK Funtions, such as S3Client->getObjectUrl() do not actually perform on the wire comms
-                 * with AWS Servers, and therefore do not return with a populated AWS\Result object with valid @metadata
-                 * Check for the presence of @metadata before extracting status code as these calls are still
-                 * instrumented.
-                 */
-                if (isset($result['@metadata'])) {
-                    $this->span->setAttributes([
-                        'http.status_code' => $result['@metadata']['statusCode'], //@phan-suppress-current-line PhanTypeMismatchDimFetch
-                    ]);
+            foreach ($this->clients as $client) {
+                $hash = spl_object_hash($client);
+                if (isset($this->instrumentedClients[$hash])) {
+                    continue;
                 }
 
-                $this->span->end();
-                $this->scope->detach();
+                $clientName = $client->getApi()->getServiceName();
+                $region = $client->getRegion();
 
-                return $result;
-            });
+                $client->getHandlerList()->prependInit(Middleware::tap(function ($cmd, $_req) use ($clientName, $region, $hash) {
+                    $tracer = $this->getTracer();
+                    $propagator = $this->getPropagator();
 
-            foreach ($this->clients as $client) {
-                $this->clientName = $client->getApi()->getServiceName();
-                $this->region = $client->getRegion();
+                    $carrier = [];
+                    /** @phan-suppress-next-line PhanTypeMismatchArgument */
+                    $span = $tracer->spanBuilder($clientName)->setSpanKind(AwsSdkInstrumentation::SPAN_KIND)->startSpan();
+                    $scope = $span->activate();
+                    $this->spanStorage[$hash] = [$span, $scope];
 
-                $client->getHandlerList()->prependInit($middleware, 'instrumentation');
-                $client->getHandlerList()->appendSign($end_middleware, 'end_instrumentation');
+                    $propagator->inject($carrier);
+
+                    /** @psalm-suppress PossiblyInvalidArgument */
+                    $span->setAttributes([
+                        'rpc.method' => $cmd->getName(),
+                        'rpc.service' => $clientName,
+                        'rpc.system' => 'aws-api',
+                        'aws.region' => $region,
+                    ]);
+                }), 'instrumentation');
+
+                $client->getHandlerList()->appendSign(Middleware::mapResult(function (ResultInterface $result) use ($hash) {
+                    if (empty($this->spanStorage[$hash])) {
+                        return $result;
+                    }
+                    [$span, $scope] = $this->spanStorage[$hash];
+                    unset($this->spanStorage[$hash]);
+
+                    /*
+                     * Some AWS SDK Functions, such as S3Client->getObjectUrl() do not actually perform on the wire comms
+                     * with AWS Servers, and therefore do not return with a populated AWS\Result object with valid @metadata
+                     * Check for the presence of @metadata before extracting status code as these calls are still
+                     * instrumented.
+                     */
+                    if (isset($result['@metadata'])) {
+                        $span->setAttributes([
+                            'http.status_code' => $result['@metadata']['statusCode'], // @phan-suppress-current-line PhanTypeMismatchDimFetch
+                        ]);
+                    }
+
+                    $span->end();
+                    $scope->detach();
+
+                    return $result;
+                }), 'end_instrumentation');
+
+                $this->instrumentedClients[$hash] = 1;
             }
         } catch (\Throwable $e) {
             return false;

src/Aws/tests/Integration/AwsSdkInstrumentationTest.php

@@ -0,0 +1,216 @@
+<?php
+
+declare(strict_types=1);
+
+namespace OpenTelemetry\Tests\Aws\Integration;
+
+use Aws\AwsClientInterface;
+use Aws\EventBridge\EventBridgeClient;
+use Aws\S3\S3Client;
+use Aws\Sqs\SqsClient;
+use OpenTelemetry\Aws\AwsSdkInstrumentation;
+use OpenTelemetry\Aws\Xray\Propagator;
+use OpenTelemetry\SDK\Trace\ReadWriteSpanInterface;
+use OpenTelemetry\SDK\Trace\TracerProvider;
+use PHPUnit\Framework\TestCase;
+
+class AwsSdkInstrumentationTest extends TestCase
+{
+    use UsesServiceTrait;
+
+    private const HANDLERS_PER_ACTIVATION = 2; // one init and one sign middleware
+
+    private AwsSdkInstrumentation $awsSdkInstrumentation;
+
+    protected function setUp(): void
+    {
+        $this->awsSdkInstrumentation = new AwsSdkInstrumentation();
+    }
+
+    public function testProperClientNameAndRegionIsPassedToSpanForSingleClientCall()
+    {
+        /** @var SqsClient $sqsClient */
+        $sqsClient = $this->getTestClient('SQS', ['region' => 'eu-west-1']);
+        /** @var S3Client $s3Client */
+        $s3Client = $this->getTestClient('S3', ['region' => 'us-east-1']);
+        $this->addMockResults($s3Client, [[]]);
+        /** @var EventBridgeClient $eventBridgeClient */
+        $eventBridgeClient = $this->getTestClient('EventBridge', ['region' => 'ap-southeast-2']);
+
+        $spanProcessor = new CollectingSpanProcessor();
+        $this->awsSdkInstrumentation->instrumentClients([$sqsClient, $s3Client, $eventBridgeClient]);
+        $this->awsSdkInstrumentation->setPropagator(new Propagator());
+        $this->awsSdkInstrumentation->setTracerProvider(new TracerProvider([$spanProcessor]));
+        $this->awsSdkInstrumentation->init();
+        $this->awsSdkInstrumentation->activate();
+
+        $s3Client->listBuckets();
+
+        $collectedSpans = $spanProcessor->getCollectedSpans();
+        $this->assertCount(1, $collectedSpans);
+
+        /** @var ReadWriteSpanInterface $span */
+        $span = reset($collectedSpans);
+        $this->assertTrue($span->hasEnded());
+
+        $attributes = $span->toSpanData()->getAttributes()->toArray();
+        $this->assertArrayHasKey('rpc.service', $attributes);
+        $this->assertSame('s3', $attributes['rpc.service']);
+        $this->assertArrayHasKey('aws.region', $attributes);
+        $this->assertSame('us-east-1', $attributes['aws.region']);
+    }
+
+    public function testProperClientNameAndRegionIsPassedToSpanForDoubleCallToSameClient()
+    {
+        /** @var SqsClient $sqsClient */
+        $sqsClient = $this->getTestClient('SQS', ['region' => 'eu-west-1']);
+        /** @var S3Client $s3Client */
+        $s3Client = $this->getTestClient('S3', ['region' => 'us-east-1']);
+        $this->addMockResults($s3Client, [[], []]);
+        /** @var EventBridgeClient $eventBridgeClient */
+        $eventBridgeClient = $this->getTestClient('EventBridge', ['region' => 'ap-southeast-2']);
+
+        $spanProcessor = new CollectingSpanProcessor();
+        $this->awsSdkInstrumentation->instrumentClients([$sqsClient, $s3Client, $eventBridgeClient]);
+        $this->awsSdkInstrumentation->setPropagator(new Propagator());
+        $this->awsSdkInstrumentation->setTracerProvider(new TracerProvider([$spanProcessor]));
+        $this->awsSdkInstrumentation->init();
+        $this->awsSdkInstrumentation->activate();
+
+        $s3Client->listBuckets();
+        $s3Client->listObjects(['Bucket' => 'foo']);
+
+        $collectedSpans = $spanProcessor->getCollectedSpans();
+        $this->assertCount(2, $collectedSpans);
+
+        /** @var ReadWriteSpanInterface $span */
+        foreach ($collectedSpans as $span) {
+            $this->assertTrue($span->hasEnded());
+            $attributes = $span->toSpanData()->getAttributes()->toArray();
+            $this->assertArrayHasKey('rpc.service', $attributes);
+            $this->assertSame('s3', $attributes['rpc.service']);
+            $this->assertArrayHasKey('aws.region', $attributes);
+            $this->assertSame('us-east-1', $attributes['aws.region']);
+        }
+    }
+
+    public function testProperClientNameAndRegionIsPassedToSpanForDoubleCallToDifferentClients()
+    {
+        /** @var SqsClient $sqsClient */
+        $sqsClient = $this->getTestClient('SQS', ['region' => 'eu-west-1']);
+        /** @var S3Client $s3Client */
+        $s3Client = $this->getTestClient('S3', ['region' => 'us-east-1']);
+        $this->addMockResults($s3Client, [[]]);
+        /** @var EventBridgeClient $eventBridgeClient */
+        $eventBridgeClient = $this->getTestClient('EventBridge', ['region' => 'ap-southeast-2']);
+        $this->addMockResults($eventBridgeClient, [[]]);
+
+        $spanProcessor = new CollectingSpanProcessor();
+        $this->awsSdkInstrumentation->instrumentClients([$sqsClient, $s3Client, $eventBridgeClient]);
+        $this->awsSdkInstrumentation->setPropagator(new Propagator());
+        $this->awsSdkInstrumentation->setTracerProvider(new TracerProvider([$spanProcessor]));
+        $this->awsSdkInstrumentation->init();
+        $this->awsSdkInstrumentation->activate();
+
+        $eventBridgeClient->putEvents([
+            'Entries' => [
+                [
+                    'Version' => 1,
+                    'EventBusName' => 'foo',
+                    'Source' => 'bar',
+                    'DetailType' => 'type',
+                    'Detail' => '{}',
+                ],
+            ],
+        ]);
+        $s3Client->listBuckets();
+
+        $collectedSpans = $spanProcessor->getCollectedSpans();
+        $this->assertCount(2, $collectedSpans);
+
+        /** @var ReadWriteSpanInterface $span */
+        $span = array_pop($collectedSpans);
+        $this->assertTrue($span->hasEnded());
+        $attributes = $span->toSpanData()->getAttributes()->toArray();
+        $this->assertArrayHasKey('rpc.service', $attributes);
+        $this->assertSame('s3', $attributes['rpc.service']);
+        $this->assertArrayHasKey('aws.region', $attributes);
+        $this->assertSame('us-east-1', $attributes['aws.region']);
+
+        /** @var ReadWriteSpanInterface $span */
+        $span = array_pop($collectedSpans);
+        $this->assertTrue($span->hasEnded());
+        $attributes = $span->toSpanData()->getAttributes()->toArray();
+        $this->assertArrayHasKey('rpc.service', $attributes);
+        $this->assertSame('eventbridge', $attributes['rpc.service']);
+        $this->assertArrayHasKey('aws.region', $attributes);
+        $this->assertSame('ap-southeast-2', $attributes['aws.region']);
+    }
+
+    public function testSpansFromDifferentClientsAreNotOverwritingOneAnother()
+    {
+        try {
+            /** @var SqsClient $sqsClient */
+            $sqsClient = $this->getTestClient('SQS', ['region' => 'eu-west-1']);
+            $this->addMockResults($sqsClient, [[]]);
+            /** @var S3Client $s3Client */
+            $s3Client = $this->getTestClient('S3', ['region' => 'us-east-1']);
+            $this->addMockResults($s3Client, [[]]);
+
+            $spanProcessor = new CollectingSpanProcessor();
+            $this->awsSdkInstrumentation->instrumentClients([$sqsClient, $s3Client]);
+            $this->awsSdkInstrumentation->setPropagator(new Propagator());
+            $this->awsSdkInstrumentation->setTracerProvider(new TracerProvider([$spanProcessor]));
+            $this->awsSdkInstrumentation->init();
+            $this->awsSdkInstrumentation->activate();
+
+            $sqsClient->listQueuesAsync();
+            $s3Client->listBucketsAsync();
+
+            $collectedSpans = $spanProcessor->getCollectedSpans();
+            $this->assertCount(2, $collectedSpans);
+
+            /** @var ReadWriteSpanInterface $span */
+            $span = array_shift($collectedSpans);
+            $attributes = $span->toSpanData()->getAttributes()->toArray();
+            $this->assertArrayHasKey('rpc.service', $attributes);
+            $this->assertSame('sqs', $attributes['rpc.service']);
+
+            /** @var ReadWriteSpanInterface $span */
+            $span = array_shift($collectedSpans);
+            $attributes = $span->toSpanData()->getAttributes()->toArray();
+            $this->assertArrayHasKey('rpc.service', $attributes);
+            $this->assertSame('s3', $attributes['rpc.service']);
+        } catch (\Throwable $throwable) {
+            /** @phpstan-ignore-next-line  */
+            $this->assertFalse(true, sprintf('Exception %s occurred: %s', get_class($throwable), $throwable->getMessage()));
+        }
+    }
+
+    public function testPreventsRepeatedInstrumentationOfSameClient()
+    {
+        $clients = [
+            'SQS' => $sqsClient = $this->getTestClient('SQS', ['region' => 'eu-west-1']),
+            'S3' => $s3Client = $this->getTestClient('S3', ['region' => 'us-east-1']),
+            'EventBridge' => $eventBridgeClient = $this->getTestClient('EventBridge', ['region' => 'ap-southeast-2']),
+        ];
+
+        $preInstrumentationHandlersCount = array_map(static fn (AwsClientInterface $client) => $client->getHandlerList()->count(), $clients);
+
+        $this->awsSdkInstrumentation->instrumentClients([$sqsClient, $eventBridgeClient]);
+        $this->awsSdkInstrumentation->init();
+        $this->awsSdkInstrumentation->activate();
+
+        $this->awsSdkInstrumentation->instrumentClients([$s3Client, $eventBridgeClient]);
+        $this->awsSdkInstrumentation->init();
+        $this->awsSdkInstrumentation->activate();
+
+        foreach ($clients as $name => $client) {
+            $this->assertSame(
+                $preInstrumentationHandlersCount[$name] + self::HANDLERS_PER_ACTIVATION,
+                $client->getHandlerList()->count(),
+                sprintf('Failed asserting that %s client was instrumented once', $name)
+            );
+        }
+    }
+}