1 // Copyright 2024 gRPC authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 #include "src/core/lib/resource_quota/connection_quota.h"
16
17 #include <grpc/support/port_platform.h>
18
19 #include <atomic>
20 #include <cstdint>
21
22 #include "absl/log/check.h"
23
24 namespace grpc_core {
25
26 ConnectionQuota::ConnectionQuota() = default;
27
SetMaxIncomingConnections(int max_incoming_connections)28 void ConnectionQuota::SetMaxIncomingConnections(int max_incoming_connections) {
29 // The maximum can only be configured once.
30 CHECK_LT(max_incoming_connections, INT_MAX);
31 CHECK(max_incoming_connections_.exchange(
32 max_incoming_connections, std::memory_order_release) == INT_MAX);
33 }
34
35 // Returns true if the incoming connection is allowed to be accepted on the
36 // server.
AllowIncomingConnection(MemoryQuotaRefPtr mem_quota,absl::string_view)37 bool ConnectionQuota::AllowIncomingConnection(MemoryQuotaRefPtr mem_quota,
38 absl::string_view /*peer*/) {
39 if (mem_quota->IsMemoryPressureHigh()) {
40 return false;
41 }
42
43 if (max_incoming_connections_.load(std::memory_order_relaxed) == INT_MAX) {
44 return true;
45 }
46
47 int curr_active_connections =
48 active_incoming_connections_.load(std::memory_order_acquire);
49 do {
50 if (curr_active_connections >=
51 max_incoming_connections_.load(std::memory_order_relaxed)) {
52 return false;
53 }
54 } while (!active_incoming_connections_.compare_exchange_weak(
55 curr_active_connections, curr_active_connections + 1,
56 std::memory_order_acq_rel, std::memory_order_relaxed));
57 return true;
58 }
59
60 // Mark connections as closed.
ReleaseConnections(int num_connections)61 void ConnectionQuota::ReleaseConnections(int num_connections) {
62 if (max_incoming_connections_.load(std::memory_order_relaxed) == INT_MAX) {
63 return;
64 }
65 CHECK(active_incoming_connections_.fetch_sub(
66 num_connections, std::memory_order_acq_rel) >= num_connections);
67 }
68
69 } // namespace grpc_core
70