Fix ArrayPool leak with JsonDocument (dotnet#59540)

macsux · Sep 24, 2021 · 2e08554 · 2e08554
1 parent 7233b56
commit 2e08554
Showing 1 changed file with 11 additions and 23 deletions.
diff --git a/src/libraries/System.Text.Json/src/System/Text/Json/Document/JsonDocument.cs b/src/libraries/System.Text.Json/src/System/Text/Json/Document/JsonDocument.cs
@@ -25,10 +25,7 @@ public sealed partial class JsonDocument : IDisposable
         private MetadataDb _parsedData;
 
         private byte[]? _extraRentedArrayPoolBytes;
-        private bool _hasExtraRentedArrayPoolBytes;
-
         private PooledByteBufferWriter? _extraPooledByteBufferWriter;
-        private bool _hasExtraPooledByteBufferWriter;
 
         private (int, string?) _lastIndexAndString = (-1, null);
 
@@ -48,27 +45,18 @@ private JsonDocument(
         {
             Debug.Assert(!utf8Json.IsEmpty);
 
-            // We never have both rented fields.
+            // Both rented values better be null if we're not disposable.
+            Debug.Assert(isDisposable ||
+                (extraRentedArrayPoolBytes == null && extraPooledByteBufferWriter == null));
+
+            // Both rented values can't be specified.
             Debug.Assert(extraRentedArrayPoolBytes == null || extraPooledByteBufferWriter == null);
 
             _utf8Json = utf8Json;
             _parsedData = parsedData;
-
-            if (_extraRentedArrayPoolBytes != null)
-            {
-                _hasExtraRentedArrayPoolBytes = true;
-                _extraRentedArrayPoolBytes = extraRentedArrayPoolBytes;
-            }
-            else if (extraPooledByteBufferWriter != null)
-            {
-                _hasExtraPooledByteBufferWriter = true;
-                _extraPooledByteBufferWriter = extraPooledByteBufferWriter;
-            }
-
+            _extraRentedArrayPoolBytes = extraRentedArrayPoolBytes;
+            _extraPooledByteBufferWriter = extraPooledByteBufferWriter;
             IsDisposable = isDisposable;
-
-            // Both rented fields better be null if we're not disposable.
-            Debug.Assert(isDisposable || (_extraRentedArrayPoolBytes == null && _extraPooledByteBufferWriter == null));
         }
 
         /// <inheritdoc />
@@ -83,9 +71,9 @@ public void Dispose()
             _parsedData.Dispose();
             _utf8Json = ReadOnlyMemory<byte>.Empty;
 
-            if (_hasExtraRentedArrayPoolBytes)
+            if (_extraRentedArrayPoolBytes != null)
             {
-                byte[]? extraRentedBytes = Interlocked.Exchange(ref _extraRentedArrayPoolBytes, null);
+                byte[]? extraRentedBytes = Interlocked.Exchange<byte[]?>(ref _extraRentedArrayPoolBytes, null);
 
                 if (extraRentedBytes != null)
                 {
@@ -95,9 +83,9 @@ public void Dispose()
                     ArrayPool<byte>.Shared.Return(extraRentedBytes);
                 }
             }
-            else if (_hasExtraPooledByteBufferWriter)
+            else if (_extraPooledByteBufferWriter != null)
             {
-                PooledByteBufferWriter? extraBufferWriter = Interlocked.Exchange(ref _extraPooledByteBufferWriter, null);
+                PooledByteBufferWriter? extraBufferWriter = Interlocked.Exchange<PooledByteBufferWriter?>(ref _extraPooledByteBufferWriter, null);
                 extraBufferWriter?.Dispose();
             }
         }