fix: NetworkTransport early and post late updates

com.unity.netcode.gameobjects/Runtime/Core/NetworkManager.cs

@@ -291,6 +291,10 @@ public void NetworkUpdate(NetworkUpdateStage updateStage)
                 case NetworkUpdateStage.EarlyUpdate:
                     {
                         UpdateTopology();
+
+                        // Handle processing any new connections or transport events
+                        NetworkConfig.NetworkTransport.EarlyUpdate();
+
                         ConnectionManager.ProcessPendingApprovals();
                         ConnectionManager.PollAndHandleNetworkEvents();
 
@@ -379,6 +383,9 @@ public void NetworkUpdate(NetworkUpdateStage updateStage)
                         // Metrics update needs to be driven by NetworkConnectionManager's update to assure metrics are dispatched after the send queue is processed.
                         MetricsManager.UpdateMetrics();
 
+                        // Handle sending any pending transport messages 
+                        NetworkConfig.NetworkTransport.PostLateUpdate();
+
                         // TODO: Determine a better way to handle this
                         NetworkObject.VerifyParentingStatus();
 

com.unity.netcode.gameobjects/Runtime/Transports/NetworkTransport.cs

@@ -107,6 +107,22 @@ protected void InvokeOnTransportEvent(NetworkEvent eventType, ulong clientId, Ar
         /// /// <param name="networkManager">optionally pass in NetworkManager</param>
         public abstract void Initialize(NetworkManager networkManager = null);
 
+        /// <summary>
+        /// Invoked by NetworkManager at the beginning of its EarlyUpdate
+        /// </summary>
+        internal virtual void EarlyUpdate()
+        {
+
+        }
+
+        /// <summary>
+        /// Invoked by NetworkManager during PostLateUpdate
+        /// </summary>
+        internal virtual void PostLateUpdate()
+        {
+
+        }
+
         protected virtual NetworkTopologyTypes OnCurrentTopology()
         {
             return NetworkTopologyTypes.ClientServer;

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

@@ -944,17 +944,10 @@ private bool ProcessEvent()
             return false;
         }
 
-        private void Update()
+        internal override void EarlyUpdate()
         {
             if (m_Driver.IsCreated)
             {
-                foreach (var kvp in m_SendQueue)
-                {
-                    SendBatchedMessages(kvp.Key, kvp.Value);
-                }
-
-                m_Driver.ScheduleUpdate().Complete();
-
                 if (m_ProtocolType == ProtocolType.RelayUnityTransport && m_Driver.GetRelayConnectionStatus() == RelayConnectionStatus.AllocationInvalid)
                 {
                     Debug.LogError("Transport failure! Relay allocation needs to be recreated, and NetworkManager restarted. " +
@@ -964,6 +957,8 @@ private void Update()
                     return;
                 }
 
+                m_Driver.ScheduleUpdate().Complete();
+
                 while (AcceptConnection() && m_Driver.IsCreated)
                 {
                     ;
@@ -973,6 +968,21 @@ private void Update()
                 {
                     ;
                 }
+            }
+        }
+
+        internal override void PostLateUpdate()
+        {
+            if (m_Driver.IsCreated)
+            {
+                foreach (var kvp in m_SendQueue)
+                {
+                    SendBatchedMessages(kvp.Key, kvp.Value);
+                }
+
+                // Schedule a flush send as the last transport action for the
+                // current frame.
+                m_Driver.ScheduleFlushSend(default).Complete();
 
 #if MULTIPLAYER_TOOLS_1_0_0_PRE_7
                 if (m_NetworkManager)
@@ -981,6 +991,7 @@ private void Update()
                 }
 #endif
             }
+            base.PostLateUpdate();
         }
 
         private void OnDestroy()
@@ -1452,6 +1463,11 @@ public override void Shutdown()
         {
             if (m_Driver.IsCreated)
             {
+                while (ProcessEvent() && m_Driver.IsCreated)
+                {
+                    ;
+                }
+
                 // Flush all send queues to the network. NGO can be configured to flush its message
                 // queue on shutdown. But this only calls the Send() method, which doesn't actually
                 // get anything to the network.

com.unity.netcode.gameobjects/Tests/Runtime/Transports/UnityTransportTestHelpers.cs

@@ -16,21 +16,56 @@ internal static class UnityTransportTestHelpers
         // so we allow even more time for these platforms.
         public const float MaxNetworkEventWaitTime = 0.5f;
 
+        private static List<UnityTransport> s_UnityTransports = new List<UnityTransport>();
+        public static void RegisterTransportInstance(UnityTransport instance)
+        {
+            s_UnityTransports.Add(instance);
+        }
+
+        public static void DeregisterTransportInstance(UnityTransport instance)
+        {
+            s_UnityTransports.Remove(instance);
+        }
+
+        public static void ClearRegisteredTransportInstances()
+        {
+            s_UnityTransports.Clear();
+        }
+
+        public static void InvokeEarlyUpdate()
+        {
+            foreach (var transport in s_UnityTransports)
+            {
+                transport.EarlyUpdate();
+            }
+        }
+
+        public static void InvokePostLateUpdate()
+        {
+            foreach (var transport in s_UnityTransports)
+            {
+                transport.PostLateUpdate();
+            }
+        }
+
         // Wait for an event to appear in the given event list (must be the very next event).
         public static IEnumerator WaitForNetworkEvent(NetworkEvent type, List<TransportEvent> events, float timeout = MaxNetworkEventWaitTime)
         {
             int initialCount = events.Count;
             float startTime = Time.realtimeSinceStartup;
-
+            var waitUntilEndofFrame = new WaitForEndOfFrame();
+            var waitForNextFrame = new WaitForFixedUpdate();
             while (Time.realtimeSinceStartup - startTime < timeout)
             {
+                yield return waitForNextFrame;
+                InvokeEarlyUpdate();
                 if (events.Count > initialCount)
                 {
                     Assert.AreEqual(type, events[initialCount].Type);
                     yield break;
                 }
-
-                yield return new WaitForSeconds(0.01f);
+                yield return waitUntilEndofFrame;
+                InvokePostLateUpdate();
             }
 
             Assert.Fail("Timed out while waiting for network event.");
@@ -55,6 +90,7 @@ public static void InitializeTransport(out UnityTransport transport, out List<Tr
             }
 
             transport.Initialize();
+            RegisterTransportInstance(transport);
         }
 
         // Information about an event generated by a transport (basically just the parameters that

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

@@ -39,6 +39,7 @@ public IEnumerator Cleanup()
         {
             if (m_Server)
             {
+                DeregisterTransportInstance(m_Server);
                 m_Server.Shutdown();
 
                 // Need to destroy the GameObject (all assigned components will get destroyed too)
@@ -47,6 +48,7 @@ public IEnumerator Cleanup()
 
             if (m_Client1)
             {
+                DeregisterTransportInstance(m_Client1);
                 m_Client1.Shutdown();
 
                 // Need to destroy the GameObject (all assigned components will get destroyed too)
@@ -55,12 +57,13 @@ public IEnumerator Cleanup()
 
             if (m_Client2)
             {
+                DeregisterTransportInstance(m_Client2);
                 m_Client2.Shutdown();
 
                 // Need to destroy the GameObject (all assigned components will get destroyed too)
                 UnityEngine.Object.DestroyImmediate(m_Client2.gameObject);
             }
-
+            ClearRegisteredTransportInstances();
             m_ServerEvents?.Clear();
             m_Client1Events?.Clear();
             m_Client2Events?.Clear();
@@ -361,9 +364,26 @@ public IEnumerator SendCompletesOnUnreliableSendQueueOverflow()
                 m_Client1.Send(m_Client1.ServerClientId, payload, NetworkDelivery.Unreliable);
             }
 
+            var waitUntilEndofFrame = new WaitForEndOfFrame();
+            var waitForNextFrame = new WaitForFixedUpdate();
+            var timeout = Time.realtimeSinceStartup + 8.0f;
+
             // Manually wait. This ends up generating quite a bit of packets and it might take a
             // while for everything to make it to the server.
-            yield return new WaitForSeconds(numSends * 0.02f);
+            // Updated: We have to simulate the EarlyUpdate and PostLateUpdate invocations by NetworkManager
+            // in order to process connnects, events, and send messages.
+            while (timeout > Time.realtimeSinceStartup)
+            {
+                yield return waitForNextFrame;
+                InvokeEarlyUpdate();
+                yield return waitUntilEndofFrame;
+                InvokePostLateUpdate();
+                // Once we hit the expected count, exit the loop
+                if ((numSends + 1) == m_ServerEvents.Count)
+                {
+                    break;
+                }
+            }
 
             // Extra event is the connect event.
             Assert.AreEqual(numSends + 1, m_ServerEvents.Count);