Merge pull request #16823 from AspirinSJL/rq

Account the memory usage of channel and call by resource quota

5 files changed, 56 insertions, 9 deletions

diff --git a/src/core/ext/transport/chttp2/server/chttp2_server.cc b/src/core/ext/transport/chttp2/server/chttp2_server.cc
index 07b304b320..bdb2339e40 100644
--- a/src/core/ext/transport/chttp2/server/chttp2_server.cc
+++ b/src/core/ext/transport/chttp2/server/chttp2_server.cc
@@ -39,6 +39,7 @@
 #include "src/core/lib/channel/handshaker_registry.h"
 #include "src/core/lib/iomgr/endpoint.h"
 #include "src/core/lib/iomgr/resolve_address.h"
+#include "src/core/lib/iomgr/resource_quota.h"
 #include "src/core/lib/iomgr/tcp_server.h"
 #include "src/core/lib/slice/slice_internal.h"
 #include "src/core/lib/surface/api_trace.h"
@@ -114,9 +115,16 @@ static void on_handshake_done(void* arg, grpc_error* error) {
   server_connection_state* connection_state =
       static_cast<server_connection_state*>(args->user_data);
   gpr_mu_lock(&connection_state->svr_state->mu);
+  grpc_resource_user* resource_user = grpc_server_get_default_resource_user(
+      connection_state->svr_state->server);
   if (error != GRPC_ERROR_NONE || connection_state->svr_state->shutdown) {
     const char* error_str = grpc_error_string(error);
     gpr_log(GPR_DEBUG, "Handshaking failed: %s", error_str);
+    grpc_resource_user* resource_user = grpc_server_get_default_resource_user(
+        connection_state->svr_state->server);
+    if (resource_user != nullptr) {
+      grpc_resource_user_free(resource_user, GRPC_RESOURCE_QUOTA_CHANNEL_SIZE);
+    }
     if (error == GRPC_ERROR_NONE && args->endpoint != nullptr) {
       // We were shut down after handshaking completed successfully, so
       // destroy the endpoint here.
@@ -134,13 +142,14 @@ static void on_handshake_done(void* arg, grpc_error* error) {
     // If the handshaking succeeded but there is no endpoint, then the
     // handshaker may have handed off the connection to some external
     // code, so we can just clean up here without creating a transport.
+    // TODO(juanlishen): Do we need to free the memory to resource user?
     if (args->endpoint != nullptr) {
-      grpc_transport* transport =
-          grpc_create_chttp2_transport(args->args, args->endpoint, false);
+      grpc_transport* transport = grpc_create_chttp2_transport(
+          args->args, args->endpoint, false, resource_user);
       grpc_server_setup_transport(
           connection_state->svr_state->server, transport,
           connection_state->accepting_pollset, args->args,
-          grpc_chttp2_transport_get_socket_uuid(transport));
+          grpc_chttp2_transport_get_socket_uuid(transport), resource_user);
       // Use notify_on_receive_settings callback to enforce the
       // handshake deadline.
       connection_state->transport =
@@ -183,6 +192,20 @@ static void on_accept(void* arg, grpc_endpoint* tcp,
     gpr_free(acceptor);
     return;
   }
+  grpc_resource_user* resource_user =
+      grpc_server_get_default_resource_user(state->server);
+  if (resource_user != nullptr &&
+      !grpc_resource_user_safe_alloc(resource_user,
+                                     GRPC_RESOURCE_QUOTA_CHANNEL_SIZE)) {
+    gpr_log(
+        GPR_ERROR,
+        "Memory quota exhausted, rejecting the connection, no handshaking.");
+    gpr_mu_unlock(&state->mu);
+    grpc_endpoint_shutdown(tcp, GRPC_ERROR_NONE);
+    grpc_endpoint_destroy(tcp);
+    gpr_free(acceptor);
+    return;
+  }
   grpc_handshake_manager* handshake_mgr = grpc_handshake_manager_create();
   grpc_handshake_manager_pending_list_add(&state->pending_handshake_mgrs,
                                           handshake_mgr);
diff --git a/src/core/ext/transport/chttp2/transport/chttp2_transport.cc b/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
index 8a481bb7d5..04874a9494 100644
--- a/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
+++ b/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
@@ -478,7 +478,8 @@ static void init_keepalive_pings_if_enabled(grpc_chttp2_transport* t) {
 
 static void init_transport(grpc_chttp2_transport* t,
                            const grpc_channel_args* channel_args,
-                           grpc_endpoint* ep, bool is_client) {
+                           grpc_endpoint* ep, bool is_client,
+                           grpc_resource_user* resource_user) {
   GPR_ASSERT(strlen(GRPC_CHTTP2_CLIENT_CONNECT_STRING) ==
              GRPC_CHTTP2_CLIENT_CONNECT_STRLEN);
 
@@ -491,6 +492,7 @@ static void init_transport(grpc_chttp2_transport* t,
   t->endpoint_reading = 1;
   t->next_stream_id = is_client ? 1 : 2;
   t->is_client = is_client;
+  t->resource_user = resource_user;
   t->deframe_state = is_client ? GRPC_DTS_FH_0 : GRPC_DTS_CLIENT_PREFIX_0;
   t->is_first_frame = true;
   grpc_connectivity_state_init(
@@ -778,6 +780,10 @@ static void destroy_stream_locked(void* sp, grpc_error* error) {
 
   s->flow_control.Destroy();
 
+  if (t->resource_user != nullptr) {
+    grpc_resource_user_free(t->resource_user, GRPC_RESOURCE_QUOTA_CALL_SIZE);
+  }
+
   GRPC_CHTTP2_UNREF_TRANSPORT(t, "stream");
 
   GRPC_CLOSURE_SCHED(s->destroy_stream_arg, GRPC_ERROR_NONE);
@@ -816,7 +822,21 @@ grpc_chttp2_stream* grpc_chttp2_parsing_accept_stream(grpc_chttp2_transport* t,
   if (t->channel_callback.accept_stream == nullptr) {
     return nullptr;
   }
-  grpc_chttp2_stream* accepting;
+  // Don't accept the stream if memory quota doesn't allow. Note that we should
+  // simply refuse the stream here instead of canceling the stream after it's
+  // accepted since the latter will create the call which costs much memory.
+  if (t->resource_user != nullptr &&
+      !grpc_resource_user_safe_alloc(t->resource_user,
+                                     GRPC_RESOURCE_QUOTA_CALL_SIZE)) {
+    gpr_log(GPR_ERROR, "Memory exhausted, rejecting the stream.");
+    grpc_slice_buffer_add(
+        &t->qbuf,
+        grpc_chttp2_rst_stream_create(
+            id, static_cast<uint32_t>(GRPC_HTTP2_REFUSED_STREAM), nullptr));
+    grpc_chttp2_initiate_write(t, GRPC_CHTTP2_INITIATE_WRITE_RST_STREAM);
+    return nullptr;
+  }
+  grpc_chttp2_stream* accepting = nullptr;
   GPR_ASSERT(t->accepting_stream == nullptr);
   t->accepting_stream = &accepting;
   t->channel_callback.accept_stream(t->channel_callback.accept_stream_user_data,
@@ -3185,10 +3205,11 @@ intptr_t grpc_chttp2_transport_get_socket_uuid(grpc_transport* transport) {
 }
 
 grpc_transport* grpc_create_chttp2_transport(
-    const grpc_channel_args* channel_args, grpc_endpoint* ep, bool is_client) {
+    const grpc_channel_args* channel_args, grpc_endpoint* ep, bool is_client,
+    grpc_resource_user* resource_user) {
   grpc_chttp2_transport* t = static_cast<grpc_chttp2_transport*>(
       gpr_zalloc(sizeof(grpc_chttp2_transport)));
-  init_transport(t, channel_args, ep, is_client);
+  init_transport(t, channel_args, ep, is_client, resource_user);
   return &t->base;
 }
 
diff --git a/src/core/ext/transport/chttp2/transport/chttp2_transport.h b/src/core/ext/transport/chttp2/transport/chttp2_transport.h
index e5872fee43..b3fe1c082e 100644
--- a/src/core/ext/transport/chttp2/transport/chttp2_transport.h
+++ b/src/core/ext/transport/chttp2/transport/chttp2_transport.h
@@ -32,7 +32,8 @@ extern grpc_core::DebugOnlyTraceFlag grpc_trace_chttp2_refcount;
 extern bool g_flow_control_enabled;
 
 grpc_transport* grpc_create_chttp2_transport(
-    const grpc_channel_args* channel_args, grpc_endpoint* ep, bool is_client);
+    const grpc_channel_args* channel_args, grpc_endpoint* ep, bool is_client,
+    grpc_resource_user* resource_user = nullptr);
 
 intptr_t grpc_chttp2_transport_get_socket_uuid(grpc_transport* transport);
 
diff --git a/src/core/ext/transport/chttp2/transport/frame_rst_stream.cc b/src/core/ext/transport/chttp2/transport/frame_rst_stream.cc
index 4bdd4309a4..a0a7534594 100644
--- a/src/core/ext/transport/chttp2/transport/frame_rst_stream.cc
+++ b/src/core/ext/transport/chttp2/transport/frame_rst_stream.cc
@@ -32,7 +32,7 @@ grpc_slice grpc_chttp2_rst_stream_create(uint32_t id, uint32_t code,
                                          grpc_transport_one_way_stats* stats) {
   static const size_t frame_size = 13;
   grpc_slice slice = GRPC_SLICE_MALLOC(frame_size);
-  stats->framing_bytes += frame_size;
+  if (stats != nullptr) stats->framing_bytes += frame_size;
   uint8_t* p = GRPC_SLICE_START_PTR(slice);
 
   // Frame size.
diff --git a/src/core/ext/transport/chttp2/transport/internal.h b/src/core/ext/transport/chttp2/transport/internal.h
index ff26dd9255..202017641b 100644
--- a/src/core/ext/transport/chttp2/transport/internal.h
+++ b/src/core/ext/transport/chttp2/transport/internal.h
@@ -285,6 +285,8 @@ struct grpc_chttp2_transport {
   grpc_endpoint* ep;
   char* peer_string;
 
+  grpc_resource_user* resource_user;
+
   grpc_combiner* combiner;
 
   grpc_closure* notify_on_receive_settings;