From 579acc351b3329da7d23b05cc87d979551c08689 Mon Sep 17 00:00:00 2001
From: Craig Tiller <ctiller@google.com>
Date: Mon, 26 Sep 2016 13:31:26 -0700
Subject: [PATCH] Ensure all memory is released back to the buffer pool

---
 src/core/lib/iomgr/buffer_pool.c   |  6 ++++
 test/core/iomgr/buffer_pool_test.c | 46 ++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/src/core/lib/iomgr/buffer_pool.c b/src/core/lib/iomgr/buffer_pool.c
index 16445a42653..699e7078db6 100644
--- a/src/core/lib/iomgr/buffer_pool.c
+++ b/src/core/lib/iomgr/buffer_pool.c
@@ -153,12 +153,14 @@ static void bpstep(grpc_exec_ctx *exec_ctx, void *bp, grpc_error *error) {
   } while (bpscavenge(exec_ctx, buffer_pool));
   bpreclaim(exec_ctx, buffer_pool, false) ||
       bpreclaim(exec_ctx, buffer_pool, true);
+  grpc_buffer_pool_internal_unref(exec_ctx, buffer_pool);
 }
 
 static void bpstep_sched(grpc_exec_ctx *exec_ctx,
                          grpc_buffer_pool *buffer_pool) {
   if (buffer_pool->step_scheduled) return;
   buffer_pool->step_scheduled = true;
+  grpc_buffer_pool_internal_ref(buffer_pool);
   grpc_combiner_execute_finally(exec_ctx, buffer_pool->combiner,
                                 &buffer_pool->bpstep_closure, GRPC_ERROR_NONE,
                                 false);
@@ -329,6 +331,10 @@ static void bu_destroy(grpc_exec_ctx *exec_ctx, void *bu, grpc_error *error) {
                       GRPC_ERROR_CANCELLED, NULL);
   grpc_exec_ctx_sched(exec_ctx, buffer_user->on_done_destroy, GRPC_ERROR_NONE,
                       NULL);
+  if (buffer_user->free_pool != 0) {
+    buffer_user->buffer_pool->free_pool += buffer_user->free_pool;
+    bpstep_sched(exec_ctx, buffer_user->buffer_pool);
+  }
   grpc_buffer_pool_internal_unref(exec_ctx, buffer_user->buffer_pool);
 }
 
diff --git a/test/core/iomgr/buffer_pool_test.c b/test/core/iomgr/buffer_pool_test.c
index 265ece72acd..9d30781c312 100644
--- a/test/core/iomgr/buffer_pool_test.c
+++ b/test/core/iomgr/buffer_pool_test.c
@@ -524,6 +524,51 @@ static void test_buffer_user_stays_allocated_until_memory_released(void) {
     grpc_exec_ctx_finish(&exec_ctx);
     GPR_ASSERT(done);
   }
+  grpc_buffer_pool_unref(p);
+}
+
+static void test_pools_merged_on_buffer_user_deletion(void) {
+  gpr_log(GPR_INFO, "** test_pools_merged_on_buffer_user_deletion **");
+  grpc_buffer_pool *p = grpc_buffer_pool_create();
+  grpc_buffer_pool_resize(p, 1024);
+  for (int i = 0; i < 10; i++) {
+    grpc_buffer_user usr;
+    grpc_buffer_user_init(&usr, p);
+    bool done = false;
+    bool reclaimer_cancelled = false;
+    {
+      grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+      grpc_buffer_user_post_reclaimer(
+          &exec_ctx, &usr, false,
+          make_unused_reclaimer(set_bool(&reclaimer_cancelled)));
+      grpc_exec_ctx_finish(&exec_ctx);
+      GPR_ASSERT(!reclaimer_cancelled);
+    }
+    {
+      bool allocated = false;
+      grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+      grpc_buffer_user_alloc(&exec_ctx, &usr, 1024, set_bool(&allocated));
+      grpc_exec_ctx_finish(&exec_ctx);
+      GPR_ASSERT(allocated);
+      GPR_ASSERT(!reclaimer_cancelled);
+    }
+    {
+      grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+      grpc_buffer_pool_unref(p);
+      grpc_buffer_user_shutdown(&exec_ctx, &usr, set_bool(&done));
+      grpc_exec_ctx_finish(&exec_ctx);
+      GPR_ASSERT(!done);
+      GPR_ASSERT(!reclaimer_cancelled);
+    }
+    {
+      grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+      grpc_buffer_user_free(&exec_ctx, &usr, 1024);
+      grpc_exec_ctx_finish(&exec_ctx);
+      GPR_ASSERT(done);
+      GPR_ASSERT(reclaimer_cancelled);
+    }
+  }
+  grpc_buffer_pool_unref(p);
 }
 
 int main(int argc, char **argv) {
@@ -545,6 +590,7 @@ int main(int argc, char **argv) {
   test_benign_reclaim_is_preferred();
   test_multiple_reclaims_can_be_triggered();
   test_buffer_user_stays_allocated_until_memory_released();
+  test_pools_merged_on_buffer_user_deletion();
   grpc_shutdown();
   return 0;
 }