Adding Kestrel connection start/stop events (#2360)

* connection start/stop events for kestrel * addressing PR comments. 1. Made the params nullable 2. change case of LocalEndPoint var
microsoft · Dec 20, 2023 · 9134c14 · 9134c14
1 parent ae0c844
commit 9134c14
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 0 deletions.
diff --git a/src/TelemetryConsumption/Kestrel/IKestrelTelemetryConsumer.cs b/src/TelemetryConsumption/Kestrel/IKestrelTelemetryConsumer.cs
@@ -10,6 +10,22 @@ namespace Yarp.Telemetry.Consumption;
 /// </summary>
 public interface IKestrelTelemetryConsumer
 {
+    /// <summary>
+    /// Called at the start of a connection.
+    /// </summary>
+    /// <param name="timestamp">Timestamp when the event was fired.</param>
+    /// <param name="connectionId">ID of the connection.</param>
+    /// <param name="localEndPoint">Local endpoint for the connection.</param>
+    /// <param name="remoteEndPoint">Remote endpoint for the connection.</param>
+    void OnConnectionStart(DateTime timestamp, string connectionId, string? localEndPoint, string? remoteEndPoint) { }
+
+    /// <summary>
+    /// Called at the end of a connection.
+    /// </summary>
+    /// <param name="timestamp">Timestamp when the event was fired.</param>
+    /// <param name="connectionId">ID of the connection.</param>
+    void OnConnectionStop(DateTime timestamp, string connectionId) { }
+
     /// <summary>
     /// Called at the start of a request.
     /// </summary>

diff --git a/src/TelemetryConsumption/Kestrel/KestrelEventListenerService.cs b/src/TelemetryConsumption/Kestrel/KestrelEventListenerService.cs
@@ -28,6 +28,30 @@ protected override void OnEvent(IKestrelTelemetryConsumer[] consumers, EventWrit
 
         switch (eventData.EventId)
         {
+            case 1:
+                Debug.Assert(eventData.EventName == "ConnectionStart" && payload.Count == 3);
+                {
+                    var connectionId = (string)payload[0];
+                    var localEndPoint = (string?)payload[1];
+                    var remoteEndPoint = (string?)payload[2];
+                    foreach (var consumer in consumers)
+                    {
+                        consumer.OnConnectionStart(eventData.TimeStamp, connectionId, localEndPoint, remoteEndPoint);
+                    }
+                }
+                break;
+
+            case 2:
+                Debug.Assert(eventData.EventName == "ConnectionStop" && payload.Count == 1);
+                {
+                    var connectionId = (string)payload[0];
+                    foreach (var consumer in consumers)
+                    {
+                        consumer.OnConnectionStop(eventData.TimeStamp, connectionId);
+                    }
+                }
+                break;
+
             case 3:
                 Debug.Assert(eventData.EventName == "RequestStart" && payload.Count == 5);
                 {

diff --git a/test/ReverseProxy.FunctionalTests/TelemetryConsumptionTests.cs b/test/ReverseProxy.FunctionalTests/TelemetryConsumptionTests.cs
@@ -123,6 +123,7 @@ await test.Invoke(async uri =>
 
         var expected = new[]
         {
+            "OnConnectionStart-Kestrel",
             "OnRequestStart-Kestrel",
             "OnForwarderInvoke",
             "OnForwarderStart",
@@ -144,6 +145,7 @@ await test.Invoke(async uri =>
             "OnContentTransferred",
             "OnForwarderStop",
             "OnRequestStop-Kestrel",
+            "OnConnectionStop-Kestrel",
         };
 
         if (!useHttpsOnDestination)
@@ -283,8 +285,10 @@ public void OnRequestStart(DateTime timestamp, string scheme, string host, int p
         public void OnConnectStart(DateTime timestamp, string address) => AddStage(nameof(OnConnectStart), timestamp);
         public void OnConnectStop(DateTime timestamp) => AddStage(nameof(OnConnectStop), timestamp);
         public void OnConnectFailed(DateTime timestamp, SocketError error, string exceptionMessage) => AddStage(nameof(OnConnectFailed), timestamp);
+        public void OnConnectionStart(DateTime timestamp, string connectionId, string localEndPoint, string remoteEndPoint) => AddStage($"{nameof(OnConnectionStart)}-Kestrel", timestamp);
         public void OnRequestStart(DateTime timestamp, string connectionId, string requestId, string httpVersion, string path, string method) => AddStage($"{nameof(OnRequestStart)}-Kestrel", timestamp);
         public void OnRequestStop(DateTime timestamp, string connectionId, string requestId, string httpVersion, string path, string method) => AddStage($"{nameof(OnRequestStop)}-Kestrel", timestamp);
+        public void OnConnectionStop(DateTime timestamp, string connectionId) => AddStage($"{nameof(OnConnectionStop)}-Kestrel", timestamp);
         public void OnRedirect(DateTime timestamp, string redirectUri) => AddStage(nameof(OnRedirect), timestamp);
     }