Merge pull request #14113 from jtattermusch/fix_csharp_pooling_leak

Fix leak caused by C# context pooling (when GrpcEnvironment is shutdown).

src/csharp/Grpc.Core.Tests/Internal/DefaultObjectPoolTest.cs

@@ -59,6 +59,16 @@ namespace Grpc.Core.Internal.Tests
             Assert.AreNotSame(origLeased, pool.Lease());
         }
 
+        [Test]
+        public void LeaseSetsReturnAction()
+        {
+            var pool = new DefaultObjectPool<TestPooledObject>(() => new TestPooledObject(), 10, 0);
+            var origLeased = pool.Lease();
+            origLeased.ReturnAction(origLeased);
+            pool.Dispose();
+            Assert.AreNotSame(origLeased, pool.Lease());
+        }
+
         [Test]
         public void Constructor()
         {
@@ -67,8 +77,14 @@ namespace Grpc.Core.Internal.Tests
             Assert.Throws<ArgumentException>(() => new DefaultObjectPool<TestPooledObject>(() => new TestPooledObject(), 10, -1));
         }
 
-        class TestPooledObject : IDisposable
+        class TestPooledObject : IPooledObject<TestPooledObject>
         {
+            public Action<TestPooledObject> ReturnAction;
+
+            public void SetReturnToPoolAction(Action<TestPooledObject> returnAction)
+            {
+                this.ReturnAction = returnAction;
+            }
 
             public void Dispose()
             {

src/csharp/Grpc.Core/GrpcEnvironment.cs

@@ -299,8 +299,8 @@ namespace Grpc.Core
         private GrpcEnvironment()
         {
             GrpcNativeInit();
-            batchContextPool = new DefaultObjectPool<BatchContextSafeHandle>(() => BatchContextSafeHandle.Create(this.batchContextPool), batchContextPoolSharedCapacity, batchContextPoolThreadLocalCapacity);
-            requestCallContextPool = new DefaultObjectPool<RequestCallContextSafeHandle>(() => RequestCallContextSafeHandle.Create(this.requestCallContextPool), requestCallContextPoolSharedCapacity, requestCallContextPoolThreadLocalCapacity);
+            batchContextPool = new DefaultObjectPool<BatchContextSafeHandle>(() => BatchContextSafeHandle.Create(), batchContextPoolSharedCapacity, batchContextPoolThreadLocalCapacity);
+            requestCallContextPool = new DefaultObjectPool<RequestCallContextSafeHandle>(() => RequestCallContextSafeHandle.Create(), requestCallContextPoolSharedCapacity, requestCallContextPoolThreadLocalCapacity);
             threadPool = new GrpcThreadPool(this, GetThreadPoolSizeOrDefault(), GetCompletionQueueCountOrDefault(), inlineHandlers);
             threadPool.Start();
         }

src/csharp/Grpc.Core/Internal/BatchContextSafeHandle.cs

@@ -33,22 +33,21 @@ namespace Grpc.Core.Internal
     /// <summary>
     /// grpcsharp_batch_context
     /// </summary>
-    internal class BatchContextSafeHandle : SafeHandleZeroIsInvalid, IOpCompletionCallback
+    internal class BatchContextSafeHandle : SafeHandleZeroIsInvalid, IOpCompletionCallback, IPooledObject<BatchContextSafeHandle>
     {
         static readonly NativeMethods Native = NativeMethods.Get();
         static readonly ILogger Logger = GrpcEnvironment.Logger.ForType<BatchContextSafeHandle>();
 
-        IObjectPool<BatchContextSafeHandle> ownedByPool;
+        Action<BatchContextSafeHandle> returnToPoolAction;
         CompletionCallbackData completionCallbackData;
 
         private BatchContextSafeHandle()
         {
         }
 
-        public static BatchContextSafeHandle Create(IObjectPool<BatchContextSafeHandle> ownedByPool = null)
+        public static BatchContextSafeHandle Create()
         {
             var ctx = Native.grpcsharp_batch_context_create();
-            ctx.ownedByPool = ownedByPool;
             return ctx;
         }
 
@@ -60,6 +59,12 @@ namespace Grpc.Core.Internal
             }
         }
 
+        public void SetReturnToPoolAction(Action<BatchContextSafeHandle> returnAction)
+        {
+            GrpcPreconditions.CheckState(returnToPoolAction == null);
+            returnToPoolAction = returnAction;
+        }
+
         public void SetCompletionCallback(BatchCompletionDelegate callback, object state)
         {
             GrpcPreconditions.CheckState(completionCallbackData.Callback == null);
@@ -109,10 +114,15 @@ namespace Grpc.Core.Internal
 
         public void Recycle()
         {
-            if (ownedByPool != null)
+            if (returnToPoolAction != null)
             {
                 Native.grpcsharp_batch_context_reset(this);
-                ownedByPool.Return(this);
+
+                var origReturnAction = returnToPoolAction;
+                // Not clearing all the references to the pool could prevent garbage collection of the pool object
+                // and thus cause memory leaks.
+                returnToPoolAction = null;
+                origReturnAction(this);
             }
             else
             {

src/csharp/Grpc.Core/Internal/DefaultObjectPool.cs

@@ -27,9 +27,10 @@ namespace Grpc.Core.Internal
     /// Pool of objects that combines a shared pool and a thread local pool.
     /// </summary>
     internal class DefaultObjectPool<T> : IObjectPool<T>
-        where T : class, IDisposable
+        where T : class, IPooledObject<T>
     {
         readonly object myLock = new object();
+        readonly Action<T> returnAction;
         readonly Func<T> itemFactory;
 
         // Queue shared between threads, access needs to be synchronized.
@@ -54,6 +55,7 @@ namespace Grpc.Core.Internal
         {
             GrpcPreconditions.CheckArgument(sharedCapacity >= 0);
             GrpcPreconditions.CheckArgument(threadLocalCapacity >= 0);
+            this.returnAction = Return;
             this.itemFactory = GrpcPreconditions.CheckNotNull(itemFactory, nameof(itemFactory));
             this.sharedQueue = new Queue<T>(sharedCapacity);
             this.sharedCapacity = sharedCapacity;
@@ -73,6 +75,13 @@ namespace Grpc.Core.Internal
         /// in the thread local pool, it will continue returning new objects created by the factory).
         /// </summary>
         public T Lease()
+        {
+            var item = LeaseInternal();
+            item.SetReturnToPoolAction(returnAction);
+            return item;
+        }
+
+        private T LeaseInternal()
         {
             var localData = threadLocalData.Value;
             if (localData.Queue.Count > 0)

src/csharp/Grpc.Core/Internal/IPooledObject.cs

@@ -0,0 +1,34 @@
+#region Copyright notice and license
+
+// Copyright 2018 gRPC authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#endregion
+
+using System;
+
+namespace Grpc.Core.Internal
+{
+    /// <summary>
+    /// An object that can be pooled in <c>IObjectPool</c>.
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    internal interface IPooledObject<T> : IDisposable
+    {
+        /// <summary>
+        /// Set the action that will be invoked to return a leased object to the pool.
+        /// </summary>
+        void SetReturnToPoolAction(Action<T> returnAction);
+    }
+}

src/csharp/Grpc.Core/Internal/RequestCallContextSafeHandle.cs

@@ -20,26 +20,26 @@ using System;
 using System.Runtime.InteropServices;
 using Grpc.Core;
 using Grpc.Core.Logging;
+using Grpc.Core.Utils;
 
 namespace Grpc.Core.Internal
 {
     /// <summary>
     /// grpcsharp_request_call_context
     /// </summary>
-    internal class RequestCallContextSafeHandle : SafeHandleZeroIsInvalid, IOpCompletionCallback
+    internal class RequestCallContextSafeHandle : SafeHandleZeroIsInvalid, IOpCompletionCallback, IPooledObject<RequestCallContextSafeHandle>
     {
         static readonly NativeMethods Native = NativeMethods.Get();
         static readonly ILogger Logger = GrpcEnvironment.Logger.ForType<RequestCallContextSafeHandle>();
-        IObjectPool<RequestCallContextSafeHandle> ownedByPool;
+        Action<RequestCallContextSafeHandle> returnToPoolAction;
 
         private RequestCallContextSafeHandle()
         {
         }
 
-        public static RequestCallContextSafeHandle Create(IObjectPool<RequestCallContextSafeHandle> ownedByPool = null)
+        public static RequestCallContextSafeHandle Create()
         {
             var ctx = Native.grpcsharp_request_call_context_create();
-            ctx.ownedByPool = ownedByPool;
             return ctx;
         }
 
@@ -51,6 +51,12 @@ namespace Grpc.Core.Internal
             }
         }
 
+        public void SetReturnToPoolAction(Action<RequestCallContextSafeHandle> returnAction)
+        {
+            GrpcPreconditions.CheckState(returnToPoolAction == null);
+            returnToPoolAction = returnAction;
+        }
+
         public RequestCallCompletionDelegate CompletionCallback { get; set; }
 
         // Gets data of server_rpc_new completion.
@@ -76,10 +82,15 @@ namespace Grpc.Core.Internal
 
         public void Recycle()
         {
-            if (ownedByPool != null)
+            if (returnToPoolAction != null)
             {
                 Native.grpcsharp_request_call_context_reset(this);
-                ownedByPool.Return(this);
+
+                var origReturnAction = returnToPoolAction;
+                // Not clearing all the references to the pool could prevent garbage collection of the pool object
+                // and thus cause memory leaks.
+                returnToPoolAction = null;
+                origReturnAction(this);
             }
             else
             {