From d9a776d4bbf9e8d3989c818c1acce8c61284289a Mon Sep 17 00:00:00 2001
From: Craig Tiller <ctiller@google.com>
Date: Thu, 30 Mar 2017 12:37:11 -0700
Subject: [PATCH] Clamp read sizes based on resource quota

---
 src/core/lib/iomgr/resource_quota.c |  9 +++++++++
 src/core/lib/iomgr/resource_quota.h |  2 ++
 src/core/lib/iomgr/tcp_posix.c      | 19 +++++++++++++------
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/core/lib/iomgr/resource_quota.c b/src/core/lib/iomgr/resource_quota.c
index 8dcd80d0011..c3ee8786517 100644
--- a/src/core/lib/iomgr/resource_quota.c
+++ b/src/core/lib/iomgr/resource_quota.c
@@ -142,6 +142,8 @@ struct grpc_resource_quota {
   /* Amount of free memory in the resource quota */
   int64_t free_pool;
 
+  gpr_atm last_size;
+
   /* Has rq_step been scheduled to occur? */
   bool step_scheduled;
   /* Are we currently reclaiming memory */
@@ -581,6 +583,7 @@ grpc_resource_quota *grpc_resource_quota_create(const char *name) {
   resource_quota->combiner = grpc_combiner_create(NULL);
   resource_quota->free_pool = INT64_MAX;
   resource_quota->size = INT64_MAX;
+  gpr_atm_no_barrier_store(&resource_quota->last_size, GPR_ATM_MAX);
   resource_quota->step_scheduled = false;
   resource_quota->reclaiming = false;
   gpr_atm_no_barrier_store(&resource_quota->memory_usage_estimation, 0);
@@ -643,11 +646,17 @@ void grpc_resource_quota_resize(grpc_resource_quota *resource_quota,
   rq_resize_args *a = gpr_malloc(sizeof(*a));
   a->resource_quota = grpc_resource_quota_ref_internal(resource_quota);
   a->size = (int64_t)size;
+  gpr_atm_no_barrier_store(&resource_quota->last_size,
+                           (gpr_atm)GPR_MIN((size_t)GPR_ATM_MAX, size));
   grpc_closure_init(&a->closure, rq_resize, a, grpc_schedule_on_exec_ctx);
   grpc_closure_sched(&exec_ctx, &a->closure, GRPC_ERROR_NONE);
   grpc_exec_ctx_finish(&exec_ctx);
 }
 
+size_t grpc_resource_quota_peek_size(grpc_resource_quota *resource_quota) {
+  return (size_t)gpr_atm_no_barrier_load(&resource_quota->last_size);
+}
+
 /*******************************************************************************
  * grpc_resource_user channel args api
  */
diff --git a/src/core/lib/iomgr/resource_quota.h b/src/core/lib/iomgr/resource_quota.h
index b9f62cbf83f..6f99be0d512 100644
--- a/src/core/lib/iomgr/resource_quota.h
+++ b/src/core/lib/iomgr/resource_quota.h
@@ -90,6 +90,8 @@ grpc_resource_quota *grpc_resource_quota_from_channel_args(
 double grpc_resource_quota_get_memory_pressure(
     grpc_resource_quota *resource_quota);
 
+size_t grpc_resource_quota_peek_size(grpc_resource_quota *resource_quota);
+
 typedef struct grpc_resource_user grpc_resource_user;
 
 grpc_resource_user *grpc_resource_user_create(
diff --git a/src/core/lib/iomgr/tcp_posix.c b/src/core/lib/iomgr/tcp_posix.c
index bfc7de59f20..5f4b38de2b9 100644
--- a/src/core/lib/iomgr/tcp_posix.c
+++ b/src/core/lib/iomgr/tcp_posix.c
@@ -133,14 +133,21 @@ static void finish_estimate(grpc_tcp *tcp) {
 }
 
 static size_t get_target_read_size(grpc_tcp *tcp) {
-  double pressure = grpc_resource_quota_get_memory_pressure(
-      grpc_resource_user_quota(tcp->resource_user));
+  grpc_resource_quota *rq = grpc_resource_user_quota(tcp->resource_user);
+  double pressure = grpc_resource_quota_get_memory_pressure(rq);
   double target =
       tcp->target_length * (pressure > 0.8 ? (1.0 - pressure) / 0.2 : 1.0);
-  return (((size_t)GPR_CLAMP(target, tcp->min_read_chunk_size,
-                             tcp->max_read_chunk_size)) +
-          255) &
-         ~(size_t)255;
+  size_t sz = (((size_t)GPR_CLAMP(target, tcp->min_read_chunk_size,
+                                  tcp->max_read_chunk_size)) +
+               255) &
+              ~(size_t)255;
+  /* don't use more than 1/16th of the overall resource quota for a single read
+   * alloc */
+  size_t rqmax = grpc_resource_quota_peek_size(rq);
+  if (sz > rqmax / 16 && rqmax > 1024) {
+    sz = rqmax / 16;
+  }
+  return sz;
 }
 
 static grpc_error *tcp_annotate_error(grpc_error *src_error, grpc_tcp *tcp) {