Fix a bug about chunk loader (#36)

Author: sunnycase

src/MineCase.Gateway/Network/ClientSession.cs

@@ -1,4 +1,5 @@
 using Microsoft.Extensions.ObjectPool;
+using MineCase.Buffers;
 using MineCase.Protocol;
 using MineCase.Server.Network;
 using Orleans;
@@ -23,12 +24,14 @@ class ClientSession : IDisposable
         private readonly OutcomingPacketObserver _outcomingPacketObserver;
         private readonly ActionBlock<UncompressedPacket> _outcomingPacketDispatcher;
         private readonly ObjectPool<UncompressedPacket> _uncompressedPacketObjectPool;
+        private readonly IBufferPool<byte> _bufferPool;
 
-        public ClientSession(TcpClient tcpClient, IGrainFactory grainFactory, ObjectPool<UncompressedPacket> uncompressedPacketObjectPool)
+        public ClientSession(TcpClient tcpClient, IGrainFactory grainFactory, IBufferPool<byte> bufferPool, ObjectPool<UncompressedPacket> uncompressedPacketObjectPool)
         {
             _sessionId = Guid.NewGuid();
             _tcpClient = tcpClient;
             _grainFactory = grainFactory;
+            _bufferPool = bufferPool;
             _uncompressedPacketObjectPool = uncompressedPacketObjectPool;
             _outcomingPacketObserver = new OutcomingPacketObserver(this);
             _outcomingPacketDispatcher = new ActionBlock<UncompressedPacket>(SendOutcomingPacket);
@@ -57,35 +60,43 @@ public async Task Startup(CancellationToken cancellationToken)
 
         private void OnClosed()
         {
-            _outcomingPacketDispatcher.Complete();
-            _tcpClient.Client.Shutdown(SocketShutdown.Send);
+            _outcomingPacketDispatcher.Post(null);
         }
 
         private async Task DispatchIncomingPacket()
         {
-            var packet = _uncompressedPacketObjectPool.Get();
-            try
+            using (var bufferScope = _bufferPool.CreateScope())
             {
-                if (_useCompression)
+                var packet = _uncompressedPacketObjectPool.Get();
+                try
                 {
-                    var compressedPacket = await CompressedPacket.DeserializeAsync(_remoteStream, null);
-                    packet = PacketCompress.Decompress(ref compressedPacket);
+                    if (_useCompression)
+                    {
+                        var compressedPacket = await CompressedPacket.DeserializeAsync(_remoteStream, null);
+                        packet = PacketCompress.Decompress(ref compressedPacket);
+                    }
+                    else
+                    {
+                        packet = await UncompressedPacket.DeserializeAsync(_remoteStream, bufferScope, packet);
+                    }
+                    await DispatchIncomingPacket(packet);
                 }
-                else
+                finally
                 {
-                    packet = await UncompressedPacket.DeserializeAsync(_remoteStream, packet);
+                    _uncompressedPacketObjectPool.Return(packet);
                 }
-                await DispatchIncomingPacket(packet);
-            }
-            finally
-            {
-                _uncompressedPacketObjectPool.Return(packet);
             }
         }
 
         private async Task SendOutcomingPacket(UncompressedPacket packet)
         {
-            if (_useCompression)
+            // Close
+            if(packet == null)
+            {
+                _tcpClient.Client.Shutdown(SocketShutdown.Send);
+                _outcomingPacketDispatcher.Complete();
+            }
+            else if (_useCompression)
             {
                 var newPacket = PacketCompress.Compress(ref packet);
                 await newPacket.SerializeAsync(_remoteStream);

src/MineCase.Gateway/Network/ConnectionRouter.cs

@@ -1,5 +1,6 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.ObjectPool;
+using MineCase.Buffers;
 using MineCase.Protocol;
 using MineCase.Server.Settings;
 using Orleans;
@@ -18,12 +19,14 @@ class ConnectionRouter
         private readonly TcpListener _listener;
         private readonly IGrainFactory _grainFactory;
         private readonly ILogger _logger;
+        private readonly IBufferPool<byte> _bufferPool;
         private readonly ObjectPool<UncompressedPacket> _uncompressedPacketObjectPool;
 
-        public ConnectionRouter(IGrainFactory grainFactory, ILoggerFactory loggerFactory, ObjectPool<UncompressedPacket> uncompressedPacketObjectPool)
+        public ConnectionRouter(IGrainFactory grainFactory, ILoggerFactory loggerFactory, IBufferPool<byte> bufferPool, ObjectPool<UncompressedPacket> uncompressedPacketObjectPool)
         {
             _grainFactory = grainFactory;
             _logger = loggerFactory.CreateLogger<ConnectionRouter>();
+            _bufferPool = bufferPool;
             _uncompressedPacketObjectPool = uncompressedPacketObjectPool;
             _listener = new TcpListener(new IPEndPoint(IPAddress.Any, 25565));
         }
@@ -44,7 +47,7 @@ private async void DispatchIncomingClient(TcpClient tcpClient, CancellationToken
             try
             {
                 _logger.LogInformation($"Incoming connection from {tcpClient.Client.RemoteEndPoint}.");
-                using (var session = new ClientSession(tcpClient, _grainFactory, _uncompressedPacketObjectPool))
+                using (var session = new ClientSession(tcpClient, _grainFactory, _bufferPool, _uncompressedPacketObjectPool))
                 {
                     await session.Startup(cancellationToken);
                 }

src/MineCase.Gateway/Program.cs

@@ -8,6 +8,8 @@
 using MineCase.Gateway.Network;
 using Microsoft.Extensions.ObjectPool;
 using MineCase.Protocol;
+using MineCase.Buffers;
+using System.Buffers;
 
 namespace MineCase.Gateway
 {
@@ -55,6 +57,7 @@ private static void ConfigureObjectPools(IServiceCollection services)
                 var provider = s.GetRequiredService<ObjectPoolProvider>();
                 return provider.Create<UncompressedPacket>();
             });
+            services.AddSingleton<IBufferPool<byte>>(s => new BufferPool<byte>(ArrayPool<byte>.Shared));
         }
 
         private static IConfiguration LoadConfiguration()

src/MineCase.Protocol/Buffers/IBufferPool.cs

@@ -0,0 +1,68 @@
+using System;
+using System.Buffers;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Text;
+
+namespace MineCase.Buffers
+{
+    public interface IBufferPool<T>
+    {
+        IBufferPoolScope<T> CreateScope();
+    }
+
+    public interface IBufferPoolScope<T> : IDisposable
+    {
+        ArraySegment<T> Rent(int length);
+    }
+
+    public class BufferPool<T> : IBufferPool<T>
+    {
+        private readonly ArrayPool<T> _arrayPool;
+        private readonly ConcurrentBag<BufferPoolScope> _scopes = new ConcurrentBag<BufferPoolScope>();
+
+        public BufferPool(ArrayPool<T> arrayPool)
+        {
+            _arrayPool = arrayPool;
+        }
+
+        public IBufferPoolScope<T> CreateScope()
+        {
+            if (!_scopes.TryTake(out var scope))
+                scope = new BufferPoolScope(this, _arrayPool);
+            return scope;
+        }
+
+        private void Return(BufferPoolScope scope)
+        {
+            _scopes.Add(scope);
+        }
+
+        private class BufferPoolScope : IBufferPoolScope<T>
+        {
+            private readonly BufferPool<T> _bufferPool;
+            private readonly ArrayPool<T> _arrayPool;
+            private readonly ConcurrentBag<T[]> _rents = new ConcurrentBag<T[]>();
+
+            public BufferPoolScope(BufferPool<T> bufferPool, ArrayPool<T> arrayPool)
+            {
+                _bufferPool = bufferPool;
+                _arrayPool = arrayPool;
+            }
+
+            public void Dispose()
+            {
+                while (_rents.TryTake(out var rent))
+                    _arrayPool.Return(rent);
+                _bufferPool.Return(this);
+            }
+
+            public ArraySegment<T> Rent(int length)
+            {
+                var rent = _arrayPool.Rent(length);
+                _rents.Add(rent);
+                return new ArraySegment<T>(rent, 0, length);
+            }
+        }
+    }
+}

src/MineCase.Protocol/Protocol/Packet.cs

@@ -4,6 +4,7 @@
 using System.Runtime.ExceptionServices;
 using System.Text;
 using System.Threading.Tasks;
+using MineCase.Buffers;
 using MineCase.Serialization;
 using Orleans.Concurrency;
 
@@ -19,11 +20,11 @@ public class UncompressedPacket
         public uint PacketId;
 
         [SerializeAs(DataType.ByteArray)]
-        public byte[] Data;
+        public ArraySegment<byte> Data;
 
         public async Task SerializeAsync(Stream stream)
         {
-            Length = (uint)Data.Length + PacketId.SizeOfVarInt();
+            Length = (uint)Data.Count + PacketId.SizeOfVarInt();
 
             using (var bw = new BinaryWriter(stream, Encoding.UTF8, true))
             {
@@ -32,10 +33,10 @@ public async Task SerializeAsync(Stream stream)
                 bw.Flush();
             }
 
-            await stream.WriteAsync(Data, 0, Data.Length);
+            await stream.WriteAsync(Data.Array, Data.Offset, Data.Count);
         }
 
-        public static async Task<UncompressedPacket> DeserializeAsync(Stream stream, UncompressedPacket packet = null)
+        public static async Task<UncompressedPacket> DeserializeAsync(Stream stream, IBufferPoolScope<byte> bufferPool, UncompressedPacket packet = null)
         {
             packet = packet ?? new UncompressedPacket();
             int packetIdLen;
@@ -45,8 +46,8 @@ public static async Task<UncompressedPacket> DeserializeAsync(Stream stream, Unc
                 packet.PacketId = br.ReadAsVarInt(out packetIdLen);
             }
 
-            packet.Data = new byte[packet.Length - packetIdLen];
-            await stream.ReadExactAsync(packet.Data, 0, packet.Data.Length);
+            packet.Data = bufferPool.Rent((int)(packet.Length - packetIdLen));
+            await stream.ReadExactAsync(packet.Data.Array, packet.Data.Offset, packet.Data.Count);
             return packet;
         }
     }

src/MineCase.Protocol/Protocol/PacketCompress.cs

@@ -11,6 +11,8 @@ public static class PacketCompress
     {
         public static UncompressedPacket Decompress(ref CompressedPacket packet)
         {
+            throw new NotImplementedException();
+            /*
             var targetPacket = default(UncompressedPacket);
             using (var br = new BinaryReader(new DeflateStream(new MemoryStream(packet.CompressedData), CompressionMode.Decompress)))
             {
@@ -19,11 +21,13 @@ public static UncompressedPacket Decompress(ref CompressedPacket packet)
             }
 
             targetPacket.Length = packet.DataLength;
-            return targetPacket;
+            return targetPacket;*/
         }
 
         public static CompressedPacket Compress(ref UncompressedPacket packet)
         {
+            throw new NotImplementedException();
+            /*
             var targetPacket = default(CompressedPacket);
             using (var stream = new MemoryStream())
             using (var bw = new BinaryWriter(new DeflateStream(stream, CompressionMode.Compress)))
@@ -36,7 +40,7 @@ public static CompressedPacket Compress(ref UncompressedPacket packet)
                 targetPacket.CompressedData = stream.ToArray();
             }
 
-            return targetPacket;
+            return targetPacket;*/
         }
     }
 }

src/MineCase.Server.Grains/Game/ChunkSenderGrain.cs

@@ -19,7 +19,7 @@ public override Task OnActivateAsync()
             return base.OnActivateAsync();
         }
 
-        public Task PostChunk(int x, int z, IReadOnlyCollection<IClientboundPacketSink> clients, IReadOnlyCollection<IUser> users)
+        public Task PostChunk(int x, int z, IReadOnlyCollection<IClientboundPacketSink> clients, IReadOnlyCollection<IUserChunkLoader> loaders)
         {
             var stream = GetStreamProvider(StreamProviders.JobsProvider).GetStream<SendChunkJob>(_jobWorkerId, StreamProviders.Namespaces.ChunkSender);
             stream.OnNextAsync(new SendChunkJob
@@ -28,7 +28,7 @@ public Task PostChunk(int x, int z, IReadOnlyCollection<IClientboundPacketSink>
                 ChunkX = x,
                 ChunkZ = z,
                 Clients = clients,
-                Users = users
+                Loaders = loaders
             }).Ignore();
             return Task.CompletedTask;
         }

src/MineCase.Server.Grains/Game/ChunkSenderJobWorker.cs

@@ -24,7 +24,7 @@ public sealed class SendChunkJob
 
         public IReadOnlyCollection<IClientboundPacketSink> Clients { get; set; }
 
-        public IReadOnlyCollection<IUser> Users { get; set; }
+        public IReadOnlyCollection<IUserChunkLoader> Loaders { get; set; }
     }
 
     internal interface IChunkSenderJobWorker : IGrainWithGuidKey
@@ -53,8 +53,8 @@ private async Task OnNextAsync(SendChunkJob job, StreamSequenceToken token)
 
             var generator = new ClientPlayPacketGenerator(new BroadcastPacketSink(job.Clients, _packetPackager));
             await generator.ChunkData(Dimension.Overworld, job.ChunkX, job.ChunkZ, await chunkColumn.GetState());
-            foreach (var user in job.Users)
-                user.OnChunkSent(job.ChunkX, job.ChunkZ).Ignore();
+            foreach (var loader in job.Loaders)
+                loader.OnChunkSent(job.ChunkX, job.ChunkZ).Ignore();
         }
 
         private class BroadcastPacketSink : IPacketSink

src/MineCase.Server.Grains/MineCase.Server.Grains.csproj

@@ -10,6 +10,7 @@
   <ItemGroup>
     <PackageReference Include="Autofac" Version="4.6.1" />
     <PackageReference Include="Microsoft.Extensions.Logging" Version="2.0.0" />
+    <PackageReference Include="Microsoft.IO.RecyclableMemoryStream" Version="1.2.2" />
     <PackageReference Include="Microsoft.Orleans.Core" Version="2.0.0-preview2-20170724" />
     <PackageReference Include="StyleCop.Analyzers" Version="1.1.0-beta004" PrivateAssets="All" />
     <PackageReference Include="Newtonsoft.Json" Version="10.0.3" />

src/MineCase.Server.Grains/Network/ClientboundPacketSinkGrain.cs

@@ -52,7 +52,7 @@ public Task SendPacket(uint packetId, Immutable<byte[]> data)
             var packet = new UncompressedPacket
             {
                 PacketId = packetId,
-                Data = data.Value
+                Data = new ArraySegment<byte>(data.Value)
             };
             _subsManager.Notify(n => n.ReceivePacket(packet));
             return Task.CompletedTask;