Merge pull request #797 from Cysharp/feature/FixIncorrectUseOfPayload…

…InHeartbeatManager Fix incorrect use of Payload within HeartbeatManager
Cysharp · Jun 28, 2024 · 133d110 · 133d110
2 parents c9d52f5 + e1ce308
commit 133d110
Show file tree

Hide file tree

Showing 8 changed files with 205 additions and 52 deletions.
diff --git a/....Unity/Assets/Scripts/MagicOnion/MagicOnion.Client/Internal.Shared/StreamingHubPayload.cs b/....Unity/Assets/Scripts/MagicOnion/MagicOnion.Client/Internal.Shared/StreamingHubPayload.cs
@@ -1,20 +1,52 @@
 #nullable enable
 using System;
 using System.Buffers;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 
 namespace MagicOnion.Internal
 {
-    internal class StreamingHubPayload : IStreamingHubPayload
+    internal class StreamingHubPayload
     {
         byte[]? buffer;
         ReadOnlyMemory<byte>? memory;
 
+#if DEBUG
+        public int Length
+        {
+            get
+            {
+                ThrowIfUninitialized();
+                return memory!.Value.Length;
+            }
+        }
+
+        public ReadOnlySpan<byte> Span
+        {
+            get
+            {
+                ThrowIfUninitialized();
+                return memory!.Value.Span;
+            }
+        }
+
+        public ReadOnlyMemory<byte> Memory
+        {
+            get
+            {
+                ThrowIfUninitialized();
+                return memory!.Value;
+            }
+        }
+
+#else
         public int Length => memory!.Value.Length;
         public ReadOnlySpan<byte> Span => memory!.Value.Span;
         public ReadOnlyMemory<byte> Memory => memory!.Value;
+#endif
 
-        void IStreamingHubPayload.Initialize(ReadOnlySpan<byte> data)
+        public void Initialize(ReadOnlySpan<byte> data)
         {
             ThrowIfUsing();
 
@@ -23,7 +55,7 @@ void IStreamingHubPayload.Initialize(ReadOnlySpan<byte> data)
             memory = buffer.AsMemory(0, (int)data.Length);
         }
 
-        void IStreamingHubPayload.Initialize(ReadOnlySequence<byte> data)
+        public void Initialize(ReadOnlySequence<byte> data)
         {
             ThrowIfUsing();
             if (data.Length > int.MaxValue) throw new InvalidOperationException("A body size of StreamingHubPayload must be less than int.MaxValue");
@@ -33,17 +65,17 @@ void IStreamingHubPayload.Initialize(ReadOnlySequence<byte> data)
             memory = buffer.AsMemory(0, (int)data.Length);
         }
 
-        void IStreamingHubPayload.Initialize(ReadOnlyMemory<byte> data)
+        public void Initialize(ReadOnlyMemory<byte> data)
         {
             ThrowIfUsing();
 
             buffer = null;
             memory = data;
         }
 
-        void IStreamingHubPayload.Uninitialize()
+        public void Uninitialize()
         {
-            ThrowIfDisposed();
+            ThrowIfUninitialized();
 
             if (buffer != null)
             {
@@ -60,22 +92,22 @@ void IStreamingHubPayload.Uninitialize()
 #if NON_UNITY && !NETSTANDARD2_0 && !NETSTANDARD2_1
         [MemberNotNull(nameof(memory))]
 #endif
-        void ThrowIfDisposed()
+        void ThrowIfUninitialized()
         {
-            if (memory is null) throw new ObjectDisposedException(nameof(StreamingHubPayload));
+            //Debug.Assert(memory is not null);
+            if (memory is null)
+            {
+                throw new InvalidOperationException("A StreamingHubPayload has been already uninitialized.");
+            }
         }
 
         void ThrowIfUsing()
         {
-            if (memory is not null) throw new InvalidOperationException(nameof(StreamingHubPayload));
+            //Debug.Assert(memory is null);
+            if (memory is not null)
+            {
+                throw new InvalidOperationException("A StreamingHubPayload is currently used by other caller.");
+            }
         }
     }
-
-    internal interface IStreamingHubPayload
-    {
-        void Initialize(ReadOnlySpan<byte> data);
-        void Initialize(ReadOnlySequence<byte> data);
-        void Initialize(ReadOnlyMemory<byte> data);
-        void Uninitialize();
-    }
 }
diff --git a/...s/Scripts/MagicOnion/MagicOnion.Client/Internal.Shared/StreamingHubPayloadPool.BuiltIn.cs b/...s/Scripts/MagicOnion/MagicOnion.Client/Internal.Shared/StreamingHubPayloadPool.BuiltIn.cs
@@ -61,28 +61,28 @@ internal class StreamingHubPayloadPool : ObjectPool<StreamingHubPayload>
 
         public void Return(StreamingHubPayload payload)
         {
-            ((IStreamingHubPayload)payload).Uninitialize();
+            payload.Uninitialize();
             ReturnCore(payload);
         }
 
         public StreamingHubPayload RentOrCreate(ReadOnlySequence<byte> data)
         {
             var payload = RentOrCreateCore(static () => new StreamingHubPayload());
-            ((IStreamingHubPayload)payload).Initialize(data);
+            payload.Initialize(data);
             return payload;
         }
 
         public StreamingHubPayload RentOrCreate(ReadOnlySpan<byte> data)
         {
             var payload = RentOrCreateCore(static () => new StreamingHubPayload());
-            ((IStreamingHubPayload)payload).Initialize(data);
+            payload.Initialize(data);
             return payload;
         }
 
         public StreamingHubPayload RentOrCreate(ReadOnlyMemory<byte> data)
         {
             var payload = RentOrCreateCore(static () => new StreamingHubPayload());
-            ((IStreamingHubPayload)payload).Initialize(data);
+            payload.Initialize(data);
             return payload;
         }
     }

diff --git a/...cripts/MagicOnion/MagicOnion.Client/Internal.Shared/StreamingHubPayloadPool.ObjectPool.cs b/...cripts/MagicOnion/MagicOnion.Client/Internal.Shared/StreamingHubPayloadPool.ObjectPool.cs
@@ -15,21 +15,21 @@ internal class StreamingHubPayloadPool
     public StreamingHubPayload RentOrCreate(ReadOnlySequence<byte> data)
     {
         var payload = pool.Get();
-        ((IStreamingHubPayload)payload).Initialize(data);
+        payload.Initialize(data);
         return payload;
     }
 
     public StreamingHubPayload RentOrCreate(ReadOnlySpan<byte> data)
     {
         var payload = pool.Get();
-        ((IStreamingHubPayload)payload).Initialize(data);
+        payload.Initialize(data);
         return payload;
     }
 
     public StreamingHubPayload RentOrCreate(ReadOnlyMemory<byte> data)
     {
         var payload = pool.Get();
-        ((IStreamingHubPayload)payload).Initialize(data);
+        payload.Initialize(data);
         return payload;
     }
 
@@ -47,7 +47,7 @@ public StreamingHubPayload Create()
 
         public bool Return(StreamingHubPayload obj)
         {
-            ((IStreamingHubPayload)obj).Uninitialize();
+            obj.Uninitialize();
             return true;
         }
     }

diff --git a/src/MagicOnion.Internal/StreamingHubPayload.cs b/src/MagicOnion.Internal/StreamingHubPayload.cs
@@ -1,20 +1,52 @@
 #nullable enable
 using System;
 using System.Buffers;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 
 namespace MagicOnion.Internal
 {
-    internal class StreamingHubPayload : IStreamingHubPayload
+    internal class StreamingHubPayload
     {
         byte[]? buffer;
         ReadOnlyMemory<byte>? memory;
 
+#if DEBUG
+        public int Length
+        {
+            get
+            {
+                ThrowIfUninitialized();
+                return memory!.Value.Length;
+            }
+        }
+
+        public ReadOnlySpan<byte> Span
+        {
+            get
+            {
+                ThrowIfUninitialized();
+                return memory!.Value.Span;
+            }
+        }
+
+        public ReadOnlyMemory<byte> Memory
+        {
+            get
+            {
+                ThrowIfUninitialized();
+                return memory!.Value;
+            }
+        }
+
+#else
         public int Length => memory!.Value.Length;
         public ReadOnlySpan<byte> Span => memory!.Value.Span;
         public ReadOnlyMemory<byte> Memory => memory!.Value;
+#endif
 
-        void IStreamingHubPayload.Initialize(ReadOnlySpan<byte> data)
+        public void Initialize(ReadOnlySpan<byte> data)
         {
             ThrowIfUsing();
 
@@ -23,7 +55,7 @@ void IStreamingHubPayload.Initialize(ReadOnlySpan<byte> data)
             memory = buffer.AsMemory(0, (int)data.Length);
         }
 
-        void IStreamingHubPayload.Initialize(ReadOnlySequence<byte> data)
+        public void Initialize(ReadOnlySequence<byte> data)
         {
             ThrowIfUsing();
             if (data.Length > int.MaxValue) throw new InvalidOperationException("A body size of StreamingHubPayload must be less than int.MaxValue");
@@ -33,17 +65,17 @@ void IStreamingHubPayload.Initialize(ReadOnlySequence<byte> data)
             memory = buffer.AsMemory(0, (int)data.Length);
         }
 
-        void IStreamingHubPayload.Initialize(ReadOnlyMemory<byte> data)
+        public void Initialize(ReadOnlyMemory<byte> data)
         {
             ThrowIfUsing();
 
             buffer = null;
             memory = data;
         }
 
-        void IStreamingHubPayload.Uninitialize()
+        public void Uninitialize()
         {
-            ThrowIfDisposed();
+            ThrowIfUninitialized();
 
             if (buffer != null)
             {
@@ -60,22 +92,22 @@ void IStreamingHubPayload.Uninitialize()
 #if NON_UNITY && !NETSTANDARD2_0 && !NETSTANDARD2_1
         [MemberNotNull(nameof(memory))]
 #endif
-        void ThrowIfDisposed()
+        void ThrowIfUninitialized()
         {
-            if (memory is null) throw new ObjectDisposedException(nameof(StreamingHubPayload));
+            //Debug.Assert(memory is not null);
+            if (memory is null)
+            {
+                throw new InvalidOperationException("A StreamingHubPayload has been already uninitialized.");
+            }
         }
 
         void ThrowIfUsing()
         {
-            if (memory is not null) throw new InvalidOperationException(nameof(StreamingHubPayload));
+            //Debug.Assert(memory is null);
+            if (memory is not null)
+            {
+                throw new InvalidOperationException("A StreamingHubPayload is currently used by other caller.");
+            }
         }
     }
-
-    internal interface IStreamingHubPayload
-    {
-        void Initialize(ReadOnlySpan<byte> data);
-        void Initialize(ReadOnlySequence<byte> data);
-        void Initialize(ReadOnlyMemory<byte> data);
-        void Uninitialize();
-    }
 }
diff --git a/src/MagicOnion.Internal/StreamingHubPayloadPool.BuiltIn.cs b/src/MagicOnion.Internal/StreamingHubPayloadPool.BuiltIn.cs
@@ -61,28 +61,28 @@ internal class StreamingHubPayloadPool : ObjectPool<StreamingHubPayload>
 
         public void Return(StreamingHubPayload payload)
         {
-            ((IStreamingHubPayload)payload).Uninitialize();
+            payload.Uninitialize();
             ReturnCore(payload);
         }
 
         public StreamingHubPayload RentOrCreate(ReadOnlySequence<byte> data)
         {
             var payload = RentOrCreateCore(static () => new StreamingHubPayload());
-            ((IStreamingHubPayload)payload).Initialize(data);
+            payload.Initialize(data);
             return payload;
         }
 
         public StreamingHubPayload RentOrCreate(ReadOnlySpan<byte> data)
         {
             var payload = RentOrCreateCore(static () => new StreamingHubPayload());
-            ((IStreamingHubPayload)payload).Initialize(data);
+            payload.Initialize(data);
             return payload;
         }
 
         public StreamingHubPayload RentOrCreate(ReadOnlyMemory<byte> data)
         {
             var payload = RentOrCreateCore(static () => new StreamingHubPayload());
-            ((IStreamingHubPayload)payload).Initialize(data);
+            payload.Initialize(data);
             return payload;
         }
     }

diff --git a/src/MagicOnion.Internal/StreamingHubPayloadPool.ObjectPool.cs b/src/MagicOnion.Internal/StreamingHubPayloadPool.ObjectPool.cs
@@ -15,21 +15,21 @@ internal class StreamingHubPayloadPool
     public StreamingHubPayload RentOrCreate(ReadOnlySequence<byte> data)
     {
         var payload = pool.Get();
-        ((IStreamingHubPayload)payload).Initialize(data);
+        payload.Initialize(data);
         return payload;
     }
 
     public StreamingHubPayload RentOrCreate(ReadOnlySpan<byte> data)
     {
         var payload = pool.Get();
-        ((IStreamingHubPayload)payload).Initialize(data);
+        payload.Initialize(data);
         return payload;
     }
 
     public StreamingHubPayload RentOrCreate(ReadOnlyMemory<byte> data)
     {
         var payload = pool.Get();
-        ((IStreamingHubPayload)payload).Initialize(data);
+        payload.Initialize(data);
         return payload;
     }
 
@@ -47,7 +47,7 @@ public StreamingHubPayload Create()
 
         public bool Return(StreamingHubPayload obj)
         {
-            ((IStreamingHubPayload)obj).Uninitialize();
+            obj.Uninitialize();
             return true;
         }
     }

diff --git a/src/MagicOnion.Server/Hubs/StreamingHubHeartbeatManager.cs b/src/MagicOnion.Server/Hubs/StreamingHubHeartbeatManager.cs
@@ -125,15 +125,14 @@ async Task StartHeartbeatAsync(PeriodicTimer runningTimer, string method)
                 writer.Write(Nil);
             }
 
-            var payload = StreamingHubPayloadPool.Shared.RentOrCreate(writer.WrittenSpan);
 
             MagicOnionServerLog.SendHeartbeat(this.logger, method);
             try
             {
                 foreach (var (contextId, handle) in contexts)
                 {
                     handle.RestartTimeoutTimer();
-                    handle.ServiceContext.QueueResponseStreamWrite(payload);
+                    handle.ServiceContext.QueueResponseStreamWrite(StreamingHubPayloadPool.Shared.RentOrCreate(writer.WrittenSpan));
                 }
             }
             catch { /* Ignore */ }