feat: Make NetworkMetricsPipelineStage public

Author: simon-lemay-unity

com.unity.netcode.gameobjects/CHANGELOG.md

@@ -10,6 +10,7 @@ Additional documentation and release notes are available at [Multiplayer Documen
 
 ### Added
 
+- The `NetworkMetricsPipelineStage` for Unity Transport is now part of the public API. This allows using it in custom implementations of `INetworkStreamDriverConstructor` that want to maintain compatibility with the multiplayer tools package.
 
 ### Changed
 

com.unity.netcode.gameobjects/Runtime/Transports/UTP/INetworkStreamDriverConstructor.cs

@@ -35,6 +35,8 @@ namespace Unity.Netcode.Transports.UTP
     ///         {
     ///             var settings = transport.GetDefaultNetworkSettings();
     ///             driver = NetworkDriver.Create(new IPCNetworkInterface(), settings);
+    /// 
+    ///             driver.RegisterPipelineStage(new NetworkMetricsPipelineStage());
     ///
     ///             transport.GetDefaultPipelineConfigurations(
     ///                 out var unreliableFragmentedPipelineStages,

com.unity.netcode.gameobjects/Runtime/Transports/UTP/NetworkMetricsPipelineStage.cs

@@ -7,27 +7,32 @@
 
 namespace Unity.Netcode.Transports.UTP
 {
+    /// <summary>
+    /// A pipeline stage that tracks some internal metrics that are then used by the multiplayer
+    /// tools package. This should only be used when creating a custom <see cref="NetworkDriver"/>
+    /// with <see cref="INetworkStreamDriverConstructor"/> if compatibility with the multiplayer
+    /// tools package is desired. In that situation, this stage needs to be registered with the
+    /// constructed driver with <see cref="NetworkDriver.RegisterPipelineStage{T}"/>.
+    /// </summary>
     [BurstCompile]
-    internal unsafe struct NetworkMetricsPipelineStage : INetworkPipelineStage
+    public unsafe struct NetworkMetricsPipelineStage : INetworkPipelineStage
     {
-        private static TransportFunctionPointer<NetworkPipelineStage.ReceiveDelegate> s_ReceiveFunction = new TransportFunctionPointer<NetworkPipelineStage.ReceiveDelegate>(Receive);
-        private static TransportFunctionPointer<NetworkPipelineStage.SendDelegate> s_SendFunction = new TransportFunctionPointer<NetworkPipelineStage.SendDelegate>(Send);
-        private static TransportFunctionPointer<NetworkPipelineStage.InitializeConnectionDelegate> s_InitializeConnectionFunction = new TransportFunctionPointer<NetworkPipelineStage.InitializeConnectionDelegate>(InitializeConnection);
-
+        /// <inheritdoc/>
         public NetworkPipelineStage StaticInitialize(byte* staticInstanceBuffer,
             int staticInstanceBufferLength,
             NetworkSettings settings)
         {
             return new NetworkPipelineStage(
-                s_ReceiveFunction,
-                s_SendFunction,
-                s_InitializeConnectionFunction,
+                new TransportFunctionPointer<NetworkPipelineStage.ReceiveDelegate>(Receive),
+                new TransportFunctionPointer<NetworkPipelineStage.SendDelegate>(Send),
+                new TransportFunctionPointer<NetworkPipelineStage.InitializeConnectionDelegate>(InitializeConnection),
                 ReceiveCapacity: 0,
                 SendCapacity: 0,
                 HeaderCapacity: 0,
                 SharedStateCapacity: UnsafeUtility.SizeOf<NetworkMetricsContext>());
         }
 
+        /// <inheritdoc/>
         public int StaticSize => 0;
 
         [BurstCompile(DisableDirectCall = true)]

com.unity.netcode.gameobjects/Runtime/Transports/UTP/UnityTransport.cs

@@ -1234,13 +1234,28 @@ private void ExtractNetworkMetricsFromPipeline(NetworkPipeline pipeline, Network
                 return;
             }
 
-            //Don't need to dispose of the buffers, they are filled with data pointers.
-            m_Driver.GetPipelineBuffers(pipeline,
-                NetworkPipelineStageId.Get<NetworkMetricsPipelineStage>(),
-                networkConnection,
-                out _,
-                out _,
-                out var sharedBuffer);
+            try
+            {
+                // Don't need to dispose of the buffers, they are filled with data pointers.
+                m_Driver.GetPipelineBuffers(pipeline,
+                    NetworkPipelineStageId.Get<NetworkMetricsPipelineStage>(),
+                    networkConnection,
+                    out _,
+                    out _,
+                    out var sharedBuffer);
+            }
+            catch (InvalidOperationException)
+            {
+                // Can happen if using a custom driver that isn't configured with the metrics stage.
+                return;
+            }
+
+            // That InvalidOperationException above is only thrown in the editor. In runtime builds
+            // we instead get default return values when the pipeline stage is invalid.
+            if (sharedBuffer == default)
+            {
+                return;
+            }
 
             unsafe
             {

com.unity.netcode.gameobjects/Tests/Editor/Transports/UnityTransportTests.cs

@@ -229,5 +229,44 @@ public void UnityTransport_HostnameValidation((string, bool) testCase)
             transport.Shutdown();
         }
 #endif
+
+        private class IPCDriverConstructor : INetworkStreamDriverConstructor
+        {
+            public void CreateDriver(
+                UnityTransport transport,
+                out NetworkDriver driver,
+                out NetworkPipeline unreliableFragmentedPipeline,
+                out NetworkPipeline unreliableSequencedFragmentedPipeline,
+                out NetworkPipeline reliableSequencedPipeline)
+            {
+                var settings = transport.GetDefaultNetworkSettings();
+                driver = NetworkDriver.Create(new IPCNetworkInterface(), settings);
+
+                driver.RegisterPipelineStage(new NetworkMetricsPipelineStage());
+
+                transport.GetDefaultPipelineConfigurations(
+                    out var unreliableFragmentedPipelineStages,
+                    out var unreliableSequencedFragmentedPipelineStages,
+                    out var reliableSequencedPipelineStages);
+
+                    unreliableFragmentedPipeline = driver.CreatePipeline(unreliableFragmentedPipelineStages);
+                    unreliableSequencedFragmentedPipeline = driver.CreatePipeline(unreliableSequencedFragmentedPipelineStages);
+                    reliableSequencedPipeline = driver.CreatePipeline(reliableSequencedPipelineStages);
+                }
+        }
+
+        [Test]
+        public void UnityTransport_CustomDriverConstructorWithDefaultPipelines()
+        {
+            UnityTransport transport = new GameObject().AddComponent<UnityTransport>();
+            UnityTransport.s_DriverConstructor = new IPCDriverConstructor();
+            transport.Initialize();
+
+            Assert.True(transport.StartServer());
+
+            transport.Shutdown();
+
+            UnityTransport.s_DriverConstructor = null;
+        }
     }
 }