Enable event tracing in SNI.dll (dotnet#650)

TrayanZapryanov · Aug 31, 2020 · 1fd8fd4 · 1fd8fd4
1 parent 30ebb76
commit 1fd8fd4
Show file tree

Hide file tree

Showing 9 changed files with 132 additions and 8 deletions.
diff --git a/BUILDGUIDE.md b/BUILDGUIDE.md
@@ -56,8 +56,8 @@ Once the environment is setup properly, execute the desired set of commands belo
 ```
 
 ```bash
-> msbuild /t:BuildAllOSes
-# Builds the driver for all Operating Systems.
+> msbuild /t:BuildNetCoreAllOS
+# Builds the .NET Core driver for all Operating Systems.
 ```
 
 ## Building Tests

diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Interop/SNINativeMethodWrapper.Windows.cs b/src/Microsoft.Data.SqlClient/netcore/src/Interop/SNINativeMethodWrapper.Windows.cs
@@ -304,7 +304,13 @@ internal struct SNI_Error
 
         [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
         private static extern uint SNIWriteSyncOverAsync(SNIHandle pConn, [In] SNIPacket pPacket);
-        #endregion
+
+        [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern bool RegisterTraceProviderWrapper(int eventKeyword);
+
+        [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern void UnregisterTraceProviderWrapper();
+	    #endregion
 
         internal static uint SniGetConnectionId(SNIHandle pConn, ref Guid connId)
         {
@@ -461,6 +467,18 @@ private static void MarshalConsumerInfo(ConsumerInfo consumerInfo, ref Sni_Consu
                 : IntPtr.Zero;
             native_consumerInfo.ConsumerKey = consumerInfo.key;
         }
+
+		internal static bool RegisterTraceProvider(int eventKeyword)
+		{
+			// Registers the TraceLogging provider, enabling it to generate events.
+			// Return true if enabled, otherwise false.
+			return RegisterTraceProviderWrapper(eventKeyword);
+		}
+
+		internal static void UnregisterTraceProvider()
+		{
+			UnregisterTraceProviderWrapper();
+		}
     }
 }
 

diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj
@@ -33,6 +33,11 @@
     <DebugSymbols>true</DebugSymbols>
   </PropertyGroup>
   <!-- Contains common items shared between NetFx and NetCore -->
+  <ItemGroup Condition="'$(OSGroup)' != 'AnyOS' AND '$(TargetsWindows)' == 'true'">
+	<Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlClientEventSource.Windows.cs">
+	  <Link>Microsoft\Data\SqlClient\SqlClientEventSource.Windows.cs</Link>
+	</Compile>
+  </ItemGroup>
   <ItemGroup Condition="'$(OSGroup)' != 'AnyOS'">
     <Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlClientEventSource.cs">
       <Link>Microsoft\Data\SqlClient\SqlClientEventSource.cs</Link>

diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj
@@ -91,6 +91,9 @@
     <Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlClientEventSource.cs">
       <Link>Microsoft\Data\SqlClient\SqlClientEventSource.cs</Link>
     </Compile>
+    <Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlClientEventSource.Windows.cs">
+      <Link>Microsoft\Data\SqlClient\SqlClientEventSource.Windows.cs</Link>
+    </Compile>
     <Compile Include="..\..\src\Microsoft\Data\SqlClient\SqlClientLogger.cs">
       <Link>Microsoft\Data\SqlClient\SqlClientLogger.cs</Link>
     </Compile>

diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/Interop/SNINativeManagedWrapperX64.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/Interop/SNINativeManagedWrapperX64.cs
@@ -133,5 +133,11 @@ internal static class SNINativeManagedWrapperX64
 
         [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
         internal static extern IntPtr SNIClientCertificateFallbackWrapper(IntPtr pCallbackContext);
+
+        [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern bool RegisterTraceProviderWrapper(int eventKeyword);
+
+        [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern void UnregisterTraceProviderWrapper();
     }
 }
diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/Interop/SNINativeManagedWrapperX86.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/Interop/SNINativeManagedWrapperX86.cs
@@ -133,5 +133,11 @@ internal static class SNINativeManagedWrapperX86
 
         [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
         internal static extern IntPtr SNIClientCertificateFallbackWrapper(IntPtr pCallbackContext);
+
+        [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern bool RegisterTraceProviderWrapper(int eventKeyword);
+
+        [DllImport(SNI, CallingConvention = CallingConvention.Cdecl)]
+        internal static extern void UnregisterTraceProviderWrapper();
     }
 }
diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/Interop/SNINativeMethodWrapper.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/Interop/SNINativeMethodWrapper.cs
@@ -51,7 +51,7 @@ internal static class SNINativeMethodWrapper
         internal const int LocalDBInvalidSqlUserInstanceDllPath = 55;
         internal const int LocalDBFailedToLoadDll = 56;
         internal const int LocalDBBadRuntime = 57;
-		internal const int SniIP6AddrStringBufferLength = 48; // from SNI layer
+        internal const int SniIP6AddrStringBufferLength = 48; // from SNI layer
 
         internal static int SniMaxComposedSpnLength
         {
@@ -725,12 +725,12 @@ internal static uint SniGetConnectionId(SNIHandle pConn, ref Guid connId)
         {
             return SNIGetInfoWrapper(pConn, QTypes.SNI_QUERY_CONN_CONNID, out connId);
         }
-       
+
         internal static uint SniGetProviderNumber(SNIHandle pConn, ref ProviderEnum provNum)
         {
             return SNIGetInfoWrapper(pConn, QTypes.SNI_QUERY_CONN_PROVIDERNUM, out provNum);
         }
-     
+
         internal static uint SniGetConnectionPort(SNIHandle pConn, ref ushort portNum)
         {
             return SNIGetInfoWrapper(pConn, QTypes.SNI_QUERY_CONN_PEERPORT, out portNum);
@@ -812,7 +812,7 @@ internal static unsafe uint SNIOpenSyncEx(ConsumerInfo consumerInfo, string cons
                 clientConsumerInfo.DNSCacheInfo.wszCachedTcpIPv4 = cachedDNSInfo?.AddrIPv4;
                 clientConsumerInfo.DNSCacheInfo.wszCachedTcpIPv6 = cachedDNSInfo?.AddrIPv6;
                 clientConsumerInfo.DNSCacheInfo.wszCachedTcpPort = cachedDNSInfo?.Port;
-                
+
                 if (spnBuffer != null)
                 {
                     fixed (byte* pin_spnBuffer = &spnBuffer[0])
@@ -1075,6 +1075,32 @@ private static void MarshalConsumerInfo(ConsumerInfo consumerInfo, ref Sni_Consu
                 : IntPtr.Zero;
             native_consumerInfo.ConsumerKey = consumerInfo.key;
         }
+
+        internal static bool RegisterTraceProvider(int eventKeyword)
+        {
+            // Registers the TraceLogging provider, enabling it to generate events.
+            // Return true if enabled, otherwise false.
+            if (s_is64bitProcess)
+            {
+                return SNINativeManagedWrapperX64.RegisterTraceProviderWrapper(eventKeyword);
+            }
+            else
+            {
+                return SNINativeManagedWrapperX86.RegisterTraceProviderWrapper(eventKeyword);
+            }
+        }
+
+        internal static void UnregisterTraceProvider()
+        {
+            if (s_is64bitProcess)
+            {
+                SNINativeManagedWrapperX64.UnregisterTraceProviderWrapper();
+            }
+            else
+            {
+                SNINativeManagedWrapperX86.UnregisterTraceProviderWrapper();
+            }
+        }
     }
 }
 

diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlClientEventSource.Windows.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlClientEventSource.Windows.cs
@@ -0,0 +1,60 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Diagnostics;
+using System.Diagnostics.Tracing;
+
+namespace Microsoft.Data.SqlClient
+{
+    internal partial class SqlClientEventSource : EventSource
+    {
+        private bool _traceLoggingProviderEnabled = false;
+
+        /// <summary>
+        /// Captures application flow traces from native networking implementation
+        /// </summary>
+        private const EventCommand SNINativeTrace = (EventCommand)8192;
+
+        /// <summary>
+        /// Captures scope trace events from native networking implementation
+        /// </summary>
+        private const EventCommand SNINativeScope = (EventCommand)16384;
+
+        /// <summary>
+        /// Disables all event tracing in native networking implementation
+        /// </summary>
+        private const EventCommand SNINativeDisable = (EventCommand)32768;
+
+        protected override void OnEventCommand(EventCommandEventArgs e)
+        {
+            // Internally, EventListener.EnableEvents sends an event command, with a reserved value of 0, -2, or -3.
+            // When a command is sent via EnableEvents or SendCommand, check if it is a user-defined value
+            // to enable or disable event tracing in sni.dll.
+            // If registration fails, all write and unregister commands will be a no-op.
+
+            // If managed networking is enabled, don't call native wrapper methods
+#if netcoreapp
+            if (AppContext.TryGetSwitch("Switch.Microsoft.Data.SqlClient.UseManagedNetworkingOnWindows", out bool isEnabled) && isEnabled)
+            {
+                return;
+            }
+#endif
+            // Only register the provider if it's not already registered. Registering a provider that is already
+            // registered can lead to unpredictable behaviour.
+            if (!_traceLoggingProviderEnabled && e.Command > 0 && (e.Command & (SNINativeTrace | SNINativeScope)) != 0)
+            {
+                int eventKeyword = (int)(e.Command & (SNINativeTrace | SNINativeScope));
+                _traceLoggingProviderEnabled = SNINativeMethodWrapper.RegisterTraceProvider(eventKeyword);
+                Debug.Assert(_traceLoggingProviderEnabled, "Failed to enable TraceLogging provider.");
+            }
+            else if (_traceLoggingProviderEnabled && (e.Command == SNINativeDisable))
+            {
+                // Only unregister the provider if it's currently registered.
+                SNINativeMethodWrapper.UnregisterTraceProvider();
+                _traceLoggingProviderEnabled = false;
+            }
+        }
+    }
+}
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlClientEventSource.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlClientEventSource.cs
@@ -8,7 +8,7 @@
 namespace Microsoft.Data.SqlClient
 {
     [EventSource(Name = "Microsoft.Data.SqlClient.EventSource")]
-    internal class SqlClientEventSource : EventSource
+    internal partial class SqlClientEventSource : EventSource
     {
         // Defines the singleton instance for the Resources ETW provider
         internal static readonly SqlClientEventSource Log = new SqlClientEventSource();