From 3fe56b52de02079eea98590adcc852dc9e3aaffd Mon Sep 17 00:00:00 2001
From: Craig Tiller <ctiller@google.com>
Date: Thu, 22 Sep 2016 17:04:53 -0700
Subject: [PATCH] Add a test of reclaimation

---
 src/core/lib/iomgr/buffer_pool.c   | 15 ++++++++
 test/core/iomgr/buffer_pool_test.c | 62 ++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/src/core/lib/iomgr/buffer_pool.c b/src/core/lib/iomgr/buffer_pool.c
index 26578e83cc8..a6824f94322 100644
--- a/src/core/lib/iomgr/buffer_pool.c
+++ b/src/core/lib/iomgr/buffer_pool.c
@@ -56,6 +56,7 @@ struct grpc_buffer_pool {
   bool step_scheduled;
   bool reclaiming;
   grpc_closure bpstep_closure;
+  grpc_closure bpreclaimation_done_closure;
 
   grpc_buffer_user *roots[GRPC_BULIST_COUNT];
 };
@@ -305,6 +306,11 @@ static void bp_resize(grpc_exec_ctx *exec_ctx, void *args, grpc_error *error) {
   gpr_free(a);
 }
 
+static void bpreclaimation_done_closure(grpc_exec_ctx *exec_ctx, void *bp,
+                                        grpc_error *error) {
+  bpstep_sched(exec_ctx, bp);
+}
+
 /*******************************************************************************
  * grpc_buffer_pool api
  */
@@ -316,6 +322,8 @@ grpc_buffer_pool *grpc_buffer_pool_create(void) {
   buffer_pool->free_pool = INT64_MAX;
   buffer_pool->size = INT64_MAX;
   grpc_closure_init(&buffer_pool->bpstep_closure, bpstep, buffer_pool);
+  grpc_closure_init(&buffer_pool->bpreclaimation_done_closure,
+                    bpreclaimation_done_closure, buffer_pool);
   for (int i = 0; i < GRPC_BULIST_COUNT; i++) {
     buffer_pool->roots[i] = NULL;
   }
@@ -440,3 +448,10 @@ void grpc_buffer_user_post_reclaimer(grpc_exec_ctx *exec_ctx,
                         &buffer_user->post_reclaimer_closure[destructive],
                         GRPC_ERROR_NONE, false);
 }
+
+void grpc_buffer_user_finish_reclaimation(grpc_exec_ctx *exec_ctx,
+                                          grpc_buffer_user *buffer_user) {
+  grpc_combiner_execute(exec_ctx, buffer_user->buffer_pool->combiner,
+                        &buffer_user->buffer_pool->bpreclaimation_done_closure,
+                        GRPC_ERROR_NONE, false);
+}
diff --git a/test/core/iomgr/buffer_pool_test.c b/test/core/iomgr/buffer_pool_test.c
index f659c5c5b4e..ba44af51444 100644
--- a/test/core/iomgr/buffer_pool_test.c
+++ b/test/core/iomgr/buffer_pool_test.c
@@ -33,6 +33,7 @@
 
 #include "src/core/lib/iomgr/buffer_pool.h"
 
+#include <grpc/support/alloc.h>
 #include <grpc/support/log.h>
 
 #include "test/core/util/test_config.h"
@@ -42,6 +43,28 @@ static void set_bool_cb(grpc_exec_ctx *exec_ctx, void *a, grpc_error *error) {
 }
 grpc_closure *set_bool(bool *p) { return grpc_closure_create(set_bool_cb, p); }
 
+typedef struct {
+  size_t size;
+  grpc_buffer_user *buffer_user;
+  grpc_closure *then;
+} reclaimer_args;
+static void reclaimer_cb(grpc_exec_ctx *exec_ctx, void *args,
+                         grpc_error *error) {
+  reclaimer_args *a = args;
+  grpc_buffer_user_free(exec_ctx, a->buffer_user, a->size);
+  grpc_buffer_user_finish_reclaimation(exec_ctx, a->buffer_user);
+  grpc_closure_run(exec_ctx, a->then, GRPC_ERROR_NONE);
+  gpr_free(a);
+}
+grpc_closure *make_reclaimer(grpc_buffer_user *buffer_user, size_t size,
+                             grpc_closure *then) {
+  reclaimer_args *a = gpr_malloc(sizeof(*a));
+  a->size = size;
+  a->buffer_user = buffer_user;
+  a->then = then;
+  return grpc_closure_create(reclaimer_cb, a);
+}
+
 static void destroy_user(grpc_buffer_user *usr) {
   grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
   bool done = false;
@@ -229,6 +252,44 @@ static void test_scavenge_blocked(void) {
   destroy_user(&usr2);
 }
 
+static void test_blocked_until_scheduled_reclaim(void) {
+  gpr_log(GPR_INFO, "** test_blocked_until_scheduled_reclaim **");
+  grpc_buffer_pool *p = grpc_buffer_pool_create();
+  grpc_buffer_pool_resize(p, 1024);
+  grpc_buffer_user usr;
+  grpc_buffer_user_init(&usr, p);
+  {
+    bool done = false;
+    grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+    grpc_buffer_user_alloc(&exec_ctx, &usr, 1024, set_bool(&done));
+    grpc_exec_ctx_finish(&exec_ctx);
+    GPR_ASSERT(done);
+  }
+  bool reclaim_done = false;
+  {
+    grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+    grpc_buffer_user_post_reclaimer(
+        &exec_ctx, &usr, false,
+        make_reclaimer(&usr, 1024, set_bool(&reclaim_done)));
+    grpc_exec_ctx_finish(&exec_ctx);
+  }
+  {
+    bool done = false;
+    grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+    grpc_buffer_user_alloc(&exec_ctx, &usr, 1024, set_bool(&done));
+    grpc_exec_ctx_finish(&exec_ctx);
+    GPR_ASSERT(reclaim_done);
+    GPR_ASSERT(done);
+  }
+  {
+    grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+    grpc_buffer_user_free(&exec_ctx, &usr, 1024);
+    grpc_exec_ctx_finish(&exec_ctx);
+  }
+  grpc_buffer_pool_unref(p);
+  destroy_user(&usr);
+}
+
 int main(int argc, char **argv) {
   grpc_test_init(argc, argv);
   grpc_init();
@@ -241,6 +302,7 @@ int main(int argc, char **argv) {
   test_async_alloc_blocked_by_size();
   test_scavenge();
   test_scavenge_blocked();
+  test_blocked_until_scheduled_reclaim();
   grpc_shutdown();
   return 0;
 }