diff --git a/CMakeLists.txt b/CMakeLists.txt index 8faf34698dd..d5fc01ad30c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2383,6 +2383,7 @@ add_library(grpc src/core/lib/promise/trace.cc src/core/lib/resource_quota/api.cc src/core/lib/resource_quota/arena.cc + src/core/lib/resource_quota/connection_quota.cc src/core/lib/resource_quota/memory_quota.cc src/core/lib/resource_quota/periodic_update.cc src/core/lib/resource_quota/resource_quota.cc @@ -3145,6 +3146,7 @@ add_library(grpc_unsecure src/core/lib/promise/trace.cc src/core/lib/resource_quota/api.cc src/core/lib/resource_quota/arena.cc + src/core/lib/resource_quota/connection_quota.cc src/core/lib/resource_quota/memory_quota.cc src/core/lib/resource_quota/periodic_update.cc src/core/lib/resource_quota/resource_quota.cc @@ -5252,6 +5254,7 @@ add_library(grpc_authorization_provider src/core/lib/promise/trace.cc src/core/lib/resource_quota/api.cc src/core/lib/resource_quota/arena.cc + src/core/lib/resource_quota/connection_quota.cc src/core/lib/resource_quota/memory_quota.cc src/core/lib/resource_quota/periodic_update.cc src/core/lib/resource_quota/resource_quota.cc @@ -10177,6 +10180,7 @@ add_executable(chunked_vector_test src/core/lib/promise/activity.cc src/core/lib/promise/trace.cc src/core/lib/resource_quota/arena.cc + src/core/lib/resource_quota/connection_quota.cc src/core/lib/resource_quota/memory_quota.cc src/core/lib/resource_quota/periodic_update.cc src/core/lib/resource_quota/resource_quota.cc @@ -14175,6 +14179,7 @@ add_executable(flow_control_test src/core/lib/iomgr/iomgr_internal.cc src/core/lib/promise/activity.cc src/core/lib/promise/trace.cc + src/core/lib/resource_quota/connection_quota.cc src/core/lib/resource_quota/memory_quota.cc src/core/lib/resource_quota/periodic_update.cc src/core/lib/resource_quota/resource_quota.cc @@ -14263,6 +14268,7 @@ add_executable(for_each_test src/core/lib/promise/activity.cc src/core/lib/promise/trace.cc src/core/lib/resource_quota/arena.cc + src/core/lib/resource_quota/connection_quota.cc src/core/lib/resource_quota/memory_quota.cc src/core/lib/resource_quota/periodic_update.cc src/core/lib/resource_quota/resource_quota.cc @@ -17709,6 +17715,7 @@ add_executable(interceptor_list_test src/core/lib/promise/activity.cc src/core/lib/promise/trace.cc src/core/lib/resource_quota/arena.cc + src/core/lib/resource_quota/connection_quota.cc src/core/lib/resource_quota/memory_quota.cc src/core/lib/resource_quota/periodic_update.cc src/core/lib/resource_quota/resource_quota.cc @@ -18900,6 +18907,7 @@ add_executable(map_pipe_test src/core/lib/promise/activity.cc src/core/lib/promise/trace.cc src/core/lib/resource_quota/arena.cc + src/core/lib/resource_quota/connection_quota.cc src/core/lib/resource_quota/memory_quota.cc src/core/lib/resource_quota/periodic_update.cc src/core/lib/resource_quota/resource_quota.cc diff --git a/Makefile b/Makefile index 4bdbbec686f..23f195b095a 100644 --- a/Makefile +++ b/Makefile @@ -1309,6 +1309,7 @@ LIBGRPC_SRC = \ src/core/lib/promise/trace.cc \ src/core/lib/resource_quota/api.cc \ src/core/lib/resource_quota/arena.cc \ + src/core/lib/resource_quota/connection_quota.cc \ src/core/lib/resource_quota/memory_quota.cc \ src/core/lib/resource_quota/periodic_update.cc \ src/core/lib/resource_quota/resource_quota.cc \ diff --git a/Package.swift b/Package.swift index 64d6d866bf7..9737494d3bc 100644 --- a/Package.swift +++ b/Package.swift @@ -1608,6 +1608,8 @@ let package = Package( "src/core/lib/resource_quota/api.h", "src/core/lib/resource_quota/arena.cc", "src/core/lib/resource_quota/arena.h", + "src/core/lib/resource_quota/connection_quota.cc", + "src/core/lib/resource_quota/connection_quota.h", "src/core/lib/resource_quota/memory_quota.cc", "src/core/lib/resource_quota/memory_quota.h", "src/core/lib/resource_quota/periodic_update.cc", diff --git a/build_autogenerated.yaml b/build_autogenerated.yaml index 967ad017f7d..e9478a2a4b8 100644 --- a/build_autogenerated.yaml +++ b/build_autogenerated.yaml @@ -1041,6 +1041,7 @@ libs: - src/core/lib/promise/try_seq.h - src/core/lib/resource_quota/api.h - src/core/lib/resource_quota/arena.h + - src/core/lib/resource_quota/connection_quota.h - src/core/lib/resource_quota/memory_quota.h - src/core/lib/resource_quota/periodic_update.h - src/core/lib/resource_quota/resource_quota.h @@ -1837,6 +1838,7 @@ libs: - src/core/lib/promise/trace.cc - src/core/lib/resource_quota/api.cc - src/core/lib/resource_quota/arena.cc + - src/core/lib/resource_quota/connection_quota.cc - src/core/lib/resource_quota/memory_quota.cc - src/core/lib/resource_quota/periodic_update.cc - src/core/lib/resource_quota/resource_quota.cc @@ -2542,6 +2544,7 @@ libs: - src/core/lib/promise/try_seq.h - src/core/lib/resource_quota/api.h - src/core/lib/resource_quota/arena.h + - src/core/lib/resource_quota/connection_quota.h - src/core/lib/resource_quota/memory_quota.h - src/core/lib/resource_quota/periodic_update.h - src/core/lib/resource_quota/resource_quota.h @@ -2961,6 +2964,7 @@ libs: - src/core/lib/promise/trace.cc - src/core/lib/resource_quota/api.cc - src/core/lib/resource_quota/arena.cc + - src/core/lib/resource_quota/connection_quota.cc - src/core/lib/resource_quota/memory_quota.cc - src/core/lib/resource_quota/periodic_update.cc - src/core/lib/resource_quota/resource_quota.cc @@ -4619,6 +4623,7 @@ libs: - src/core/lib/promise/try_seq.h - src/core/lib/resource_quota/api.h - src/core/lib/resource_quota/arena.h + - src/core/lib/resource_quota/connection_quota.h - src/core/lib/resource_quota/memory_quota.h - src/core/lib/resource_quota/periodic_update.h - src/core/lib/resource_quota/resource_quota.h @@ -4916,6 +4921,7 @@ libs: - src/core/lib/promise/trace.cc - src/core/lib/resource_quota/api.cc - src/core/lib/resource_quota/arena.cc + - src/core/lib/resource_quota/connection_quota.cc - src/core/lib/resource_quota/memory_quota.cc - src/core/lib/resource_quota/periodic_update.cc - src/core/lib/resource_quota/resource_quota.cc @@ -7689,6 +7695,7 @@ targets: - src/core/lib/promise/seq.h - src/core/lib/promise/trace.h - src/core/lib/resource_quota/arena.h + - src/core/lib/resource_quota/connection_quota.h - src/core/lib/resource_quota/memory_quota.h - src/core/lib/resource_quota/periodic_update.h - src/core/lib/resource_quota/resource_quota.h @@ -7738,6 +7745,7 @@ targets: - src/core/lib/promise/activity.cc - src/core/lib/promise/trace.cc - src/core/lib/resource_quota/arena.cc + - src/core/lib/resource_quota/connection_quota.cc - src/core/lib/resource_quota/memory_quota.cc - src/core/lib/resource_quota/periodic_update.cc - src/core/lib/resource_quota/resource_quota.cc @@ -9785,6 +9793,7 @@ targets: - src/core/lib/promise/race.h - src/core/lib/promise/seq.h - src/core/lib/promise/trace.h + - src/core/lib/resource_quota/connection_quota.h - src/core/lib/resource_quota/memory_quota.h - src/core/lib/resource_quota/periodic_update.h - src/core/lib/resource_quota/resource_quota.h @@ -9839,6 +9848,7 @@ targets: - src/core/lib/iomgr/iomgr_internal.cc - src/core/lib/promise/activity.cc - src/core/lib/promise/trace.cc + - src/core/lib/resource_quota/connection_quota.cc - src/core/lib/resource_quota/memory_quota.cc - src/core/lib/resource_quota/periodic_update.cc - src/core/lib/resource_quota/resource_quota.cc @@ -9928,6 +9938,7 @@ targets: - src/core/lib/promise/trace.h - src/core/lib/promise/try_seq.h - src/core/lib/resource_quota/arena.h + - src/core/lib/resource_quota/connection_quota.h - src/core/lib/resource_quota/memory_quota.h - src/core/lib/resource_quota/periodic_update.h - src/core/lib/resource_quota/resource_quota.h @@ -9978,6 +9989,7 @@ targets: - src/core/lib/promise/activity.cc - src/core/lib/promise/trace.cc - src/core/lib/resource_quota/arena.cc + - src/core/lib/resource_quota/connection_quota.cc - src/core/lib/resource_quota/memory_quota.cc - src/core/lib/resource_quota/periodic_update.cc - src/core/lib/resource_quota/resource_quota.cc @@ -11578,6 +11590,7 @@ targets: - src/core/lib/promise/seq.h - src/core/lib/promise/trace.h - src/core/lib/resource_quota/arena.h + - src/core/lib/resource_quota/connection_quota.h - src/core/lib/resource_quota/memory_quota.h - src/core/lib/resource_quota/periodic_update.h - src/core/lib/resource_quota/resource_quota.h @@ -11628,6 +11641,7 @@ targets: - src/core/lib/promise/activity.cc - src/core/lib/promise/trace.cc - src/core/lib/resource_quota/arena.cc + - src/core/lib/resource_quota/connection_quota.cc - src/core/lib/resource_quota/memory_quota.cc - src/core/lib/resource_quota/periodic_update.cc - src/core/lib/resource_quota/resource_quota.cc @@ -12256,6 +12270,7 @@ targets: - src/core/lib/promise/trace.h - src/core/lib/promise/try_seq.h - src/core/lib/resource_quota/arena.h + - src/core/lib/resource_quota/connection_quota.h - src/core/lib/resource_quota/memory_quota.h - src/core/lib/resource_quota/periodic_update.h - src/core/lib/resource_quota/resource_quota.h @@ -12306,6 +12321,7 @@ targets: - src/core/lib/promise/activity.cc - src/core/lib/promise/trace.cc - src/core/lib/resource_quota/arena.cc + - src/core/lib/resource_quota/connection_quota.cc - src/core/lib/resource_quota/memory_quota.cc - src/core/lib/resource_quota/periodic_update.cc - src/core/lib/resource_quota/resource_quota.cc diff --git a/config.m4 b/config.m4 index 4600f743302..71180188c82 100644 --- a/config.m4 +++ b/config.m4 @@ -684,6 +684,7 @@ if test "$PHP_GRPC" != "no"; then src/core/lib/promise/trace.cc \ src/core/lib/resource_quota/api.cc \ src/core/lib/resource_quota/arena.cc \ + src/core/lib/resource_quota/connection_quota.cc \ src/core/lib/resource_quota/memory_quota.cc \ src/core/lib/resource_quota/periodic_update.cc \ src/core/lib/resource_quota/resource_quota.cc \ diff --git a/config.w32 b/config.w32 index dcf40dbb7d3..b9e6c920b2e 100644 --- a/config.w32 +++ b/config.w32 @@ -649,6 +649,7 @@ if (PHP_GRPC != "no") { "src\\core\\lib\\promise\\trace.cc " + "src\\core\\lib\\resource_quota\\api.cc " + "src\\core\\lib\\resource_quota\\arena.cc " + + "src\\core\\lib\\resource_quota\\connection_quota.cc " + "src\\core\\lib\\resource_quota\\memory_quota.cc " + "src\\core\\lib\\resource_quota\\periodic_update.cc " + "src\\core\\lib\\resource_quota\\resource_quota.cc " + diff --git a/gRPC-C++.podspec b/gRPC-C++.podspec index 84a7bfef691..bad4a958ef8 100644 --- a/gRPC-C++.podspec +++ b/gRPC-C++.podspec @@ -1147,6 +1147,7 @@ Pod::Spec.new do |s| 'src/core/lib/promise/try_seq.h', 'src/core/lib/resource_quota/api.h', 'src/core/lib/resource_quota/arena.h', + 'src/core/lib/resource_quota/connection_quota.h', 'src/core/lib/resource_quota/memory_quota.h', 'src/core/lib/resource_quota/periodic_update.h', 'src/core/lib/resource_quota/resource_quota.h', @@ -2410,6 +2411,7 @@ Pod::Spec.new do |s| 'src/core/lib/promise/try_seq.h', 'src/core/lib/resource_quota/api.h', 'src/core/lib/resource_quota/arena.h', + 'src/core/lib/resource_quota/connection_quota.h', 'src/core/lib/resource_quota/memory_quota.h', 'src/core/lib/resource_quota/periodic_update.h', 'src/core/lib/resource_quota/resource_quota.h', diff --git a/gRPC-Core.podspec b/gRPC-Core.podspec index 35adc20a350..ca1624e7ac2 100644 --- a/gRPC-Core.podspec +++ b/gRPC-Core.podspec @@ -1722,6 +1722,8 @@ Pod::Spec.new do |s| 'src/core/lib/resource_quota/api.h', 'src/core/lib/resource_quota/arena.cc', 'src/core/lib/resource_quota/arena.h', + 'src/core/lib/resource_quota/connection_quota.cc', + 'src/core/lib/resource_quota/connection_quota.h', 'src/core/lib/resource_quota/memory_quota.cc', 'src/core/lib/resource_quota/memory_quota.h', 'src/core/lib/resource_quota/periodic_update.cc', @@ -3191,6 +3193,7 @@ Pod::Spec.new do |s| 'src/core/lib/promise/try_seq.h', 'src/core/lib/resource_quota/api.h', 'src/core/lib/resource_quota/arena.h', + 'src/core/lib/resource_quota/connection_quota.h', 'src/core/lib/resource_quota/memory_quota.h', 'src/core/lib/resource_quota/periodic_update.h', 'src/core/lib/resource_quota/resource_quota.h', diff --git a/grpc.gemspec b/grpc.gemspec index eaed281aa80..12d4db8b99b 100644 --- a/grpc.gemspec +++ b/grpc.gemspec @@ -1614,6 +1614,8 @@ Gem::Specification.new do |s| s.files += %w( src/core/lib/resource_quota/api.h ) s.files += %w( src/core/lib/resource_quota/arena.cc ) s.files += %w( src/core/lib/resource_quota/arena.h ) + s.files += %w( src/core/lib/resource_quota/connection_quota.cc ) + s.files += %w( src/core/lib/resource_quota/connection_quota.h ) s.files += %w( src/core/lib/resource_quota/memory_quota.cc ) s.files += %w( src/core/lib/resource_quota/memory_quota.h ) s.files += %w( src/core/lib/resource_quota/periodic_update.cc ) diff --git a/grpc.gyp b/grpc.gyp index 1f7a89d8d52..df31e5f4f3b 100644 --- a/grpc.gyp +++ b/grpc.gyp @@ -883,6 +883,7 @@ 'src/core/lib/promise/trace.cc', 'src/core/lib/resource_quota/api.cc', 'src/core/lib/resource_quota/arena.cc', + 'src/core/lib/resource_quota/connection_quota.cc', 'src/core/lib/resource_quota/memory_quota.cc', 'src/core/lib/resource_quota/periodic_update.cc', 'src/core/lib/resource_quota/resource_quota.cc', @@ -1422,6 +1423,7 @@ 'src/core/lib/promise/trace.cc', 'src/core/lib/resource_quota/api.cc', 'src/core/lib/resource_quota/arena.cc', + 'src/core/lib/resource_quota/connection_quota.cc', 'src/core/lib/resource_quota/memory_quota.cc', 'src/core/lib/resource_quota/periodic_update.cc', 'src/core/lib/resource_quota/resource_quota.cc', @@ -2228,6 +2230,7 @@ 'src/core/lib/promise/trace.cc', 'src/core/lib/resource_quota/api.cc', 'src/core/lib/resource_quota/arena.cc', + 'src/core/lib/resource_quota/connection_quota.cc', 'src/core/lib/resource_quota/memory_quota.cc', 'src/core/lib/resource_quota/periodic_update.cc', 'src/core/lib/resource_quota/resource_quota.cc', diff --git a/include/grpc/impl/channel_arg_names.h b/include/grpc/impl/channel_arg_names.h index b42bf7026ef..7ba19bebb17 100644 --- a/include/grpc/impl/channel_arg_names.h +++ b/include/grpc/impl/channel_arg_names.h @@ -395,6 +395,10 @@ * factory. */ #define GRPC_ARG_EVENT_ENGINE_USE_MEMORY_ALLOCATOR_FACTORY \ "grpc.event_engine_use_memory_allocator_factory" +/** Configure the max number of allowed incoming connections to the server. + * If unspecified, it is unlimited */ +#define GRPC_ARG_MAX_ALLOWED_INCOMING_CONNECTIONS \ + "grpc.max_allowed_incoming_connections" /** \} */ #endif /* GRPC_IMPL_CHANNEL_ARG_NAMES_H */ diff --git a/package.xml b/package.xml index fa35a2c4033..1ba949ae253 100644 --- a/package.xml +++ b/package.xml @@ -1596,6 +1596,8 @@ + + diff --git a/src/core/BUILD b/src/core/BUILD index f0c58279faf..46708db643f 100644 --- a/src/core/BUILD +++ b/src/core/BUILD @@ -1399,6 +1399,23 @@ grpc_cc_library( ], ) +grpc_cc_library( + name = "connection_quota", + srcs = [ + "lib/resource_quota/connection_quota.cc", + ], + hdrs = [ + "lib/resource_quota/connection_quota.h", + ], + external_deps = ["absl/base:core_headers"], + deps = [ + "memory_quota", + "ref_counted", + "//:gpr", + "//:ref_counted_ptr", + ], +) + grpc_cc_library( name = "resource_quota_trace", srcs = [ @@ -1426,6 +1443,7 @@ grpc_cc_library( "@grpc:alt_grpc_base_legacy", ], deps = [ + "connection_quota", "memory_quota", "ref_counted", "thread_quota", @@ -6747,6 +6765,7 @@ grpc_cc_library( "channel_args", "channel_args_endpoint_config", "closure", + "connection_quota", "error", "error_utils", "grpc_insecure_credentials", diff --git a/src/core/ext/transport/chttp2/server/chttp2_server.cc b/src/core/ext/transport/chttp2/server/chttp2_server.cc index dea28445a4f..7ac39d01cff 100644 --- a/src/core/ext/transport/chttp2/server/chttp2_server.cc +++ b/src/core/ext/transport/chttp2/server/chttp2_server.cc @@ -72,6 +72,7 @@ #include "src/core/lib/iomgr/tcp_server.h" #include "src/core/lib/iomgr/unix_sockets_posix.h" #include "src/core/lib/iomgr/vsock.h" +#include "src/core/lib/resource_quota/connection_quota.h" #include "src/core/lib/resource_quota/memory_quota.h" #include "src/core/lib/resource_quota/resource_quota.h" #include "src/core/lib/security/credentials/credentials.h" @@ -283,6 +284,7 @@ class Chttp2ServerListener : public Server::ListenerInterface { grpc_closure* on_destroy_done_ ABSL_GUARDED_BY(mu_) = nullptr; RefCountedPtr channelz_listen_socket_; MemoryQuotaRefPtr memory_quota_; + ConnectionQuotaRefPtr connection_quota_; }; // @@ -504,6 +506,18 @@ void Chttp2ServerListener::ActiveConnection::HandshakingState::OnHandshakeDone( } else { // Remove the connection from the connections_ map since OnClose() // will not be invoked when a config fetcher is set. + auto connection_quota = + self->connection_->listener_->connection_quota_->Ref() + .release(); + auto on_close_transport = [](void* arg, + grpc_error_handle /*handle*/) { + ConnectionQuota* connection_quota = + static_cast(arg); + connection_quota->ReleaseConnections(1); + connection_quota->Unref(); + }; + on_close = GRPC_CLOSURE_CREATE(on_close_transport, connection_quota, + grpc_schedule_on_exec_ctx_); cleanup_connection = true; } grpc_chttp2_transport_start_reading(transport, args->read_buffer, @@ -658,6 +672,7 @@ void Chttp2ServerListener::ActiveConnection::OnClose( self->drain_grace_timer_handle_.reset(); } } + self->listener_->connection_quota_->ReleaseConnections(1); self->Unref(); } @@ -762,7 +777,14 @@ Chttp2ServerListener::Chttp2ServerListener( : server_(server), args_modifier_(args_modifier), args_(args), - memory_quota_(args.GetObject()->memory_quota()) { + memory_quota_(args.GetObject()->memory_quota()), + connection_quota_(MakeRefCounted()) { + auto max_allowed_incoming_connections = + args.GetInt(GRPC_ARG_MAX_ALLOWED_INCOMING_CONNECTIONS); + if (max_allowed_incoming_connections.has_value()) { + connection_quota_->SetMaxIncomingConnections( + max_allowed_incoming_connections.value()); + } GRPC_CLOSURE_INIT(&tcp_server_shutdown_complete_, TcpServerShutdownComplete, this, grpc_schedule_on_exec_ctx); } @@ -821,6 +843,14 @@ void Chttp2ServerListener::OnAccept(void* arg, grpc_endpoint* tcp, grpc_endpoint_destroy(tcp); gpr_free(acceptor); }; + if (!self->connection_quota_->AllowIncomingConnection( + self->memory_quota_, grpc_endpoint_get_peer(tcp))) { + grpc_error_handle error = GRPC_ERROR_CREATE( + "Rejected incoming connection because configured connection quota " + "limits have been exceeded."); + endpoint_cleanup(error); + return; + } if (self->server_->config_fetcher() != nullptr) { if (connection_manager == nullptr) { grpc_error_handle error = GRPC_ERROR_CREATE( diff --git a/src/core/lib/resource_quota/connection_quota.cc b/src/core/lib/resource_quota/connection_quota.cc new file mode 100644 index 00000000000..75d3c7b8edd --- /dev/null +++ b/src/core/lib/resource_quota/connection_quota.cc @@ -0,0 +1,71 @@ +// Copyright 2024 gRPC authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "src/core/lib/resource_quota/connection_quota.h" + +#include +#include + +#include + +namespace grpc_core { + +ConnectionQuota::ConnectionQuota() = default; + +void ConnectionQuota::SetMaxIncomingConnections(int max_incoming_connections) { + // The maximum can only be configured once. + GPR_ASSERT(max_incoming_connections < INT_MAX); + GPR_ASSERT(max_incoming_connections_.exchange(max_incoming_connections, + std::memory_order_release) == + INT_MAX); +} + +// Returns true if the incoming connection is allowed to be accepted on the +// server. +bool ConnectionQuota::AllowIncomingConnection(MemoryQuotaRefPtr mem_quota, + absl::string_view /*peer*/) { + if (mem_quota->IsMemoryPressureHigh()) { + return false; + } + + if (max_incoming_connections_.load(std::memory_order_relaxed) == INT_MAX) { + return true; + } + + int curr_active_connections = + active_incoming_connections_.load(std::memory_order_acquire); + do { + if (curr_active_connections >= + max_incoming_connections_.load(std::memory_order_relaxed)) { + return false; + } + } while (!active_incoming_connections_.compare_exchange_weak( + curr_active_connections, curr_active_connections + 1, + std::memory_order_acq_rel, std::memory_order_relaxed)); + return true; +} + +// Mark connections as closed. +void ConnectionQuota::ReleaseConnections(int num_connections) { + if (max_incoming_connections_.load(std::memory_order_relaxed) == INT_MAX) { + return; + } + GPR_ASSERT(active_incoming_connections_.fetch_sub( + num_connections, std::memory_order_acq_rel) >= + num_connections); +} + +} // namespace grpc_core diff --git a/src/core/lib/resource_quota/connection_quota.h b/src/core/lib/resource_quota/connection_quota.h new file mode 100644 index 00000000000..32815f57600 --- /dev/null +++ b/src/core/lib/resource_quota/connection_quota.h @@ -0,0 +1,61 @@ +// Copyright 2024 gRPC authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef GRPC_SRC_CORE_LIB_RESOURCE_QUOTA_CONNECTION_QUOTA_H +#define GRPC_SRC_CORE_LIB_RESOURCE_QUOTA_CONNECTION_QUOTA_H + +#include + +#include +#include + +#include "absl/base/thread_annotations.h" + +#include "src/core/lib/gprpp/ref_counted.h" +#include "src/core/lib/gprpp/ref_counted_ptr.h" +#include "src/core/lib/gprpp/sync.h" +#include "src/core/lib/resource_quota/memory_quota.h" + +namespace grpc_core { + +// Tracks the amount of threads in a resource quota. +class ConnectionQuota : public RefCounted { + public: + ConnectionQuota(); + ~ConnectionQuota() override = default; + + ConnectionQuota(const ConnectionQuota&) = delete; + ConnectionQuota& operator=(const ConnectionQuota&) = delete; + + // Set the maximum number of allowed incoming connections on the server. + void SetMaxIncomingConnections(int max_incoming_connections); + + // Returns true if the incoming connection is allowed to be accepted on the + // server. + bool AllowIncomingConnection(MemoryQuotaRefPtr mem_quota, + absl::string_view peer); + + // Mark connections as closed. + void ReleaseConnections(int num_connections); + + private: + std::atomic active_incoming_connections_{0}; + std::atomic max_incoming_connections_{std::numeric_limits::max()}; +}; + +using ConnectionQuotaRefPtr = RefCountedPtr; + +} // namespace grpc_core + +#endif // GRPC_SRC_CORE_LIB_RESOURCE_QUOTA_CONNECTION_QUOTA_H diff --git a/src/python/grpcio/grpc_core_dependencies.py b/src/python/grpcio/grpc_core_dependencies.py index 857b62a8259..f3f4a909fd5 100644 --- a/src/python/grpcio/grpc_core_dependencies.py +++ b/src/python/grpcio/grpc_core_dependencies.py @@ -658,6 +658,7 @@ CORE_SOURCE_FILES = [ 'src/core/lib/promise/trace.cc', 'src/core/lib/resource_quota/api.cc', 'src/core/lib/resource_quota/arena.cc', + 'src/core/lib/resource_quota/connection_quota.cc', 'src/core/lib/resource_quota/memory_quota.cc', 'src/core/lib/resource_quota/periodic_update.cc', 'src/core/lib/resource_quota/resource_quota.cc', diff --git a/test/cpp/end2end/resource_quota_end2end_stress_test.cc b/test/cpp/end2end/resource_quota_end2end_stress_test.cc index 83af5a2fe27..4a8f28deaeb 100644 --- a/test/cpp/end2end/resource_quota_end2end_stress_test.cc +++ b/test/cpp/end2end/resource_quota_end2end_stress_test.cc @@ -53,7 +53,6 @@ class EchoClientUnaryReactor : public grpc::ClientUnaryReactor { EchoClientUnaryReactor(ClientContext* ctx, EchoTestService::Stub* stub, const std::string payload, Status* status) : ctx_(ctx), payload_(payload), status_(status) { - ctx_->set_wait_for_ready(true); request_.set_message(payload); stub->async()->Echo(ctx_, &request_, &response_, this); StartCall(); @@ -120,6 +119,7 @@ class End2EndResourceQuotaUnaryTest : public ::testing::Test { Status status; auto stub = EchoTestService::NewStub( CreateChannel(server_address_, grpc::InsecureChannelCredentials())); + ctx.set_wait_for_ready(true); EchoClientUnaryReactor reactor(&ctx, stub.get(), payload_, &status); reactor.Await(); } @@ -145,6 +145,78 @@ class End2EndResourceQuotaUnaryTest : public ::testing::Test { TEST_F(End2EndResourceQuotaUnaryTest, MultipleUnaryRPCTest) { MakeGrpcCalls(); } +class End2EndConnectionQuotaTest : public ::testing::TestWithParam { + protected: + End2EndConnectionQuotaTest() { + port_ = grpc_pick_unused_port_or_die(); + server_address_ = absl::StrCat("[::]:", port_); + payload_ = std::string(kPayloadSizeBytes, 'a'); + ServerBuilder builder; + builder.AddListeningPort(server_address_, InsecureServerCredentials()); + builder.AddChannelArgument(GRPC_ARG_MAX_ALLOWED_INCOMING_CONNECTIONS, + GetParam()); + builder.AddChannelArgument( + GRPC_ARG_HTTP2_MIN_RECV_PING_INTERVAL_WITHOUT_DATA_MS, 10000); + builder.RegisterService(&grpc_service_); + server_ = builder.BuildAndStart(); + } + + ~End2EndConnectionQuotaTest() override { server_->Shutdown(); } + + std::unique_ptr CreateGrpcChannelStub() { + grpc::ChannelArguments args; + args.SetInt(GRPC_ARG_USE_LOCAL_SUBCHANNEL_POOL, 1); + args.SetInt(GRPC_ARG_ENABLE_RETRIES, 0); + args.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS, 20000); + args.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS, 10000); + args.SetInt(GRPC_ARG_HTTP2_MIN_SENT_PING_INTERVAL_WITHOUT_DATA_MS, 15000); + args.SetInt(GRPC_ARG_KEEPALIVE_PERMIT_WITHOUT_CALLS, 1); + + return EchoTestService::NewStub( + CreateCustomChannel(absl::StrCat("ipv6:[::1]:", port_), + grpc::InsecureChannelCredentials(), args)); + } + + void TestExceedingConnectionQuota() { + const int kNumConnections = 2 * GetParam(); + std::vector> stubs; + stubs.reserve(kNumConnections); + for (int i = 0; i < kNumConnections; i++) { + stubs.push_back(CreateGrpcChannelStub()); + } + for (int i = 0; i < kNumConnections; ++i) { + ClientContext ctx; + Status status; + ctx.set_wait_for_ready(false); + EchoClientUnaryReactor reactor(&ctx, stubs[i].get(), payload_, &status); + reactor.Await(); + // The first half RPCs should succeed. + if (i < kNumConnections / 2) { + EXPECT_TRUE(status.ok()); + + } else { + // The second half should fail because they would attempt to create a + // new connection and fail since it would exceed the connection quota + // limit set at the server. + EXPECT_FALSE(status.ok()); + } + } + } + + int port_; + std::unique_ptr server_; + string server_address_; + GrpcCallbackServiceImpl grpc_service_; + std::string payload_; +}; + +TEST_P(End2EndConnectionQuotaTest, ConnectionQuotaTest) { + TestExceedingConnectionQuota(); +} + +INSTANTIATE_TEST_SUITE_P(ConnectionQuotaParamTest, End2EndConnectionQuotaTest, + ::testing::ValuesIn({10, 100})); + } // namespace testing } // namespace grpc diff --git a/tools/doxygen/Doxyfile.c++.internal b/tools/doxygen/Doxyfile.c++.internal index 1fbf624c5d3..39f3d505ed3 100644 --- a/tools/doxygen/Doxyfile.c++.internal +++ b/tools/doxygen/Doxyfile.c++.internal @@ -2613,6 +2613,8 @@ src/core/lib/resource_quota/api.cc \ src/core/lib/resource_quota/api.h \ src/core/lib/resource_quota/arena.cc \ src/core/lib/resource_quota/arena.h \ +src/core/lib/resource_quota/connection_quota.cc \ +src/core/lib/resource_quota/connection_quota.h \ src/core/lib/resource_quota/memory_quota.cc \ src/core/lib/resource_quota/memory_quota.h \ src/core/lib/resource_quota/periodic_update.cc \ diff --git a/tools/doxygen/Doxyfile.core.internal b/tools/doxygen/Doxyfile.core.internal index 9f4eb9b2777..7ef0084dfc4 100644 --- a/tools/doxygen/Doxyfile.core.internal +++ b/tools/doxygen/Doxyfile.core.internal @@ -2388,6 +2388,8 @@ src/core/lib/resource_quota/api.cc \ src/core/lib/resource_quota/api.h \ src/core/lib/resource_quota/arena.cc \ src/core/lib/resource_quota/arena.h \ +src/core/lib/resource_quota/connection_quota.cc \ +src/core/lib/resource_quota/connection_quota.h \ src/core/lib/resource_quota/memory_quota.cc \ src/core/lib/resource_quota/memory_quota.h \ src/core/lib/resource_quota/periodic_update.cc \