RebootKit/Runtime/Engine/Code/Network/NetworkPacketQueue.cs

using System;
using System.Collections.Generic;
using RebootKit.Engine.Foundation;
using RebootKit.Engine.Simulation;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine.Assertions;
using UnityEngine.Pool;

namespace RebootKit.Engine.Network {
    struct NetworkPacketHeader : ISerializableEntity {
        public int MagicNumber;
        public ushort Version;
        public ushort EntityCount;

        public static int GetEntityCountOffset() {
            return sizeof(int) + sizeof(ushort);
        }

        public void Serialize(NetworkBufferWriter writer) {
            writer.Write(MagicNumber);
            writer.Write(Version);
            writer.Write(EntityCount);
        }

        public void Deserialize(NetworkBufferReader reader) {
            reader.Read(out MagicNumber);
            reader.Read(out Version);
            reader.Read(out EntityCount);
        }

        public int GetMaxBytes() {
            return sizeof(int) + sizeof(ushort) * 2; // MagicNumber, Version, EntityCount
        }
    }

    class NetworkPacket : IDisposable {
        public static readonly IObjectPool<NetworkPacket> Pool = new ObjectPool<NetworkPacket>(
             () => {
                 NetworkPacket packet = new NetworkPacket();
                 packet.Data = default;
                 packet.Writer = default;
                 return packet;
             },
             packet => {
                 // Packet is initialized after being retrieved from the pool
             },
             packet => {
                 packet.Dispose();
             },
             packet => {
                 packet.Dispose();
             },
             true,
             16
            );

        public NativeArray<byte> Data;
        public NetworkBufferWriter Writer;

        public ushort EntityCount { get; private set; }

        public void IncrementEntityCount() {
            int originalPosition = Writer.Position;

            EntityCount += 1;

            Writer.Position = NetworkPacketHeader.GetEntityCountOffset(); // Reset position to write the entity count
            Writer.Write(EntityCount);
            Writer.Position = originalPosition;
        }

        public void Dispose() {
            Data.Dispose();
            Writer.Dispose();
            
            EntityCount = 0;
        }
    }
    
    enum NetworkDataType : byte {
        None = 0x00,
        ActorCoreState = 0x01,
        ActorTransformSync = 0x02,
        ActorState = 0x03,
        ActorEvent = 0x04,
        ActorCommand = 0x05,
        SynchronizeActor = 0x07,
        SpawnActor = 0x08,
    }

    struct NetworkDataHeader : ISerializableEntity {
        public NetworkDataType Type;
        public ulong ActorID;
        public int DataSize;
        
        public int GetMaxBytes() {
            return sizeof(ulong) + sizeof(byte) + sizeof(int);
        }
        
        public void Serialize(NetworkBufferWriter writer) {
            writer.Write((byte) Type);
            writer.Write(ActorID);
            writer.Write(DataSize);
        }

        public void Deserialize(NetworkBufferReader reader) {
            reader.Read(out byte typeByte);
            Type = (NetworkDataType) typeByte;
            reader.Read(out ActorID);
            reader.Read(out DataSize);
        }
    }

    class NetworkPacketQueue : IDisposable {
        static readonly Logger s_Logger = new Logger(nameof(NetworkPacketQueue));

        readonly int m_PacketMaxSize;
        readonly ushort m_Version;

        internal readonly List<NetworkPacket> NetworkPackets = new List<NetworkPacket>();

        public NetworkPacketQueue(int packetMaxSize, ushort version = 1) {
            m_PacketMaxSize = packetMaxSize;
            m_Version = version;
            Assert.IsTrue(m_PacketMaxSize > 0, "Packet maximum size must be greater than zero.");
        }

        public void Dispose() {
            foreach (NetworkPacket packet in NetworkPackets) {
                packet.Data.Dispose();
            }

            NetworkPackets.Clear();
        }

        public void Clear() {
            foreach (NetworkPacket packet in NetworkPackets) {
                packet.Dispose();
            }

            NetworkPackets.Clear();
        }

        public void WriteActorState(ulong actorID, IActorData entity) {
            Assert.IsTrue(entity.GetMaxBytes() <= m_PacketMaxSize,
                          $"Entity size {entity.GetMaxBytes()} exceeds packet max size {m_PacketMaxSize}.");

            NetworkDataHeader header = new NetworkDataHeader {
                Type = NetworkDataType.ActorState,
                ActorID = actorID,
                DataSize = entity.GetMaxBytes()
            };

            int bytesToWrite = header.GetMaxBytes() + entity.GetMaxBytes();

            NetworkPacket packet = GetPacketToWriteTo(bytesToWrite);
            header.Serialize(packet.Writer);
            entity.Serialize(packet.Writer);
            packet.IncrementEntityCount();
        }

        public void WriteActorTransformState(ulong actorID, ActorTransformSyncData transformData) {
            NetworkDataHeader header = new NetworkDataHeader {
                Type = NetworkDataType.ActorTransformSync,
                ActorID = actorID,
                DataSize = transformData.GetMaxBytes()
            };

            int bytesToWrite = header.GetMaxBytes() + transformData.GetMaxBytes();

            NetworkPacket packet = GetPacketToWriteTo(bytesToWrite);
            header.Serialize(packet.Writer);
            transformData.Serialize(packet.Writer);
            packet.IncrementEntityCount();
        }

        public void WriteActorCoreState(ulong actorID, ActorCoreStateSnapshot coreState) {
            NetworkDataHeader header = new NetworkDataHeader {
                Type = NetworkDataType.ActorCoreState,
                ActorID = actorID,
                DataSize = coreState.GetMaxBytes()
            };

            int bytesToWrite = header.GetMaxBytes() + coreState.GetMaxBytes();

            NetworkPacket packet = GetPacketToWriteTo(bytesToWrite);
            header.Serialize(packet.Writer);
            coreState.Serialize(packet.Writer);
            packet.IncrementEntityCount();
        }

        public void WriteSpawnActor(FixedString64Bytes assetGUID,
                                    ulong actorID,
                                    ActorCoreStateSnapshot coreState,
                                    IActorData actorData) {
            NetworkDataHeader header = new NetworkDataHeader {
                Type = NetworkDataType.SpawnActor,
                ActorID = actorID,
                DataSize = 0
            };

            header.DataSize += sizeof(byte) * 64; // assetGUID
            header.DataSize += coreState.GetMaxBytes();
            header.DataSize += sizeof(ushort);
            header.DataSize += actorData.GetMaxBytes();

            NetworkPacket packet = GetPacketToWriteTo(header.GetMaxBytes() + header.DataSize);
            header.Serialize(packet.Writer);

            packet.Writer.Write(assetGUID);
            coreState.Serialize(packet.Writer);

            packet.Writer.Write((ushort) actorData.GetMaxBytes());
            actorData.Serialize(packet.Writer);

            packet.IncrementEntityCount();
        }

        public void WriteActorSynchronize(ulong actorID,
                                          ActorCoreStateSnapshot coreState,
                                          IActorData actorData) {
            NetworkDataHeader header = new NetworkDataHeader {
                Type = NetworkDataType.SynchronizeActor,
                ActorID = actorID,
                DataSize = 0
            };

            header.DataSize += coreState.GetMaxBytes();
            header.DataSize += sizeof(ushort);
            header.DataSize += actorData.GetMaxBytes();

            NetworkPacket packet = GetPacketToWriteTo(header.GetMaxBytes() + header.DataSize);
            header.Serialize(packet.Writer);

            coreState.Serialize(packet.Writer);

            packet.Writer.Write((ushort) actorData.GetMaxBytes());
            actorData.Serialize(packet.Writer);

            packet.IncrementEntityCount();
        }

        public void WriteActorEvent(ActorEvent actorEvent) {
            NetworkDataHeader header = new NetworkDataHeader {
                Type = NetworkDataType.ActorEvent,
                ActorID = actorEvent.ActorID,
                DataSize = actorEvent.GetMaxBytes()
            };

            NetworkPacket packet = GetPacketToWriteTo(header.GetMaxBytes() + header.DataSize);
            header.Serialize(packet.Writer);
            actorEvent.Serialize(packet.Writer);

            packet.IncrementEntityCount();           
        }
        
        public void WriteActorCommand(ActorCommand actorCommand) {
            NetworkDataHeader header = new NetworkDataHeader {
                Type = NetworkDataType.ActorCommand,
                ActorID = actorCommand.ActorID,
                DataSize = actorCommand.GetMaxBytes()
            };

            NetworkPacket packet = GetPacketToWriteTo(header.GetMaxBytes() + header.DataSize);
            header.Serialize(packet.Writer);
            actorCommand.Serialize(packet.Writer);

            packet.IncrementEntityCount();
        }

        NetworkPacket GetPacketToWriteTo(int bytesToWrite) {
            foreach (NetworkPacket networkPacket in NetworkPackets) {
                if (networkPacket.Writer.WillFit(bytesToWrite)) {
                    return networkPacket;
                }
            }
            
            Assert.IsTrue(bytesToWrite < m_PacketMaxSize,
                          $"Packet size {bytesToWrite} exceeds maximum allowed size {m_PacketMaxSize}.");

            NetworkPacket packet = NetworkPacket.Pool.Get();
            packet.Data = new NativeArray<byte>(m_PacketMaxSize, Allocator.Persistent);

            unsafe {
                void* ptr = packet.Data.GetUnsafePtr();
                UnsafeUtility.MemClear(ptr, sizeof(byte) * packet.Data.Length);
            }

            packet.Writer = new NetworkBufferWriter(packet.Data, 0);

            NetworkPacketHeader header = new NetworkPacketHeader {
                MagicNumber = RConsts.k_NetworkPacketMagicNumber,
                Version = m_Version,
                EntityCount = 0 // Will be updated later
            };

            header.Serialize(packet.Writer);
            NetworkPackets.Add(packet);
            return packet;
        }
    }
}