A bit more refactoring

newrelic · Jun 24, 2024 · c858368 · c858368
1 parent 1b6496c
commit c858368
Showing 1 changed file with 14 additions and 8 deletions.
diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs
@@ -56,29 +56,35 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
  agent.Logger.Debug("AwsSdkPipelineWrapper: requestContext.Request is null. Returning NoOp delegate.");
  return Delegates.NoOp;
  }
- dynamic webRequest = requestContext.Request;
-
- ISegment segment;
 
+ MessageBrokerAction action;
+ var insertDistributedTraceHeaders = false;
  switch (requestType)
  {
  case "SendMessageRequest":
  case "SendMessageBatchRequest":
- segment = SqsHelper.GenerateSegment(transaction, instrumentedMethodCall.MethodCall, requestQueueUrl, MessageBrokerAction.Produce);
- // This needs to happen at the end
- //SqsHelper.InsertDistributedTraceHeaders(transaction, webRequest);
+ action = MessageBrokerAction.Produce;
+ insertDistributedTraceHeaders = true;
  break;
  case "ReceiveMessageRequest":
- segment = SqsHelper.GenerateSegment(transaction, instrumentedMethodCall.MethodCall, requestQueueUrl, MessageBrokerAction.Consume);
+ action = MessageBrokerAction.Consume;
  break;
  case "PurgeQueueRequest":
- segment = SqsHelper.GenerateSegment(transaction, instrumentedMethodCall.MethodCall, requestQueueUrl, MessageBrokerAction.Purge);
+ action = MessageBrokerAction.Purge;
  break;
  default:
  agent.Logger.Debug($"AwsSdkPipelineWrapper: Request type {requestType} is not supported. Returning NoOp delegate.");
  return Delegates.NoOp;
  }
 
+ ISegment segment = SqsHelper.GenerateSegment(transaction, instrumentedMethodCall.MethodCall, requestQueueUrl, action);
+ if (insertDistributedTraceHeaders)
+ {
+ // This needs to happen at the end
+ //dynamic webRequest = requestContext.Request;
+ //SqsHelper.InsertDistributedTraceHeaders(transaction, webRequest);
+ }
+
  return Delegates.GetDelegateFor(segment);
  }
  }