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 <grpc/support/port_platform.h>
16
17 #include "src/core/lib/resource_quota/connection_quota.h"
18
19 #include <atomic>
20 #include <cstdint>
21
22 #include <grpc/support/log.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 GPR_ASSERT(max_incoming_connections < INT_MAX);
31 GPR_ASSERT(max_incoming_connections_.exchange(max_incoming_connections,
32 std::memory_order_release) ==
33 INT_MAX);
34 }
35
36 // Returns true if the incoming connection is allowed to be accepted on the
37 // server.
AllowIncomingConnection(MemoryQuotaRefPtr mem_quota,absl::string_view)38 bool ConnectionQuota::AllowIncomingConnection(MemoryQuotaRefPtr mem_quota,
39 absl::string_view /*peer*/) {
40 if (mem_quota->IsMemoryPressureHigh()) {
41 return false;
42 }
43
44 if (max_incoming_connections_.load(std::memory_order_relaxed) == INT_MAX) {
45 return true;
46 }
47
48 int curr_active_connections =
49 active_incoming_connections_.load(std::memory_order_acquire);
50 do {
51 if (curr_active_connections >=
52 max_incoming_connections_.load(std::memory_order_relaxed)) {
53 return false;
54 }
55 } while (!active_incoming_connections_.compare_exchange_weak(
56 curr_active_connections, curr_active_connections + 1,
57 std::memory_order_acq_rel, std::memory_order_relaxed));
58 return true;
59 }
60
61 // Mark connections as closed.
ReleaseConnections(int num_connections)62 void ConnectionQuota::ReleaseConnections(int num_connections) {
63 if (max_incoming_connections_.load(std::memory_order_relaxed) == INT_MAX) {
64 return;
65 }
66 GPR_ASSERT(active_incoming_connections_.fetch_sub(
67 num_connections, std::memory_order_acq_rel) >=
68 num_connections);
69 }
70
71 } // namespace grpc_core
72