refactoring

Author: rithikanarayan

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs

@@ -71,7 +71,7 @@ public object OnDelegateBegin<TArg1>(object sender, ref TArg1 arg)
             LambdaCommon.Log("DelegateWrapper Running OnDelegateBegin");
 
             Scope scope;
-            object state = null;
+            object requestid = null;
             var proxyInstance = arg.DuckCast<IInvocationRequest>();
             if (proxyInstance == null)
             {
@@ -82,11 +82,11 @@ public object OnDelegateBegin<TArg1>(object sender, ref TArg1 arg)
             {
                 var jsonString = ConvertPayloadStream(proxyInstance.InputStream);
                 scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), jsonString, proxyInstance.LambdaContext);
-                state = proxyInstance.LambdaContext?.AwsRequestId;
+                requestid = proxyInstance.LambdaContext?.AwsRequestId;
             }
 
             LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateBegin");
-            return new CallTargetState(scope, state);
+            return new CallTargetState(scope, requestid);
         }
 
         public void OnException(object sender, Exception ex)

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs

@@ -24,7 +24,7 @@ internal abstract class LambdaCommon
     private const string PlaceholderOperationName = "placeholder-operation";
     private const double ServerlessMaxWaitingFlushTime = 3;
     private const string LogLevelEnvName = "DD_LOG_LEVEL";
-    private const string LambdaRuntimeAwsRequestHeaderId = "lambda-runtime-aws-request-id";
+    private const string LambdaRuntimeAwsRequestIdHeader = "lambda-runtime-aws-request-id";
 
     internal static Scope CreatePlaceholderScope(Tracer tracer, NameValueHeadersCollection headers)
     {
@@ -45,7 +45,12 @@ internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder
     {
         var request = requestBuilder.GetStartInvocationRequest();
         WriteRequestPayload(request, data);
-        WriteRequestHeaders(request, context);
+        WriteRequestHeaders(request, context.ClientContext?.Custom);
+        if (context.AwsRequestId != null)
+        {
+            request.Headers.Add(LambdaRuntimeAwsRequestIdHeader, context.AwsRequestId);
+        }
+
         using var response = (HttpWebResponse)request.GetResponse();
 
         var headers = response.Headers.Wrap();
@@ -121,23 +126,11 @@ private static void WriteRequestPayload(WebRequest request, string data)
         dataStream.Close();
     }
 
-    private static void WriteRequestHeaders(WebRequest request, ILambdaContext context)
+    private static void WriteRequestHeaders(WebRequest request, IDictionary<string, string> context)
     {
-        if (context != null)
+        foreach (var kv in context)
         {
-            var clientContext = context.ClientContext?.Custom;
-            if (clientContext != null)
-            {
-                foreach (var kv in clientContext)
-                {
-                    request.Headers.Add(kv.Key, kv.Value);
-                }
-            }
-
-            if (context.AwsRequestId != null)
-            {
-                request.Headers.Add(LambdaRuntimeAwsRequestHeaderId, context.AwsRequestId);
-            }
+            request.Headers.Add(kv.Key, kv.Value);
         }
     }