feat: speed up tcp state collector

Vostok.Metrics.System/Host/Interop.cs

@@ -0,0 +1 @@
+

Vostok.Metrics.System/Host/TcpStateCollector.cs

@@ -1,22 +1,277 @@
-using System.Collections.Generic;
+using System;
+using System.Collections.Generic;
+using System.Net;
 using System.Net.NetworkInformation;
+using System.Net.Sockets;
+using System.Runtime.InteropServices;
 
 namespace Vostok.Metrics.System.Host
 {
+    internal readonly struct TcpInfo
+    {
+        public TcpInfo(TcpState state, IPEndPoint localEndPoint, IPEndPoint remoteEndPoint)
+        {
+            State = state;
+            LocalEndPoint = localEndPoint;
+            RemoteEndPoint = remoteEndPoint;
+        }
+
+        public readonly TcpState State;
+        public readonly IPEndPoint LocalEndPoint;
+        public readonly IPEndPoint RemoteEndPoint;
+    }
+
     internal class TcpStateCollector
     {
         public void Collect(HostMetrics metrics)
         {
             var states = new Dictionary<TcpState, int>();
-
+
+            #if NET6_0_OR_GREATER
+
+            IterateOverTcpConnections(info =>
+            {
+                if (!states.ContainsKey(info.State))
+                    states[info.State] = 0;
+                states[info.State]++;
+            });
+
+            #else
+
             foreach (var tcpConnection in IPGlobalProperties.GetIPGlobalProperties().GetActiveTcpConnections())
             {
                 if (!states.ContainsKey(tcpConnection.State))
                     states[tcpConnection.State] = 0;
                 states[tcpConnection.State]++;
             }
+
+            #endif
+
+
 
             metrics.TcpStates = states;
         }
+
+        #if NET6_0_OR_GREATER
+        private unsafe void IterateOverTcpConnections(Action<TcpInfo> handle)
+        {
+            var size = 0U;
+            uint result;
+
+            if (Socket.OSSupportsIPv4)
+            {
+                result = Interop.IpHlpApi.GetTcpTable(IntPtr.Zero, ref size, true);
+
+                while (result == Interop.IpHlpApi.ERROR_INSUFFICIENT_BUFFER)
+                {
+                    var buffer = Marshal.AllocHGlobal((int)size);
+                    try
+                    {
+                        result = Interop.IpHlpApi.GetTcpTable(buffer, ref size, true);
+
+                        if (result == Interop.IpHlpApi.ERROR_SUCCESS)
+                        {
+                            var span = new ReadOnlySpan<byte>((byte*)buffer, (int)size);
+
+                            // The table info just gives us the number of rows.
+                            ref readonly var tcpTableInfo = ref MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcpTable>(span);
+
+                            if (tcpTableInfo.numberOfEntries > 0)
+                            {
+                                var row = MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcpRow>(span);
+                                // Skip over the tableinfo to get the inline rows.
+                                span = span[sizeof(Interop.IpHlpApi.MibTcpTable)..];
+
+                                for (var i = 0; i < tcpTableInfo.numberOfEntries; i++)
+                                {
+                                    var state = row.state;
+
+                                    // Port is returned in Big-Endian - most significant bit on left.
+                                    // Unfortunately, its done at the word level and not the DWORD level.
+                                    var localPort = row.localPort1 << 8 | row.localPort2;
+                                    var remotePort = state == TcpState.Listen ? 0 : row.remotePort1 << 8 | row.remotePort2;
+
+                                    var info = new TcpInfo(state, new IPEndPoint(row.localAddr, localPort), new IPEndPoint(row.remoteAddr, remotePort));
+
+                                    handle(info);
+
+                                    span = span[sizeof(Interop.IpHlpApi.MibTcpRow)..];
+                                }
+                            }
+                        }
+                    }
+                    finally
+                    {
+                        Marshal.FreeHGlobal(buffer);
+                    }
+                }
+
+                // If we don't have any ipv4 interfaces detected, just continue.
+                if (result != Interop.IpHlpApi.ERROR_SUCCESS && result != Interop.IpHlpApi.ERROR_NO_DATA)
+                {
+                    throw new NetworkInformationException((int)result);
+                }
+            }
+
+            if (Socket.OSSupportsIPv6)
+            {
+                // Get the buffer size needed.
+                size = 0;
+                result = Interop.IpHlpApi.GetExtendedTcpTable(IntPtr.Zero, ref size, true,
+                                                                        (uint)AddressFamily.InterNetworkV6,
+                                                                        Interop.IpHlpApi.TcpTableClass.TcpTableOwnerPidAll, 0);
+
+                while (result == Interop.IpHlpApi.ERROR_INSUFFICIENT_BUFFER)
+                {
+                    // Allocate the buffer and get the TCP table.
+                    IntPtr buffer = Marshal.AllocHGlobal((int)size);
+                    try
+                    {
+                        result = Interop.IpHlpApi.GetExtendedTcpTable(buffer, ref size, true,
+                                                                                (uint)AddressFamily.InterNetworkV6,
+                                                                                Interop.IpHlpApi.TcpTableClass.TcpTableOwnerPidAll, 0);
+                        if (result == Interop.IpHlpApi.ERROR_SUCCESS)
+                        {
+                            var span = new ReadOnlySpan<byte>((byte*)buffer, (int)size);
+
+                            // The table info just gives us the number of rows.
+                            ref readonly Interop.IpHlpApi.MibTcp6TableOwnerPid tcpTable6OwnerPid = ref MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcp6TableOwnerPid>(span);
+
+                            if (tcpTable6OwnerPid.numberOfEntries > 0)
+                            {
+                                // Skip over the tableinfo to get the inline rows.
+                                span = span.Slice(sizeof(Interop.IpHlpApi.MibTcp6TableOwnerPid));
+
+                                for (var i = 0; i < tcpTable6OwnerPid.numberOfEntries; i++)
+                                {
+                                    var row = MemoryMarshal.AsRef<Interop.IpHlpApi.MibTcp6RowOwnerPid>(span);
+                                    var state = row.state;
+
+                                    // Port is returned in Big-Endian - most significant bit on left.
+                                    // Unfortunately, its done at the word level and not the DWORD level.
+                                    var localPort = row.localPort1 << 8 | row.localPort2;
+                                    var remotePort = ((state == TcpState.Listen) ? 0 : row.remotePort1 << 8 | row.remotePort2);
+
+                                    var info = new TcpInfo(
+                                        state,
+                                        new IPEndPoint(new IPAddress(row.localAddrAsSpan, row.localScopeId), localPort),
+                                        new IPEndPoint(new IPAddress(row.remoteAddrAsSpan, row.remoteScopeId), remotePort));
+
+                                    handle(info);
+                                    // We increment the pointer to the next row.
+                                    span = span.Slice(sizeof(Interop.IpHlpApi.MibTcp6RowOwnerPid));
+                                }
+                            }
+                        }
+                    }
+                    finally
+                    {
+                        Marshal.FreeHGlobal(buffer);
+                    }
+                }
+
+                // If we don't have any ipv6 interfaces detected, just continue.
+                if (result != Interop.IpHlpApi.ERROR_SUCCESS && result != Interop.IpHlpApi.ERROR_NO_DATA)
+                {
+                    throw new NetworkInformationException((int)result);
+                }
+            }
+        }
+        #endif
+    }
+
+    #if NET6_0_OR_GREATER
+internal static class Interop
+{
+    private class Libraries
+    {
+        public const string IpHlpApi = "iphlpapi.dll";
+    }
+
+    internal static class IpHlpApi
+    {
+        public const int ERROR_INSUFFICIENT_BUFFER = 0x007A;
+        public const int ERROR_SUCCESS = 0x0;
+        public const int ERROR_NO_DATA = 0xE8;
+
+        [StructLayout(LayoutKind.Sequential)]
+        internal struct MibTcpTable
+        {
+            internal uint numberOfEntries;
+        }
+
+        [StructLayout(LayoutKind.Sequential)]
+        internal struct MibTcpRow
+        {
+            internal TcpState state;
+            internal uint localAddr;
+            internal byte localPort1;
+            internal byte localPort2;
+            // Ports are only 16 bit values (in network WORD order, 3,4,1,2).
+            // There are reports where the high order bytes have garbage in them.
+            internal byte ignoreLocalPort3;
+            internal byte ignoreLocalPort4;
+            internal uint remoteAddr;
+            internal byte remotePort1;
+            internal byte remotePort2;
+            // Ports are only 16 bit values (in network WORD order, 3,4,1,2).
+            // There are reports where the high order bytes have garbage in them.
+            internal byte ignoreRemotePort3;
+            internal byte ignoreRemotePort4;
+        }
+
+        [StructLayout(LayoutKind.Sequential)]
+        internal struct MibTcp6TableOwnerPid
+        {
+            internal uint numberOfEntries;
+        }
+
+        [StructLayout(LayoutKind.Sequential)]
+        internal unsafe struct MibTcp6RowOwnerPid
+        {
+            internal fixed byte localAddr[16];
+            internal uint localScopeId;
+            internal byte localPort1;
+            internal byte localPort2;
+            // Ports are only 16 bit values (in network WORD order, 3,4,1,2).
+            // There are reports where the high order bytes have garbage in them.
+            internal byte ignoreLocalPort3;
+            internal byte ignoreLocalPort4;
+            internal fixed byte remoteAddr[16];
+            internal uint remoteScopeId;
+            internal byte remotePort1;
+            internal byte remotePort2;
+            // Ports are only 16 bit values (in network WORD order, 3,4,1,2).
+            // There are reports where the high order bytes have garbage in them.
+            internal byte ignoreRemotePort3;
+            internal byte ignoreRemotePort4;
+            internal TcpState state;
+            internal uint owningPid;
+
+            internal ReadOnlySpan<byte> localAddrAsSpan => MemoryMarshal.CreateSpan(ref localAddr[0], 16);
+            internal ReadOnlySpan<byte> remoteAddrAsSpan => MemoryMarshal.CreateSpan(ref remoteAddr[0], 16);
+        }
+
+        internal enum TcpTableClass
+        {
+            TcpTableBasicListener = 0,
+            TcpTableBasicConnections = 1,
+            TcpTableBasicAll = 2,
+            TcpTableOwnerPidListener = 3,
+            TcpTableOwnerPidConnections = 4,
+            TcpTableOwnerPidAll = 5,
+            TcpTableOwnerModuleListener = 6,
+            TcpTableOwnerModuleConnections = 7,
+            TcpTableOwnerModuleAll = 8
+        }
+
+        [DllImport(Libraries.IpHlpApi)]
+        internal static extern uint GetTcpTable(IntPtr pTcpTable, ref uint dwOutBufLen, bool order);
+
+        [DllImport(Libraries.IpHlpApi)]
+        internal static extern uint GetExtendedTcpTable(IntPtr pTcpTable, ref uint dwOutBufLen, bool order,
+                                                        uint IPVersion, TcpTableClass tableClass, uint reserved);
     }
+}
+#endif
 }