#1034: refactored in order to support "nothing or all" trace data inj…

…ection + adapted tests accordingly

src/OpenTelemetry.Contrib.Instrumentation.AWS/Implementation/AWSMessagingUtils.cs

@@ -14,61 +14,53 @@
 // limitations under the License.
 // </copyright>
 
-using Amazon.Runtime;
-using SNS = Amazon.SimpleNotificationService.Model;
-using SQS = Amazon.SQS.Model;
+using System.Collections.Generic;
+using OpenTelemetry.Context.Propagation;
 
 namespace OpenTelemetry.Contrib.Instrumentation.AWS.Implementation;
 
 internal static class AWSMessagingUtils
 {
-    private static readonly AWSMessageAttributeHelper SqsAttributeHelper = new(new SqsMessageAttributeFormatter());
-    private static readonly AWSMessageAttributeHelper SnsAttributeHelper = new(new SnsMessageAttributeFormatter());
+    // SQS/SNS message attributes collection size limit according to
+    // https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-message-metadata.html and
+    // https://docs.aws.amazon.com/sns/latest/dg/sns-message-attributes.html
+    private const int MaxMessageAttributes = 10;
 
-    internal static void SqsMessageAttributeSetter(IRequestContext context, string name, string value)
+    internal static void Inject(IRequestContextAdapter requestAdapter, PropagationContext propagationContext)
     {
-        var parameters = context.Request?.ParameterCollection;
-        if (parameters == null ||
-            !parameters.ContainsKey("MessageBody") ||
-            context.OriginalRequest is not SQS::SendMessageRequest originalRequest)
+        if (!requestAdapter.HasMessageBody ||
+            !requestAdapter.HasOriginalRequest)
         {
             return;
         }
 
-        // Add trace data to parameters collection of the request.
-        if (SqsAttributeHelper.TryAddParameter(parameters, name, value))
-        {
-            // Add trace data to message attributes dictionary of the original request.
-            // This dictionary must be in sync with parameters collection to pass through the MD5 hash matching check.
-            if (!originalRequest.MessageAttributes.ContainsKey(name))
-            {
-                originalRequest.MessageAttributes.Add(
-                    name, new SQS::MessageAttributeValue
-                    { DataType = "String", StringValue = value });
-            }
-        }
-    }
+        var carrier = new Dictionary<string, string>();
+        Propagators.DefaultTextMapPropagator.Inject(propagationContext, carrier, Setter);
 
-    internal static void SnsMessageAttributeSetter(IRequestContext context, string name, string value)
-    {
-        var parameters = context.Request?.ParameterCollection;
-        if (parameters == null ||
-            !parameters.ContainsKey("Message") ||
-            context.OriginalRequest is not SNS::PublishRequest originalRequest)
+        int attributesCount = requestAdapter.AttributesCount;
+        if (carrier.Count + attributesCount > MaxMessageAttributes)
         {
+            // TODO: Add logging (event source).
             return;
         }
 
-        if (SnsAttributeHelper.TryAddParameter(parameters, name, value))
+        int nextAttributeIndex = attributesCount + 1;
+        foreach (var param in carrier)
         {
-            // Add trace data to message attributes dictionary of the original request.
-            // This dictionary must be in sync with parameters collection to pass through the MD5 hash matching check.
-            if (!originalRequest.MessageAttributes.ContainsKey(name))
+            if (requestAdapter.ContainsAttribute(param.Key))
             {
-                originalRequest.MessageAttributes.Add(
-                    name, new SNS::MessageAttributeValue
-                    { DataType = "String", StringValue = value });
+                continue;
             }
+
+            requestAdapter.AddAttribute(param.Key, param.Value, nextAttributeIndex);
+            nextAttributeIndex++;
+
+            // Add trace data to message attributes dictionary of the original request.
+            // This dictionary must be in sync with parameters collection to pass through the MD5 hash matching check.
+            requestAdapter.AddAttributeToOriginalRequest(param.Key, param.Value);
         }
     }
+
+    private static void Setter(IDictionary<string, string> carrier, string name, string value) =>
+        carrier[name] = value;
 }

src/OpenTelemetry.Contrib.Instrumentation.AWS/Implementation/AWSTracingPipelineHandler.cs

@@ -202,11 +202,11 @@ private void AddRequestSpecificInformation(Activity activity, IRequestContext re
         }
         else if (AWSServiceType.IsSqsService(service))
         {
-            Propagators.DefaultTextMapPropagator.Inject(new PropagationContext(activity.Context, Baggage.Current), requestContext, AWSMessagingUtils.SqsMessageAttributeSetter);
+            AWSMessagingUtils.Inject(new SqsRequestContextAdapter(requestContext), new PropagationContext(activity.Context, Baggage.Current));
         }
         else if (AWSServiceType.IsSnsService(service))
         {
-            Propagators.DefaultTextMapPropagator.Inject(new PropagationContext(activity.Context, Baggage.Current), requestContext, AWSMessagingUtils.SnsMessageAttributeSetter);
+            AWSMessagingUtils.Inject(new SnsRequestContextAdapter(requestContext), new PropagationContext(activity.Context, Baggage.Current));
         }
     }
 

src/OpenTelemetry.Contrib.Instrumentation.AWS/Implementation/IAWSMessageAttributeFormatter.cs → src/OpenTelemetry.Contrib.Instrumentation.AWS/Implementation/IRequestContextAdapter.cs

@@ -1,4 +1,4 @@
-// <copyright file="IAWSMessageAttributeFormatter.cs" company="OpenTelemetry Authors">
+// <copyright file="IRequestContextAdapter.cs" company="OpenTelemetry Authors">
 // Copyright The OpenTelemetry Authors
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,14 +14,18 @@
 // limitations under the License.
 // </copyright>
 
-using System.Text.RegularExpressions;
-
 namespace OpenTelemetry.Contrib.Instrumentation.AWS.Implementation;
-internal interface IAWSMessageAttributeFormatter
+internal interface IRequestContextAdapter
 {
-    Regex AttributeNameRegex { get; }
+    bool HasMessageBody { get; }
+
+    bool HasOriginalRequest { get; }
+
+    int AttributesCount { get; }
+
+    bool ContainsAttribute(string name);
 
-    string AttributeNamePrefix { get; }
+    void AddAttribute(string name, string value, int nextAttributeIndex);
 
-    int? GetAttributeIndex(string attributeName);
+    void AddAttributeToOriginalRequest(string name, string value);
 }

src/OpenTelemetry.Contrib.Instrumentation.AWS/Implementation/SnsRequestContextAdapter.cs

@@ -0,0 +1,59 @@
+// <copyright file="SnsRequestContextAdapter.cs" company="OpenTelemetry Authors">
+// Copyright The OpenTelemetry Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// </copyright>
+
+using Amazon.Runtime;
+using Amazon.Runtime.Internal;
+using Amazon.SimpleNotificationService.Model;
+
+namespace OpenTelemetry.Contrib.Instrumentation.AWS.Implementation;
+internal class SnsRequestContextAdapter : IRequestContextAdapter
+{
+    private readonly ParameterCollection parameters;
+    private readonly PublishRequest originalRequest;
+
+    public SnsRequestContextAdapter(IRequestContext context)
+    {
+        this.parameters = context.Request?.ParameterCollection;
+        this.originalRequest = context.OriginalRequest as PublishRequest;
+    }
+
+    public bool HasMessageBody =>
+        this.parameters?.ContainsKey("Message") ?? false;
+
+    public bool HasOriginalRequest => this.originalRequest != null;
+
+    public int AttributesCount =>
+        this.originalRequest?.MessageAttributes.Count ?? 0;
+
+    public void AddAttribute(string name, string value, int nextAttributeIndex)
+    {
+        if (this.parameters == null)
+        {
+            return;
+        }
+
+        var attributePrefix = "MessageAttributes.entry." + nextAttributeIndex;
+        this.parameters.Add(attributePrefix + ".Name", name);
+        this.parameters.Add(attributePrefix + ".Value.DataType", "String");
+        this.parameters.Add(attributePrefix + ".Value.StringValue", value);
+    }
+
+    public void AddAttributeToOriginalRequest(string name, string value) =>
+        this.originalRequest?.MessageAttributes.Add(name, new MessageAttributeValue { DataType = "String", StringValue = value });
+
+    public bool ContainsAttribute(string name)
+        => this.originalRequest?.MessageAttributes.ContainsKey(name) ?? false;
+}