Fix System.Net.Sockets telemetry

src/libraries/Common/tests/System/Diagnostics/Tracing/TestEventListener.cs

@@ -2,97 +2,99 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Generic;
+using System.Reflection;
 using System.Threading.Tasks;
 
 namespace System.Diagnostics.Tracing
 {
     /// <summary>Simple event listener than invokes a callback for each event received.</summary>
     internal sealed class TestEventListener : EventListener
     {
-        private readonly string _targetSourceName;
-        private readonly Guid _targetSourceGuid;
-        private readonly EventLevel _level;
+        private class Settings
+        {
+            public EventLevel Level;
+            public EventKeywords Keywords;
+        }
+
+        private readonly Dictionary<string, Settings> _names = new Dictionary<string, Settings>();
+        private readonly Dictionary<Guid, Settings> _guids = new Dictionary<Guid, Settings>();
+
         private readonly double? _eventCounterInterval;
 
         private Action<EventWrittenEventArgs> _eventWritten;
-        private List<EventSource> _tmpEventSourceList = new List<EventSource>();
+        private readonly List<EventSource> _eventSourceList = new List<EventSource>();
 
         public TestEventListener(string targetSourceName, EventLevel level, double? eventCounterInterval = null)
         {
-            // Store the arguments
-            _targetSourceName = targetSourceName;
-            _level = level;
             _eventCounterInterval = eventCounterInterval;
-
-            LoadSourceList();
+            AddSource(targetSourceName, level);
         }
 
         public TestEventListener(Guid targetSourceGuid, EventLevel level, double? eventCounterInterval = null)
         {
-            // Store the arguments
-            _targetSourceGuid = targetSourceGuid;
-            _level = level;
             _eventCounterInterval = eventCounterInterval;
-
-            LoadSourceList();
+            AddSource(targetSourceGuid, level);
         }
 
-        private void LoadSourceList()
+        public void AddSource(string name, EventLevel level, EventKeywords keywords = EventKeywords.All) =>
+            AddSource(name, null, level, keywords);
+
+        public void AddSource(Guid guid, EventLevel level, EventKeywords keywords = EventKeywords.All) =>
+            AddSource(null, guid, level, keywords);
+
+        private void AddSource(string name, Guid? guid, EventLevel level, EventKeywords keywords)
         {
-            // The base constructor, which is called before this constructor,
-            // will invoke the virtual OnEventSourceCreated method for each
-            // existing EventSource, which means OnEventSourceCreated will be
-            // called before _targetSourceGuid and _level have been set.  As such,
-            // we store a temporary list that just exists from the moment this instance
-            // is created (instance field initializers run before the base constructor)
-            // and until we finish construction... in that window, OnEventSourceCreated
-            // will store the sources into the list rather than try to enable them directly,
-            // and then here we can enumerate that list, then clear it out.
-            List<EventSource> sources;
-            lock (_tmpEventSourceList)
+            lock (_eventSourceList)
             {
-                sources = _tmpEventSourceList;
-                _tmpEventSourceList = null;
-            }
-            foreach (EventSource source in sources)
-            {
-                EnableSourceIfMatch(source);
+                var settings = new Settings()
+                {
+                    Level = level,
+                    Keywords = keywords
+                };
+
+                if (name is not null)
+                    _names.Add(name, settings);
+
+                if (guid.HasValue)
+                    _guids.Add(guid.Value, settings);
+
+                foreach (EventSource source in _eventSourceList)
+                {
+                    if (name == source.Name || guid == source.Guid)
+                    {
+                        EnableEventSource(source, level, keywords);
+                    }
+                }
             }
         }
 
+        public void AddActivityTracking() =>
+            AddSource("System.Threading.Tasks.TplEventSource", EventLevel.Informational, (EventKeywords)0x80 /* TasksFlowActivityIds */);
+
         protected override void OnEventSourceCreated(EventSource eventSource)
         {
-            List<EventSource> tmp = _tmpEventSourceList;
-            if (tmp != null)
+            lock (_eventSourceList)
             {
-                lock (tmp)
+                _eventSourceList.Add(eventSource);
+
+                if (_names.TryGetValue(eventSource.Name, out Settings settings) ||
+                    _guids.TryGetValue(eventSource.Guid, out settings))
                 {
-                    if (_tmpEventSourceList != null)
-                    {
-                        _tmpEventSourceList.Add(eventSource);
-                        return;
-                    }
+                    EnableEventSource(eventSource, settings.Level, settings.Keywords);
                 }
             }
-
-            EnableSourceIfMatch(eventSource);
         }
 
-        private void EnableSourceIfMatch(EventSource source)
+        private void EnableEventSource(EventSource source, EventLevel level, EventKeywords keywords)
         {
-            if (source.Name.Equals(_targetSourceName) ||
-                source.Guid.Equals(_targetSourceGuid))
+            var args = new Dictionary<string, string>();
+
+            if (_eventCounterInterval != null)
             {
-                if (_eventCounterInterval != null)
-                {
-                    var args = new Dictionary<string, string> { { "EventCounterIntervalSec", _eventCounterInterval?.ToString() } };
-                    EnableEvents(source, _level, EventKeywords.All, args);
-                }
-                else
-                {
-                    EnableEvents(source, _level);
-                }
+                args.Add("EventCounterIntervalSec", _eventCounterInterval.ToString());
             }
+
+            EnableEvents(source, level, keywords, args);
         }
 
         public void RunWithCallback(Action<EventWrittenEventArgs> handler, Action body)
@@ -109,8 +111,14 @@ public async Task RunWithCallbackAsync(Action<EventWrittenEventArgs> handler, Fu
             finally { _eventWritten = null; }
         }
 
+        // Workaround for being able to inspect the ActivityId property after storing EventWrittenEventArgs
+        // [ActiveIssue("https://github.com/dotnet/runtime/issues/42128")]
+        private static readonly FieldInfo s_activityIdFieldInfo =
+            typeof(EventWrittenEventArgs).GetField("m_activityId", BindingFlags.NonPublic | BindingFlags.Instance);
+
         protected override void OnEventWritten(EventWrittenEventArgs eventData)
         {
+            s_activityIdFieldInfo.SetValue(eventData, eventData.ActivityId);
             _eventWritten?.Invoke(eventData);
         }
     }

src/libraries/System.Net.NameResolution/src/System/Net/Dns.cs

@@ -466,34 +466,9 @@ private static Task GetHostEntryOrAddressesCoreAsync(string hostName, bool justR
 
                 if (NameResolutionTelemetry.Log.IsEnabled())
                 {
-                    ValueStopwatch stopwatch = NameResolutionTelemetry.Log.BeforeResolution(hostName);
-
-                    Task coreTask;
-                    try
-                    {
-                        coreTask = NameResolutionPal.GetAddrInfoAsync(hostName, justAddresses);
-                    }
-                    catch when (LogFailure(stopwatch))
-                    {
-                        Debug.Fail("LogFailure should return false");
-                        throw;
-                    }
-
-                    coreTask.ContinueWith(
-                        (task, state) =>
-                        {
-                            NameResolutionTelemetry.Log.AfterResolution(
-                                stopwatch: (ValueStopwatch)state!,
-                                successful: task.IsCompletedSuccessfully);
-                        },
-                        state: stopwatch,
-                        cancellationToken: default,
-                        TaskContinuationOptions.ExecuteSynchronously,
-                        TaskScheduler.Default);
-
-                    // coreTask is not actually a base Task, but Task<IPHostEntry> / Task<IPAddress[]>
-                    // We have to return it and not the continuation
-                    return coreTask;
+                    return justAddresses
+                        ? (Task)GetAddrInfoWithTelemetryAsync<IPAddress[]>(hostName, justAddresses)
+                        : (Task)GetAddrInfoWithTelemetryAsync<IPHostEntry>(hostName, justAddresses);
                 }
                 else
                 {
@@ -506,6 +481,23 @@ private static Task GetHostEntryOrAddressesCoreAsync(string hostName, bool justR
                 RunAsync(s => GetHostEntryCore((string)s), hostName);
         }
 
+        private static async Task<T> GetAddrInfoWithTelemetryAsync<T>(string hostName, bool justAddresses)
+            where T : class
+        {
+            ValueStopwatch stopwatch = NameResolutionTelemetry.Log.BeforeResolution(hostName);
+
+            T? result = null;
+            try
+            {
+                result = await ((Task<T>)NameResolutionPal.GetAddrInfoAsync(hostName, justAddresses)).ConfigureAwait(false);
+                return result;
+            }
+            finally
+            {
+                NameResolutionTelemetry.Log.AfterResolution(stopwatch, successful: result is not null);
+            }
+        }
+
         private static Task<TResult> RunAsync<TResult>(Func<object, TResult> func, object arg) =>
             Task.Factory.StartNew(func!, arg, CancellationToken.None, TaskCreationOptions.DenyChildAttach, TaskScheduler.Default);
 

src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.cs

@@ -1110,19 +1110,15 @@ public Socket Accept()
                 ref socketAddress.InternalSize,
                 out acceptedSocketHandle);
 
+            if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AfterAccept(errorCode);
+
             // Throw an appropriate SocketException if the native call fails.
             if (errorCode != SocketError.Success)
             {
-                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AcceptFailedAndStop(errorCode, null);
-
                 Debug.Assert(acceptedSocketHandle.IsInvalid);
                 UpdateAcceptSocketErrorForDisposed(ref errorCode);
                 UpdateStatusAfterSocketErrorAndThrowException(errorCode);
             }
-            else
-            {
-                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AcceptStop();
-            }
 
             Debug.Assert(!acceptedSocketHandle.IsInvalid);
 
@@ -2140,7 +2136,13 @@ private bool CanUseConnectEx(EndPoint remoteEP)
 
         internal IAsyncResult UnsafeBeginConnect(EndPoint remoteEP, AsyncCallback? callback, object? state, bool flowContext = false)
         {
-            if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.ConnectStart(remoteEP);
+            if (SocketsTelemetry.Log.IsEnabled())
+            {
+                SocketsTelemetry.Log.ConnectStart(remoteEP);
+
+                // Ignore flowContext when using Telemetry to avoid losing Activity tracking
+                flowContext = true;
+            }
 
             if (CanUseConnectEx(remoteEP))
             {
@@ -2363,7 +2365,15 @@ public void Disconnect(bool reuseSocket)
         //    int - Return code from async Connect, 0 for success, SocketError.NotConnected otherwise
         public void EndConnect(IAsyncResult asyncResult)
         {
-            ThrowIfDisposed();
+            if (Disposed)
+            {
+                if (SocketsTelemetry.Log.IsEnabled() && asyncResult is not MultipleAddressConnectAsyncResult)
+                {
+                    SocketsTelemetry.Log.AfterConnect(SocketError.OperationAborted);
+                }
+
+                ThrowObjectDisposedException();
+            }
 
             // Validate input parameters.
             if (asyncResult == null)
@@ -2391,10 +2401,14 @@ asyncResult as MultipleAddressConnectAsyncResult ??
             if (NetEventSource.Log.IsEnabled()) NetEventSource.Info(this, $"asyncResult:{asyncResult}");
 
             Exception? ex = castedAsyncResult.Result as Exception;
-            if (ex != null || (SocketError)castedAsyncResult.ErrorCode != SocketError.Success)
+
+            if (SocketsTelemetry.Log.IsEnabled() && castedAsyncResult is not MultipleAddressConnectAsyncResult)
             {
-                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.ConnectFailedAndStop((SocketError)castedAsyncResult.ErrorCode, ex?.Message);
+                SocketsTelemetry.Log.AfterConnect((SocketError)castedAsyncResult.ErrorCode, ex?.Message);
+            }
 
+            if (ex != null || (SocketError)castedAsyncResult.ErrorCode != SocketError.Success)
+            {
                 if (ex == null)
                 {
                     SocketError errorCode = (SocketError)castedAsyncResult.ErrorCode;
@@ -2409,8 +2423,6 @@ asyncResult as MultipleAddressConnectAsyncResult ??
                 ExceptionDispatchInfo.Throw(ex);
             }
 
-            if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.ConnectStop();
-
             if (NetEventSource.Log.IsEnabled()) NetEventSource.Connected(this, LocalEndPoint, RemoteEndPoint);
         }
 
@@ -3540,14 +3552,12 @@ private IAsyncResult BeginAcceptCommon(Socket? acceptSocket, int receiveSize, As
             // Throw an appropriate SocketException if the native call fails synchronously.
             if (!CheckErrorAndUpdateStatus(errorCode))
             {
-                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AcceptFailedAndStop(errorCode, null);
+                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AfterAccept(errorCode);
 
                 UpdateAcceptSocketErrorForDisposed(ref errorCode);
                 throw new SocketException((int)errorCode);
             }
 
-            if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AcceptStop();
-
             // Finish the flow capture, maybe complete here.
             asyncResult.FinishPostingAsyncOp(ref Caches.AcceptClosureCache);
 
@@ -3573,7 +3583,12 @@ public Socket EndAccept(IAsyncResult asyncResult)
         }
         private Socket EndAcceptCommon(out byte[]? buffer, out int bytesTransferred, IAsyncResult asyncResult)
         {
-            ThrowIfDisposed();
+            if (Disposed)
+            {
+                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AfterAccept(SocketError.OperationAborted);
+
+                ThrowObjectDisposedException();
+            }
 
             // Validate input parameters.
             if (asyncResult == null)
@@ -3594,21 +3609,20 @@ private Socket EndAcceptCommon(out byte[]? buffer, out int bytesTransferred, IAs
             bytesTransferred = (int)castedAsyncResult.BytesTransferred;
             buffer = castedAsyncResult.Buffer;
 
+            if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.BytesReceived(bytesTransferred);
+
             castedAsyncResult.EndCalled = true;
 
             // Throw an appropriate SocketException if the native call failed asynchronously.
             SocketError errorCode = (SocketError)castedAsyncResult.ErrorCode;
+
+            if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AfterAccept(errorCode);
+
             if (errorCode != SocketError.Success)
             {
-                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AcceptFailedAndStop(errorCode, null);
-
                 UpdateAcceptSocketErrorForDisposed(ref errorCode);
                 UpdateStatusAfterSocketErrorAndThrowException(errorCode);
             }
-            else
-            {
-                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AcceptStop();
-            }
 
             if (NetEventSource.Log.IsEnabled()) NetEventSource.Accepted(socket, socket.RemoteEndPoint, socket.LocalEndPoint);
             return socket;
@@ -4223,11 +4237,11 @@ private void DoConnect(EndPoint endPointSnapshot, Internals.SocketAddress socket
 
             SocketError errorCode = SocketPal.Connect(_handle, socketAddress.Buffer, socketAddress.Size);
 
+            if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.AfterConnect(errorCode);
+
             // Throw an appropriate SocketException if the native call fails.
             if (errorCode != SocketError.Success)
             {
-                if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.ConnectFailedAndStop(errorCode, null);
-
                 UpdateConnectSocketErrorForDisposed(ref errorCode);
                 // Update the internal state of this socket according to the error before throwing.
                 SocketException socketException = SocketExceptionFactory.CreateSocketException((int)errorCode, endPointSnapshot);
@@ -4236,8 +4250,6 @@ private void DoConnect(EndPoint endPointSnapshot, Internals.SocketAddress socket
                 throw socketException;
             }
 
-            if (SocketsTelemetry.Log.IsEnabled()) SocketsTelemetry.Log.ConnectStop();
-
             if (_rightEndPoint == null)
             {
                 // Save a copy of the EndPoint so we can use it for Create().