fix: continue with origin trace info when pub/sub retry (#507)

SkyAPM · Sep 15, 2022 · 06c60d9 · 06c60d9
1 parent 7b3119d
commit 06c60d9
Showing 1 changed file with 31 additions and 6 deletions.
diff --git a/src/SkyApm.Diagnostics.CAP/CapDiagnosticProcessor.cs b/src/SkyApm.Diagnostics.CAP/CapDiagnosticProcessor.cs
@@ -106,11 +106,24 @@ public void ErrorPublishStore([Object] CapEventDataPubStore eventData)
         [DiagnosticName(CapEvents.BeforePublish)]
         public void BeforePublish([Object] CapEventDataPubSend eventData)
         {
-            _localSegmentContextAccessor.Context = _contexts[eventData.TransportMessage.GetId()];
-
+            SegmentContext context = null;
             var host = eventData.BrokerAddress.Endpoint.Replace("-1", "5672");
-            var context = _tracingContext.CreateExitSegmentContext(OperateNamePrefix + eventData.Operation + ProducerOperateNameSuffix,
-                host, new CapCarrierHeaderCollection(eventData.TransportMessage));
+            if (_contexts.TryGetValue(eventData.TransportMessage.GetId(),out var ctx))
+            {
+                _localSegmentContextAccessor.Context = ctx;
+                context = _tracingContext.CreateExitSegmentContext(OperateNamePrefix + eventData.Operation + ProducerOperateNameSuffix,
+                    host, new CapCarrierHeaderCollection(eventData.TransportMessage));
+
+            }
+            else
+            {
+                // may be come from retry loop
+                var carrierHeader = new CapCarrierHeaderCollection(eventData.TransportMessage);
+                var eventName = OperateNamePrefix + eventData.Operation + ProducerOperateNameSuffix;
+                var operationName = OperateNamePrefix + eventName + ConsumerOperateNameSuffix;
+                context = _tracingContext.CreateEntrySegmentContext(operationName, carrierHeader);
+            }
+
 
             context.Span.SpanLayer = SpanLayer.MQ;
             context.Span.Component = GetComponent(eventData.BrokerAddress, true);
@@ -205,9 +218,21 @@ public void CapErrorConsume([Object] CapEventDataSubStore eventData)
         [DiagnosticName(CapEvents.BeforeSubscriberInvoke)]
         public void CapBeforeSubscriberInvoke([Object] CapEventDataSubExecute eventData)
         {
-            _entrySegmentContextAccessor.Context = _contexts[eventData.Message.GetId() + eventData.Message.GetGroup()];
+            SegmentContext context = null;
+            if (_contexts.TryGetValue(eventData.Message.GetId() + eventData.Message.GetGroup(),out var ctx))
+            {
+                _entrySegmentContextAccessor.Context = ctx;
+                context = _tracingContext.CreateLocalSegmentContext("Subscriber Invoke: " + eventData.MethodInfo.Name);
+            }
+            else
+            {
+                // may be come from retry loop
+                var carrierHeader = new CapCarrierHeaderCollection(eventData.Message);
+                var eventName = eventData.Message.GetGroup() + "/" + eventData.Operation;
+                var operationName = OperateNamePrefix + eventName + ConsumerOperateNameSuffix;
+                context = _tracingContext.CreateEntrySegmentContext(operationName, carrierHeader);
+            }
 
-            var context = _tracingContext.CreateLocalSegmentContext("Subscriber Invoke: " + eventData.MethodInfo.Name);
             context.Span.SpanLayer = SpanLayer.MQ;
             context.Span.Component = Components.CAP;
             context.Span.AddLog(LogEvent.Event("Subscriber Invoke Start"));