3 files changed, 37 insertions, 5 deletions

diff --git a/test/core/transport/bdp_estimator_test.c b/test/core/transport/bdp_estimator_test.c
index f55a3ca643..122e097cc4 100644
--- a/test/core/transport/bdp_estimator_test.c
+++ b/test/core/transport/bdp_estimator_test.c
@@ -33,6 +33,7 @@
 
 #include "src/core/lib/transport/bdp_estimator.h"
 
+#include <grpc/grpc.h>
 #include <grpc/support/alloc.h>
 #include <grpc/support/log.h>
 #include <grpc/support/string_util.h>
@@ -64,6 +65,8 @@ static void add_samples(grpc_bdp_estimator *estimator, int64_t *samples,
     GPR_ASSERT(grpc_bdp_estimator_add_incoming_bytes(estimator, samples[i]) ==
                false);
   }
+  gpr_sleep_until(gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
+                               gpr_time_from_millis(1, GPR_TIMESPAN)));
   grpc_bdp_estimator_complete_ping(estimator);
 }
 
@@ -123,11 +126,11 @@ static void test_get_estimate_random_values(size_t n) {
   gpr_log(GPR_INFO, "test_get_estimate_random_values(%" PRIdPTR ")", n);
   grpc_bdp_estimator est;
   grpc_bdp_estimator_init(&est, "test");
-  int min = INT_MAX;
-  int max = 65535;  // Windows rand() has limited range, make sure the ASSERT
-                    // passes
+  const int kMaxSample = 65535;
+  int min = kMaxSample;
+  int max = 0;
   for (size_t i = 0; i < n; i++) {
-    int sample = rand();
+    int sample = rand() % (kMaxSample + 1);
     if (sample < min) min = sample;
     if (sample > max) max = sample;
     add_sample(&est, sample);
@@ -141,6 +144,7 @@ static void test_get_estimate_random_values(size_t n) {
 
 int main(int argc, char **argv) {
   grpc_test_init(argc, argv);
+  grpc_init();
   test_noop();
   test_get_estimate_no_samples();
   test_get_estimate_1_sample();
@@ -149,5 +153,6 @@ int main(int argc, char **argv) {
   for (size_t i = 3; i < 1000; i = i * 3 / 2) {
     test_get_estimate_random_values(i);
   }
+  grpc_shutdown();
   return 0;
 }
diff --git a/test/core/util/trickle_endpoint.c b/test/core/util/trickle_endpoint.c
index 0848147158..02ba257abe 100644
--- a/test/core/util/trickle_endpoint.c
+++ b/test/core/util/trickle_endpoint.c
@@ -44,6 +44,8 @@
 #include <grpc/support/useful.h>
 #include "src/core/lib/slice/slice_internal.h"
 
+#define WRITE_BUFFER_SIZE (2 * 1024 * 1024)
+
 typedef struct {
   grpc_endpoint base;
   double bytes_per_second;
@@ -55,6 +57,7 @@ typedef struct {
   grpc_slice_buffer writing_buffer;
   grpc_error *error;
   bool writing;
+  grpc_closure *write_cb;
 } trickle_endpoint;
 
 static void te_read(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
@@ -63,6 +66,15 @@ static void te_read(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
   grpc_endpoint_read(exec_ctx, te->wrapped, slices, cb);
 }
 
+static void maybe_call_write_cb_locked(grpc_exec_ctx *exec_ctx,
+                                       trickle_endpoint *te) {
+  if (te->write_cb != NULL && (te->error != GRPC_ERROR_NONE ||
+                               te->write_buffer.length <= WRITE_BUFFER_SIZE)) {
+    grpc_closure_sched(exec_ctx, te->write_cb, GRPC_ERROR_REF(te->error));
+    te->write_cb = NULL;
+  }
+}
+
 static void te_write(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
                      grpc_slice_buffer *slices, grpc_closure *cb) {
   trickle_endpoint *te = (trickle_endpoint *)ep;
@@ -70,11 +82,13 @@ static void te_write(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
     grpc_slice_ref_internal(slices->slices[i]);
   }
   gpr_mu_lock(&te->mu);
+  GPR_ASSERT(te->write_cb == NULL);
   if (te->write_buffer.length == 0) {
     te->last_write = gpr_now(GPR_CLOCK_MONOTONIC);
   }
   grpc_slice_buffer_addn(&te->write_buffer, slices->slices, slices->count);
-  grpc_closure_sched(exec_ctx, cb, GRPC_ERROR_REF(te->error));
+  te->write_cb = cb;
+  maybe_call_write_cb_locked(exec_ctx, te);
   gpr_mu_unlock(&te->mu);
 }
 
@@ -102,6 +116,7 @@ static void te_shutdown(grpc_exec_ctx *exec_ctx, grpc_endpoint *ep,
   if (te->error == GRPC_ERROR_NONE) {
     te->error = GRPC_ERROR_REF(why);
   }
+  maybe_call_write_cb_locked(exec_ctx, te);
   gpr_mu_unlock(&te->mu);
   grpc_endpoint_shutdown(exec_ctx, te->wrapped, why);
 }
@@ -157,6 +172,7 @@ grpc_endpoint *grpc_trickle_endpoint_create(grpc_endpoint *wrap,
   te->base.vtable = &vtable;
   te->wrapped = wrap;
   te->bytes_per_second = bytes_per_second;
+  te->write_cb = NULL;
   gpr_mu_init(&te->mu);
   grpc_slice_buffer_init(&te->write_buffer);
   grpc_slice_buffer_init(&te->writing_buffer);
@@ -187,9 +203,18 @@ size_t grpc_trickle_endpoint_trickle(grpc_exec_ctx *exec_ctx,
       grpc_endpoint_write(
           exec_ctx, te->wrapped, &te->writing_buffer,
           grpc_closure_create(te_finish_write, te, grpc_schedule_on_exec_ctx));
+      maybe_call_write_cb_locked(exec_ctx, te);
     }
   }
   size_t backlog = te->write_buffer.length;
   gpr_mu_unlock(&te->mu);
   return backlog;
 }
+
+size_t grpc_trickle_get_backlog(grpc_endpoint *ep) {
+  trickle_endpoint *te = (trickle_endpoint *)ep;
+  gpr_mu_lock(&te->mu);
+  size_t backlog = te->write_buffer.length;
+  gpr_mu_unlock(&te->mu);
+  return backlog;
+}
diff --git a/test/core/util/trickle_endpoint.h b/test/core/util/trickle_endpoint.h
index 7e8d9d91e3..e513774eb4 100644
--- a/test/core/util/trickle_endpoint.h
+++ b/test/core/util/trickle_endpoint.h
@@ -43,4 +43,6 @@ grpc_endpoint *grpc_trickle_endpoint_create(grpc_endpoint *wrap,
 size_t grpc_trickle_endpoint_trickle(grpc_exec_ctx *exec_ctx,
                                      grpc_endpoint *endpoint);
 
+size_t grpc_trickle_get_backlog(grpc_endpoint *endpoint);
+
 #endif