[Messenger] Added X-Ray trace header support to the SQS transport

WaylandAce · WaylandAce · commit 5fa7ff95413e · 2021-04-10T00:55:17.000+03:00
diff --git a/src/Symfony/Component/Messenger/Bridge/AmazonSqs/CHANGELOG.md b/src/Symfony/Component/Messenger/Bridge/AmazonSqs/CHANGELOG.md
@@ -6,6 +6,7 @@ CHANGELOG
 
  * Added new `debug` option to log HTTP requests and responses.
  * Allowed for receiver & sender injection into AmazonSqsTransport
+ * Add X-Ray trace header support to the SQS transport
 
 5.2.0
 -----
diff --git a/src/Symfony/Component/Messenger/Bridge/AmazonSqs/Tests/Transport/AmazonSqsSenderTest.php b/src/Symfony/Component/Messenger/Bridge/AmazonSqs/Tests/Transport/AmazonSqsSenderTest.php
@@ -15,6 +15,7 @@
 use Symfony\Component\Messenger\Bridge\AmazonSqs\Tests\Fixtures\DummyMessage;
 use Symfony\Component\Messenger\Bridge\AmazonSqs\Transport\AmazonSqsFifoStamp;
 use Symfony\Component\Messenger\Bridge\AmazonSqs\Transport\AmazonSqsSender;
+use Symfony\Component\Messenger\Bridge\AmazonSqs\Transport\AmazonSqsXrayTraceHeaderStamp;
 use Symfony\Component\Messenger\Bridge\AmazonSqs\Transport\Connection;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Transport\Serialization\SerializerInterface;
@@ -53,4 +54,22 @@ public function testSendWithAmazonSqsFifoStamp()
         $sender = new AmazonSqsSender($connection, $serializer);
         $sender->send($envelope);
     }
+
+    public function testSendWithAmazonSqsXrayTraceHeaderStamp()
+    {
+        $envelope = (new Envelope(new DummyMessage('Oy')))
+            ->with($stamp = new AmazonSqsXrayTraceHeaderStamp('traceHeader'));
+
+        $encoded = ['body' => '...', 'headers' => ['type' => DummyMessage::class]];
+
+        $connection = $this->createMock(Connection::class);
+        $connection->expects($this->once())->method('send')
+            ->with($encoded['body'], $encoded['headers'], 0, null, null, $stamp->getTraceId());
+
+        $serializer = $this->createMock(SerializerInterface::class);
+        $serializer->method('encode')->with($envelope)->willReturnOnConsecutiveCalls($encoded);
+
+        $sender = new AmazonSqsSender($connection, $serializer);
+        $sender->send($envelope);
+    }
 }
diff --git a/src/Symfony/Component/Messenger/Bridge/AmazonSqs/Transport/AmazonSqsSender.php b/src/Symfony/Component/Messenger/Bridge/AmazonSqs/Transport/AmazonSqsSender.php
@@ -45,6 +45,7 @@ public function send(Envelope $envelope): Envelope
 
         $messageGroupId = null;
         $messageDeduplicationId = null;
+        $xrayTraceId = null;
 
         /** @var AmazonSqsFifoStamp|null $amazonSqsFifoStamp */
         $amazonSqsFifoStamp = $envelope->last(AmazonSqsFifoStamp::class);
@@ -53,13 +54,20 @@ public function send(Envelope $envelope): Envelope
             $messageDeduplicationId = $amazonSqsFifoStamp->getMessageDeduplicationId();
         }
 
+        /** @var AmazonSqsXrayTraceHeaderStamp|null $amazonSqsXrayTraceHeaderStamp */
+        $amazonSqsXrayTraceHeaderStamp = $envelope->last(AmazonSqsXrayTraceHeaderStamp::class);
+        if (null !== $amazonSqsXrayTraceHeaderStamp) {
+            $xrayTraceId = $amazonSqsXrayTraceHeaderStamp->getTraceId();
+        }
+
         try {
             $this->connection->send(
                 $encodedMessage['body'],
                 $encodedMessage['headers'] ?? [],
                 $delay,
                 $messageGroupId,
-                $messageDeduplicationId
+                $messageDeduplicationId,
+                $xrayTraceId
             );
         } catch (HttpException $e) {
             throw new TransportException($e->getMessage(), 0, $e);
diff --git a/src/Symfony/Component/Messenger/Bridge/AmazonSqs/Transport/AmazonSqsXrayTraceHeaderStamp.php b/src/Symfony/Component/Messenger/Bridge/AmazonSqs/Transport/AmazonSqsXrayTraceHeaderStamp.php
@@ -0,0 +1,29 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\Component\Messenger\Bridge\AmazonSqs\Transport;
+
+use Symfony\Component\Messenger\Stamp\NonSendableStampInterface;
+
+final class AmazonSqsXrayTraceHeaderStamp implements NonSendableStampInterface
+{
+    private $traceId;
+
+    public function __construct(string $traceId)
+    {
+        $this->traceId = $traceId;
+    }
+
+    public function getTraceId(): string
+    {
+        return $this->traceId;
+    }
+}
diff --git a/src/Symfony/Component/Messenger/Bridge/AmazonSqs/Transport/Connection.php b/src/Symfony/Component/Messenger/Bridge/AmazonSqs/Transport/Connection.php
@@ -11,10 +11,12 @@
 
 namespace Symfony\Component\Messenger\Bridge\AmazonSqs\Transport;
 
+use AsyncAws\Sqs\Enum\MessageSystemAttributeName;
 use AsyncAws\Sqs\Enum\QueueAttributeName;
 use AsyncAws\Sqs\Result\ReceiveMessageResult;
 use AsyncAws\Sqs\SqsClient;
 use AsyncAws\Sqs\ValueObject\MessageAttributeValue;
+use AsyncAws\Sqs\ValueObject\MessageSystemAttributeValue;
 use Psr\Log\LoggerInterface;
 use Symfony\Component\Messenger\Exception\InvalidArgumentException;
 use Symfony\Component\Messenger\Exception\TransportException;
@@ -312,7 +314,7 @@ public function getMessageCount(): int
         return (int) ($attributes[QueueAttributeName::APPROXIMATE_NUMBER_OF_MESSAGES] ?? 0);
     }
 
-    public function send(string $body, array $headers, int $delay = 0, ?string $messageGroupId = null, ?string $messageDeduplicationId = null): void
+    public function send(string $body, array $headers, int $delay = 0, ?string $messageGroupId = null, ?string $messageDeduplicationId = null, ?string $xrayTraceId = null): void
     {
         if ($this->configuration['auto_setup']) {
             $this->setup();
@@ -323,6 +325,7 @@ public function send(string $body, array $headers, int $delay = 0, ?string $mess
             'MessageBody' => $body,
             'DelaySeconds' => $delay,
             'MessageAttributes' => [],
+            'MessageSystemAttributes' => [],
         ];
 
         $specialHeaders = [];
@@ -346,6 +349,13 @@ public function send(string $body, array $headers, int $delay = 0, ?string $mess
             ]);
         }
 
+        if (null !== $xrayTraceId) {
+            $parameters['MessageSystemAttributes'][MessageSystemAttributeName::AWSTRACE_HEADER] = new MessageSystemAttributeValue([
+                'DataType' => 'String',
+                'StringValue' => $xrayTraceId,
+            ]);
+        }
+
         if (self::isFifoQueue($this->configuration['queue_name'])) {
             $parameters['MessageGroupId'] = null !== $messageGroupId ? $messageGroupId : __METHOD__;
             $parameters['MessageDeduplicationId'] = null !== $messageDeduplicationId ? $messageDeduplicationId : sha1(json_encode(['body' => $body, 'headers' => $headers]));
