Add MQ profiling support

mythz · mythz · commit 62bf31deb45f · 2022-07-11T16:20:28.000+08:00
diff --git a/ServiceStack.Aws/src/ServiceStack.Aws/Sqs/SqsMqMessageProducer.cs b/ServiceStack.Aws/src/ServiceStack.Aws/Sqs/SqsMqMessageProducer.cs
@@ -31,6 +31,7 @@ public void Publish<T>(T messageBody)
         {
             if (messageBody is IMessage message)
             {
+                Diagnostics.ServiceStack.Init(message);
                 Publish(message.ToInQueueName(), message);
             }
             else
diff --git a/ServiceStack.Azure/src/ServiceStack.Azure/Messaging/ServiceBusMqMessageProducer.cs b/ServiceStack.Azure/src/ServiceStack.Azure/Messaging/ServiceBusMqMessageProducer.cs
@@ -42,6 +42,7 @@ public void Publish<T>(T messageBody)
             // Ensure we're publishing an IMessage
             if (messageBody is IMessage message)
             {
+                Diagnostics.ServiceStack.Init(message);
                 Publish(message.ToInQueueName(), message);
             }
             else
diff --git a/ServiceStack.Text/src/ServiceStack.Text/Diagnostics.cs b/ServiceStack.Text/src/ServiceStack.Text/Diagnostics.cs
@@ -41,6 +41,11 @@ public static class ServiceStack
             public const string WriteGatewayBefore = Prefix + nameof(WriteGatewayBefore);
             public const string WriteGatewayAfter = Prefix + nameof(WriteGatewayAfter);
             public const string WriteGatewayError = Prefix + nameof(WriteGatewayError);
+            
+            public const string WriteMqRequestBefore = Prefix + nameof(WriteMqRequestBefore);
+            public const string WriteMqRequestAfter = Prefix + nameof(WriteMqRequestAfter);
+            public const string WriteMqRequestError = Prefix + nameof(WriteMqRequestError);
+            public const string WriteMqRequestPublish = Prefix + nameof(WriteMqRequestPublish);
         }
         
         public static class Client
@@ -108,6 +113,9 @@ public static class Activity
         public const string OperationId = nameof(OperationId);
         public const string UserId = nameof(UserId);
         public const string Tag = nameof(Tag);
+        public const string MqBegin = nameof(MqBegin);
+        public const string MqEnd = nameof(MqEnd);
+
     }
 
     private DiagnosticListener servicestack { get; set; } = new(Listeners.ServiceStack);
diff --git a/ServiceStack/src/ServiceStack.Client/ClientDiagnosticUtils.cs b/ServiceStack/src/ServiceStack.Client/ClientDiagnosticUtils.cs
@@ -0,0 +1,23 @@
+using System.Diagnostics;
+using System.Runtime.CompilerServices;
+using ServiceStack.Messaging;
+
+namespace ServiceStack;
+
+// sync with CommonDiagnosticUtils
+public static class ClientDiagnosticUtils
+{
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static void InitMessage(this DiagnosticListener listener, IMessage msg)
+    {
+        if (listener.IsEnabled(Diagnostics.Events.ServiceStack.WriteMqRequestBefore))
+        {
+            var activity = Activity.Current;
+            if (activity != null)
+            {
+                msg.TraceId ??= activity.GetTraceId();
+                msg.Tag ??= activity.GetTag();
+            }
+        }
+    }
+}
diff --git a/ServiceStack/src/ServiceStack.Client/Messaging/RedisMessageProducer.cs b/ServiceStack/src/ServiceStack.Client/Messaging/RedisMessageProducer.cs
@@ -27,19 +27,16 @@ public IRedisNativeClient ReadWriteClient
         {
             get
             {
-                if (this.readWriteClient == null)
-                {
-                    this.readWriteClient = (IRedisNativeClient)clientsManager.GetClient();
-                }
+                this.readWriteClient ??= (IRedisNativeClient)clientsManager.GetClient();
                 return readWriteClient;
             }
         }
 
         public void Publish<T>(T messageBody)
         {
-            var message = messageBody as IMessage;
-            if (message != null)
+            if (messageBody is IMessage message)
             {
+                Diagnostics.ServiceStack.InitMessage(message);
                 Publish(message.ToInQueueName(), message);
             }
             else
diff --git a/ServiceStack/src/ServiceStack.Client/Messaging/RedisMessageQueueClient.cs b/ServiceStack/src/ServiceStack.Client/Messaging/RedisMessageQueueClient.cs
@@ -36,6 +36,7 @@ public void Publish<T>(T messageBody)
         {
             if (messageBody is IMessage message)
             {
+                Diagnostics.ServiceStack.InitMessage(message);
                 Publish(message.ToInQueueName(), message);
             }
             else
diff --git a/ServiceStack/src/ServiceStack.Client/ServiceClientBase.cs b/ServiceStack/src/ServiceStack.Client/ServiceClientBase.cs
@@ -1127,6 +1127,7 @@ public void Publish<T>(T requestDto)
 
         public void Publish<T>(IMessage<T> message)
         {
+            Diagnostics.ServiceStack.InitMessage(message);
             var requestDto = message.GetBody();
 
             if (message.CreatedDate != default(DateTime))
diff --git a/ServiceStack/src/ServiceStack.Common/CommonDiagnosticUtils.cs b/ServiceStack/src/ServiceStack.Common/CommonDiagnosticUtils.cs
@@ -0,0 +1,24 @@
+using System;
+using System.Diagnostics;
+using System.Runtime.CompilerServices;
+using ServiceStack.Messaging;
+
+namespace ServiceStack;
+
+// sync with ClientDiagnosticUtils
+public static class CommonDiagnosticUtils
+{
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public static void Init(this DiagnosticListener listener, IMessage msg)
+    {
+        if (listener.IsEnabled(Diagnostics.Events.ServiceStack.WriteMqRequestBefore))
+        {
+            var activity = Activity.Current;
+            if (activity != null)
+            {
+                msg.TraceId ??= activity.GetTraceId();
+                msg.Tag ??= activity.GetTag();
+            }
+        }
+    }
+}
diff --git a/ServiceStack/src/ServiceStack.Interfaces/Messaging/IMessage.cs b/ServiceStack/src/ServiceStack.Interfaces/Messaging/IMessage.cs
@@ -23,6 +23,8 @@ public interface IMessage
         string Tag { get; set; }
 
         object Body { get; set; }
+        
+        string TraceId { get; set; }
     }
 
     public interface IMessage<T>
diff --git a/ServiceStack/src/ServiceStack.Interfaces/Messaging/MessageFactory.cs b/ServiceStack/src/ServiceStack.Interfaces/Messaging/MessageFactory.cs
@@ -51,6 +51,7 @@ public class Message : IMessage
         public string Tag { get; set; }
         public Dictionary<string, string> Meta { get; set; }
         public object Body { get; set; }
+        public string TraceId { get; set; }
 
         public static Message<T> Create<T>(T request) => new(request);
     }
diff --git a/ServiceStack/src/ServiceStack.RabbitMq/RabbitMqProducer.cs b/ServiceStack/src/ServiceStack.RabbitMq/RabbitMqProducer.cs
@@ -37,7 +37,7 @@ public IModel Channel
         {
             get
             {
-                if (channel == null || !channel.IsOpen)
+                if (channel is not { IsOpen: true })
                 {
                     channel = Connection.OpenChannel();
                     //prefetch size is no supported by RabbitMQ
@@ -57,6 +57,7 @@ public virtual void Publish<T>(T messageBody)
         {
             if (messageBody is IMessage message)
             {
+                Diagnostics.ServiceStack.Init(message);
                 Publish(message.ToInQueueName(), message);
             }
             else
diff --git a/ServiceStack/src/ServiceStack/Messaging/BackgroundMqService.cs b/ServiceStack/src/ServiceStack/Messaging/BackgroundMqService.cs
@@ -743,6 +743,7 @@ public void Publish<T>(T messageBody)
         {
             if (messageBody is IMessage message)
             {
+                Diagnostics.ServiceStack.Init(message);
                 Publish(message.ToInQueueName(), message);
             }
             else
diff --git a/ServiceStack/src/ServiceStack/Messaging/InMemoryTransientMessageFactory.cs b/ServiceStack/src/ServiceStack/Messaging/InMemoryTransientMessageFactory.cs
@@ -56,6 +56,7 @@ public void Publish<T>(T messageBody)
             {
                 if (messageBody is IMessage message)
                 {
+                    Diagnostics.ServiceStack.Init(message);
                     Publish(message.ToInQueueName(), message);
                 }
                 else
diff --git a/ServiceStack/src/ServiceStack/Messaging/MessageHandler.cs b/ServiceStack/src/ServiceStack/Messaging/MessageHandler.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading.Tasks;
 using ServiceStack.Logging;
@@ -66,18 +67,17 @@ public void Process(IMessageQueueClient mqClient)
 
         public int ProcessQueue(IMessageQueueClient mqClient, string queueName, Func<bool> doNext = null)
         {
-            var msgsProcessed = 0;
+            var messagesProcessed = 0;
             try
             {
-                IMessage<T> message;
-                while ((message = mqClient.GetAsync<T>(queueName)) != null)
+                while (mqClient.GetAsync<T>(queueName) is { } message)
                 {
                     ProcessMessage(mqClient, message);
 
-                    msgsProcessed++;
+                    messagesProcessed++;
 
                     if (doNext != null && !doNext()) 
-                        return msgsProcessed;
+                        return messagesProcessed;
                 }
             }
             catch (TaskCanceledException) {}
@@ -86,7 +86,7 @@ public int ProcessQueue(IMessageQueueClient mqClient, string queueName, Func<boo
                 Log.Error("Error serializing message from mq server: " + ex.Message, ex);
             }
 
-            return msgsProcessed;
+            return messagesProcessed;
         }
 
         public IMessageHandlerStats GetStats()
@@ -100,7 +100,7 @@ private void DefaultInExceptionHandler(IMessageHandler mqHandler, IMessage<T> me
         {
             Log.Error("Message exception handler threw an error", ex);
 
-            bool requeue = !(ex is UnRetryableMessagingException)
+            bool requeue = ex is not UnRetryableMessagingException
                 && message.RetryAttempts < retryCount;
 
             if (requeue)
@@ -124,6 +124,28 @@ public void ProcessMessage(IMessageQueueClient mqClient, IMessage<T> message)
             this.MqClient = mqClient;
             bool msgHandled = false;
 
+
+            Activity origActivity = null;
+            Activity activity = null;
+            if (Diagnostics.ServiceStack.IsEnabled(Diagnostics.Events.ServiceStack.WriteMqRequestBefore))
+            {
+                origActivity = Activity.Current;
+                if (origActivity == null)
+                {
+                    var traceId = message.TraceId;
+                    if (traceId != null)
+                    {
+                        activity = new Activity(Diagnostics.Activity.MqBegin);
+                        activity.SetParentId(traceId);
+                        if (message.Tag != null)
+                            activity.AddTag(Diagnostics.Activity.Tag, message.Tag);
+
+                        Diagnostics.ServiceStack.StartActivity(activity, new ServiceStackMqActivityArgs { Message = message, Activity = activity });
+                    }
+                }
+            }
+            
+            var id = Diagnostics.ServiceStack.WriteMqRequestBefore(message);
             try
             {
                 var response = processMessageFn(message);
@@ -141,17 +163,20 @@ public void ProcessMessage(IMessageQueueClient mqClient, IMessage<T> message)
                 
                 if (responseEx != null)
                 {
+                    Diagnostics.ServiceStack.WriteMqRequestError(id, message, responseEx);
                     TotalMessagesFailed++;
 
                     if (message.ReplyTo != null)
                     {
                         var replyClient = ReplyClientFactory(message.ReplyTo);
                         if (replyClient != null)
                         {
+                            Diagnostics.ServiceStack.WriteMqRequestPublish(id, replyClient, message.ReplyTo, response);
                             replyClient.SendOneWay(message.ReplyTo, response);
                         }
                         else
                         {
+                            Diagnostics.ServiceStack.WriteMqRequestPublish(id, mqClient, message.ReplyTo, response);
                             var responseDto = response.GetResponseDto();
                             mqClient.Publish(message.ReplyTo, MessageFactory.Create(responseDto));
                         }
@@ -162,6 +187,7 @@ public void ProcessMessage(IMessageQueueClient mqClient, IMessage<T> message)
                     processInExceptionFn(this, message, responseEx);
                     return;
                 }
+                Diagnostics.ServiceStack.WriteMqRequestAfter(id, message);
 
                 this.TotalMessagesProcessed++;
 
@@ -185,6 +211,7 @@ public void ProcessMessage(IMessageQueueClient mqClient, IMessage<T> message)
                         var messageOptions = (MessageOption) message.Options;
                         if (messageOptions.Has(MessageOption.NotifyOneWay))
                         {
+                            Diagnostics.ServiceStack.WriteMqRequestPublish(id, mqClient, QueueNames<T>.Out, response);
                             mqClient.Notify(QueueNames<T>.Out, message);
                         }
                     }
@@ -214,9 +241,6 @@ public void ProcessMessage(IMessageQueueClient mqClient, IMessage<T> message)
                                 return;
                         }
 
-                        // Leave as-is to work around a Mono 2.6.7 compiler bug
-                        if (!responseType.IsUserType()) 
-                            return;
                         mqReplyTo = new QueueNames(responseType).In;
                     }
 
@@ -225,6 +249,7 @@ public void ProcessMessage(IMessageQueueClient mqClient, IMessage<T> message)
                     {
                         try
                         {
+                            Diagnostics.ServiceStack.WriteMqRequestPublish(id, replyClient, mqReplyTo, response);
                             replyClient.SendOneWay(mqReplyTo, response);
                             return;
                         }
@@ -244,6 +269,7 @@ public void ProcessMessage(IMessageQueueClient mqClient, IMessage<T> message)
                         responseMessage = MessageFactory.Create(response);
 
                     responseMessage.ReplyId = message.Id;
+                    Diagnostics.ServiceStack.WriteMqRequestPublish(id, mqClient, responseMessage.ReplyId.ToString(), response);
                     mqClient.Publish(mqReplyTo, responseMessage);
                 }
             }
@@ -270,6 +296,11 @@ public void ProcessMessage(IMessageQueueClient mqClient, IMessage<T> message)
 
                 this.TotalNormalMessagesReceived++;
                 LastMessageProcessed = DateTime.UtcNow;
+
+                if (activity != null)
+                {
+                    Diagnostics.ServiceStack.StopActivity(activity, new ServiceStackMqActivityArgs { Message = message, Activity = activity });
+                }
             }
         }
 
diff --git a/ServiceStack/src/ServiceStack/ProfilingFeature.cs b/ServiceStack/src/ServiceStack/ProfilingFeature.cs
diff --git a/ServiceStack/src/ServiceStack/ServiceStackDiagnostics.cs b/ServiceStack/src/ServiceStack/ServiceStackDiagnostics.cs
diff --git a/ServiceStack/src/ServiceStack/ServiceStackHost.Runtime.cs b/ServiceStack/src/ServiceStack/ServiceStackHost.Runtime.cs
diff --git a/ServiceStack/src/ServiceStack/modules/admin-ui/components/Profiling.html b/ServiceStack/src/ServiceStack/modules/admin-ui/components/Profiling.html
diff --git a/ServiceStack/tests/NorthwindAuto/Configure.Profiling.cs b/ServiceStack/tests/NorthwindAuto/Configure.Profiling.cs
diff --git a/ServiceStack/tests/NorthwindAuto/admin-ui/components/Profiling.html b/ServiceStack/tests/NorthwindAuto/admin-ui/components/Profiling.html

Original file line number	Diff line number	Diff line change
`@@ -31,6 +31,7 @@ public void Publish<T>(T messageBody)`
`31`	`31`	`{`
`32`	`32`	`if (messageBody is IMessage message)`
`33`	`33`	`{`
	`34`	`+ Diagnostics.ServiceStack.Init(message);`
`34`	`35`	`Publish(message.ToInQueueName(), message);`
`35`	`36`	`}`
`36`	`37`	`else`
Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,7 @@ public void Publish<T>(T messageBody)`
`42`	`42`	`// Ensure we're publishing an IMessage`
`43`	`43`	`if (messageBody is IMessage message)`
`44`	`44`	`{`
	`45`	`+ Diagnostics.ServiceStack.Init(message);`
`45`	`46`	`Publish(message.ToInQueueName(), message);`
`46`	`47`	`}`
`47`	`48`	`else`
Original file line number	Diff line number	Diff line change
`@@ -27,19 +27,16 @@ public IRedisNativeClient ReadWriteClient`
`27`	`27`	`{`
`28`	`28`	`get`
`29`	`29`	`{`
`30`		`- if (this.readWriteClient == null)`
`31`		`- {`
`32`		`- this.readWriteClient = (IRedisNativeClient)clientsManager.GetClient();`
`33`		`- }`
	`30`	`+ this.readWriteClient ??= (IRedisNativeClient)clientsManager.GetClient();`
`34`	`31`	`return readWriteClient;`
`35`	`32`	`}`
`36`	`33`	`}`
`37`	`34`
`38`	`35`	`public void Publish<T>(T messageBody)`
`39`	`36`	`{`
`40`		`- var message = messageBody as IMessage;`
`41`		`- if (message != null)`
	`37`	`+ if (messageBody is IMessage message)`
`42`	`38`	`{`
	`39`	`+ Diagnostics.ServiceStack.InitMessage(message);`
`43`	`40`	`Publish(message.ToInQueueName(), message);`
`44`	`41`	`}`
`45`	`42`	`else`
Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@ public void Publish<T>(T messageBody)`
`36`	`36`	`{`
`37`	`37`	`if (messageBody is IMessage message)`
`38`	`38`	`{`
	`39`	`+ Diagnostics.ServiceStack.InitMessage(message);`
`39`	`40`	`Publish(message.ToInQueueName(), message);`
`40`	`41`	`}`
`41`	`42`	`else`
Original file line number	Diff line number	Diff line change
`@@ -1127,6 +1127,7 @@ public void Publish<T>(T requestDto)`
`1127`	`1127`
`1128`	`1128`	`public void Publish<T>(IMessage<T> message)`
`1129`	`1129`	`{`
	`1130`	`+ Diagnostics.ServiceStack.InitMessage(message);`
`1130`	`1131`	`var requestDto = message.GetBody();`
`1131`	`1132`
`1132`	`1133`	`if (message.CreatedDate != default(DateTime))`
Original file line number	Diff line number	Diff line change
`@@ -23,6 +23,8 @@ public interface IMessage`
`23`	`23`	`string Tag { get; set; }`
`24`	`24`
`25`	`25`	`object Body { get; set; }`
	`26`	`+`
	`27`	`+ string TraceId { get; set; }`
`26`	`28`	`}`
`27`	`29`
`28`	`30`	`public interface IMessage<T>`
Original file line number	Diff line number	Diff line change
`@@ -51,6 +51,7 @@ public class Message : IMessage`
`51`	`51`	`public string Tag { get; set; }`
`52`	`52`	`public Dictionary<string, string> Meta { get; set; }`
`53`	`53`	`public object Body { get; set; }`
	`54`	`+ public string TraceId { get; set; }`
`54`	`55`
`55`	`56`	`public static Message<T> Create<T>(T request) => new(request);`
`56`	`57`	`}`