Get current context from ConcurrentDictionary (#499)

SkyAPM · Aug 23, 2022 · b4fc68a · b4fc68a
1 parent 596b357
commit b4fc68a
Showing 1 changed file with 22 additions and 15 deletions.
diff --git a/src/SkyApm.Diagnostics.MongoDB/MongoDiagnosticProcessor.cs b/src/SkyApm.Diagnostics.MongoDB/MongoDiagnosticProcessor.cs
@@ -18,28 +18,31 @@
 
 using MongoDB.Driver.Core.Events;
 using SkyApm.Tracing;
+using SkyApm.Tracing.Segments;
 using System;
+using System.Collections.Concurrent;
 
 namespace SkyApm.Diagnostics.MongoDB
 {
     public class MongoDiagnosticsProcessor : ITracingDiagnosticProcessor
     {
+        private readonly ConcurrentDictionary<int, SegmentContext> _contextMap = new ConcurrentDictionary<int, SegmentContext>();
+
         public string ListenerName => "MongoSourceListener";
         private readonly ITracingContext _tracingContext;
-        private readonly IExitSegmentContextAccessor _contextAccessor;
 
-        public MongoDiagnosticsProcessor(ITracingContext tracingContext,
-            IExitSegmentContextAccessor contextAccessor)
+        public MongoDiagnosticsProcessor(ITracingContext tracingContext)
         {
             _tracingContext = tracingContext;
-            _contextAccessor = contextAccessor;
         }
 
         [DiagnosticName("MongoActivity.Start")]
         public void BeforeExecuteCommand([Object] CommandStartedEvent @event)
         {
             var operationName = DiagnosticsActivityEventSubscriber.GetCollectionName(@event);
             var context = _tracingContext.CreateExitSegmentContext(operationName, @event.ConnectionId.ServerId.EndPoint.ToString());
+            _contextMap.TryAdd(@event.RequestId, context); 
+
             context.Span.SpanLayer = Tracing.Segments.SpanLayer.DB;
             context.Span.Component = Common.Components.MongoDBCLIENT;
             context.Span.AddTag("db.system", "mongodb");
@@ -53,24 +56,28 @@ public void BeforeExecuteCommand([Object] CommandStartedEvent @event)
 
         [DiagnosticName("MongoActivity.Stop")]
         public void AfterExecuteCommand([Object] CommandSucceededEvent @event)
-        { 
-            var context = _contextAccessor.Context;
-            context?.Span.AddTag(Common.Tags.STATUS_CODE, "ok");
+        {
+            if (_contextMap.TryRemove(@event.RequestId, out var context))
+            {
+                context?.Span.AddTag(Common.Tags.STATUS_CODE, "ok");
 
-            _tracingContext.Release(context);
+                _tracingContext.Release(context);
+            }
         }
 
         [DiagnosticName("MongoActivity.Failed")]
         public void FailedExecuteCommand([Object] CommandFailedEvent @event)
         {
-            var context = _contextAccessor.Context;
-            context?.Span.AddTag("status_description", @event.Failure.Message);
-            context?.Span.AddTag("error.type", @event.Failure.GetType().FullName);
-            context?.Span.AddTag("error.msg", @event.Failure.Message);
-            context?.Span.AddTag("error.stack", @event.Failure.StackTrace);
-            context?.Span.AddTag(Common.Tags.STATUS_CODE, "error");
+            if (_contextMap.TryRemove(@event.RequestId, out var context))
+            {
+                context?.Span.AddTag("status_description", @event.Failure.Message);
+                context?.Span.AddTag("error.type", @event.Failure.GetType().FullName);
+                context?.Span.AddTag("error.msg", @event.Failure.Message);
+                context?.Span.AddTag("error.stack", @event.Failure.StackTrace);
+                context?.Span.AddTag(Common.Tags.STATUS_CODE, "error");
 
-            _tracingContext.Release(context);
+                _tracingContext.Release(context);
+            }
         }
 
     }