RebootKit/Runtime/Engine/Code/Main/NetworkSystem.cs

using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Threading;
using Cysharp.Threading.Tasks;
using NUnit.Framework;
using RebootKit.Engine.Simulation;
using Unity.Collections;
using Unity.Netcode;
using UnityEngine;
using Logger = RebootKit.Engine.Foundation.Logger;

namespace RebootKit.Engine.Main {
    enum NetworkClientSyncState {
        NotReady,
        LoadingWorld,
        PreparingForActorsSync,
        SyncingActors,
        Ready
    }

    struct NetworkClientState : INetworkSerializable {
        public ulong ClientID;
        public NetworkClientSyncState SyncState;
        public int ActorsSyncPacketsLeft;

        public bool IsReady {
            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            get {
                return SyncState == NetworkClientSyncState.Ready;
            }
        }
        
        public void NetworkSerialize<T>(BufferSerializer<T> serializer) where T : IReaderWriter {
            serializer.SerializeValue(ref ClientID);
            serializer.SerializeValue(ref SyncState);
            serializer.SerializeValue(ref ActorsSyncPacketsLeft);
        }
    }

    public class NetworkSystem : NetworkBehaviour {
        static readonly Logger s_Logger = new Logger(nameof(NetworkSystem));

        [field: SerializeField] public ActorsManager Actors { get; private set; }

        internal readonly Dictionary<ulong, NetworkClientState> Clients = new Dictionary<ulong, NetworkClientState>();

        public FixedString512Bytes WorldID { get; private set; } = new FixedString512Bytes("");
        bool m_IsChangingWorld = false;

        public ulong LocalClientID {
            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            get {
                return NetworkManager.Singleton.LocalClientId;
            }
        }

        //
        // @MARK: Unity callbacks
        //
        void Awake() {
            RR.NetworkSystemInstance = this;
        }

        //
        // @MARK: NetworkBehaviour callbacks 
        //
        public override void OnNetworkSpawn() {
            base.OnNetworkSpawn();

            NetworkManager.Singleton.OnClientConnectedCallback += OnClientConnected;
            NetworkManager.Singleton.OnClientDisconnectCallback += OnClientDisconnect;
        }

        public override void OnNetworkDespawn() {
            base.OnNetworkDespawn();

            NetworkManager.Singleton.OnClientConnectedCallback -= OnClientConnected;
            NetworkManager.Singleton.OnClientDisconnectCallback -= OnClientDisconnect;
        }

        void OnClientConnected(ulong clientID) {
            if (!IsServer) {
                return;
            }

            s_Logger.Info($"OnClientConnected: {clientID}");

            NetworkClientState newClientState = new NetworkClientState {
                ClientID = clientID,
                SyncState = NetworkClientSyncState.NotReady
            };
            Clients.Add(clientID, newClientState);

            if (clientID != NetworkManager.Singleton.LocalClientId) {
                foreach (NetworkClientState state in Clients.Values) {
                    UpdateClientStateRpc(state, RpcTarget.Single(clientID, RpcTargetUse.Temp));
                }
            }

            if (!WorldID.IsEmpty) {
                s_Logger.Info($"Synchronizing world load for client {clientID} with world ID '{WorldID}'");
                ClientLoadWorldRpc(WorldID.ToString(), RpcTarget.Single(clientID, RpcTargetUse.Temp));
            }
        }

        void OnClientDisconnect(ulong clientID) {
            s_Logger.Info($"OnClientDisconnect: {clientID}");
            Clients.Remove(clientID);
        }

        //
        // @MARK: Server API
        //
        public void KickClient(ulong clientID, string reason = "Kicked by server") {
            if (!IsServer) {
                s_Logger.Error("Only server can kick clients.");
                return;
            }

            if (NetworkManager.Singleton.ConnectedClients.TryGetValue(clientID, out NetworkClient client)) {
                NetworkManager.Singleton.DisconnectClient(clientID, reason);
                s_Logger.Info($"Kicked client {clientID}: {reason}");
            } else {
                s_Logger.Error($"Client {clientID} not found.");
            }
        }

        public void SetCurrentWorld(string worldID) {
            if (!IsServer) {
                s_Logger.Error("Only server can set the current world.");
                return;
            }

            if (m_IsChangingWorld) {
                s_Logger.Error($"Already changing world to '{WorldID}'. Please wait until the current world change is complete.");
                return;
            }

            WorldConfigAsset worldConfigAsset = RR.GetWorldConfigAsset(worldID);
            if (worldConfigAsset is null) {
                s_Logger.Error($"Failed to set current world: World config asset for '{worldID}' not found.");
                return;
            }

            WorldID = worldID;
            
            foreach ((ulong _, NetworkClientState clientState) in Clients.ToList()) {
                NetworkClientState state = clientState;
                state.SyncState = NetworkClientSyncState.LoadingWorld;
                UpdateClientState(state);
            }

            ServerLoadWorldAsync(worldConfigAsset, destroyCancellationToken).Forget();
        }

        async UniTask ServerLoadWorldAsync(WorldConfigAsset asset, CancellationToken cancellationToken) {
            s_Logger.Info($"ServerLoadWorldAsync: {asset.Config.name}");

            m_IsChangingWorld = true;

            RR.World.Unload();
            RR.CloseMainMenu();

            await RR.World.LoadAsync(asset.Config, cancellationToken);

            m_IsChangingWorld = false;

            if (!TryGetClientState(NetworkManager.Singleton.LocalClientId, out NetworkClientState localClientState)) {
                s_Logger.Error($"Local client state not found for client ID {NetworkManager.Singleton.LocalClientId}.");
                RR.Disconnect();
                return;
            }

            localClientState.SyncState = NetworkClientSyncState.Ready;
            UpdateClientState(localClientState);

            RR.GameInstance.PlayerBecameReady(localClientState.ClientID);

            ClientLoadWorldRpc(asset.name, RpcTarget.NotMe);
        }

        [Rpc(SendTo.SpecifiedInParams, Delivery = RpcDelivery.Reliable)]
        void ClientLoadWorldRpc(string worldID, RpcParams rpcParams) {
            WorldConfigAsset worldConfigAsset = RR.GetWorldConfigAsset(worldID);
            if (worldConfigAsset is null) {
                s_Logger.Error($"World config asset for '{worldID}' not found.");
                RR.Disconnect();
                return;
            }

            ClientLoadWorldAsync(worldID, destroyCancellationToken).Forget();
        }

        async UniTask ClientLoadWorldAsync(string worldID, CancellationToken cancellationToken) {
            s_Logger.Info($"ClientLoadWorldAsync: {worldID}");

            WorldConfigAsset worldConfigAsset = RR.GetWorldConfigAsset(worldID);
            if (worldConfigAsset is null) {
                s_Logger.Error($"World config asset for '{worldID}' not found.");
                return;
            }

            RR.World.Unload();
            RR.CloseMainMenu();

            await RR.World.LoadAsync(worldConfigAsset.Config, cancellationToken);

            WorldID = worldID;
            ClientLoadedWorldRpc(worldID);
        }

        [Rpc(SendTo.Server, Delivery = RpcDelivery.Reliable)]
        void ClientLoadedWorldRpc(string worldID, RpcParams rpcParams = default) {
            ulong clientID = rpcParams.Receive.SenderClientId;

            if (!WorldID.Equals(worldID)) {
                s_Logger.Error($"Client {clientID} tried to load world '{worldID}', but server is in world '{WorldID}'.");
                NetworkManager.Singleton.DisconnectClient(clientID, "World mismatch!");
                return;
            }

            if (Clients.TryGetValue(clientID, out NetworkClientState clientState)) {
                Actors.InitializeActorsForClient(clientID);
            } else {
                NetworkManager.Singleton.DisconnectClient(clientID, "Client is not registered!");
            }
        }

        //
        // @MARK: Internal
        //
        internal bool TryGetClientState(ulong clientID, out NetworkClientState clientState) {
            return Clients.TryGetValue(clientID, out clientState);
        }

        internal void UpdateClientState(NetworkClientState clientState) {
            if (!IsServer) {
                s_Logger.Error("UpdateClientState can only be called on the server.");
                return;
            }
            
            Clients[clientState.ClientID] = clientState;
            UpdateClientStateRpc(clientState, RpcTarget.NotServer);
        }
        
        [Rpc(SendTo.SpecifiedInParams, Delivery = RpcDelivery.Reliable)]
        void UpdateClientStateRpc(NetworkClientState newState, RpcParams rpcParams) {
            Clients[newState.ClientID] = newState;
        }

        internal void ClientSynchronizedActors(ulong clientID) {
            if (TryGetClientState(clientID, out NetworkClientState state)) {
                state.SyncState = NetworkClientSyncState.Ready;
                UpdateClientState(state);

                RR.GameInstance.PlayerBecameReady(clientID);
            } else {
                s_Logger.Error($"Client state for {clientID} not found.");
            }
        }

        internal int GetReadyClientsCount() {
            int count = 0;
            foreach (NetworkClientState clientState in Clients.Values) {
                if (clientState.IsReady) {
                    count++;
                }
            }
            return count;
        }
    }
}