GCC Code Coverage Report

src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
57 / 66
86.4%
Branches:
51 / 78
65.4%
Functions:
9 / 10
90.0%
ps/rdma/control_plane.h
Line Branch Exec Source
1 #pragma once
2
3 #include <atomic>
4 #include <chrono>
5 #include <condition_variable>
6 #include <cstdint>
7 #include <memory>
8 #include <mutex>
9 #include <string>
10 #include <unordered_map>
11
12 #include "ps/rdma/raw_verbs_transport.h"
13
14 namespace grpc {
15 class Server;
16 }
17
18 namespace petps {
19
20 class RdmaControlPlaneService;
21
22 struct RdmaControlPlaneEndpoint {
23 std::string host = "127.0.0.1";
24 int port = 25100;
25 int timeout_ms = 30000;
26 std::string deployment_id;
27 std::uint64_t deployment_epoch = 0;
28 std::uint32_t protocol_version = recstore::kRdmaDeploymentProtocolVersion;
29 std::string configuration_digest;
30 std::string fabric_digest;
31 // Optional on client-only endpoints; the server uses it to bound READY ids.
32 int num_servers = 0;
33 // Control-plane health evaluation thresholds (server side only).
34 int health_stale_ms = 3000;
35 int health_dead_ms = 9000;
36 };
37
38 inline constexpr std::uint32_t kRawVerbsMetadataMagic = 0x52444d41;
39 inline constexpr std::size_t kMaxRawVerbsMetadataBytes = 4096;
40
41 // Per-shard liveness reported by GetServerHealth: heartbeats feed `last_seen`
42 // and a query evaluates it against the stale/dead thresholds.
43 enum class RdmaServerHealth : std::uint8_t {
44 kServing = 0,
45 kSuspect = 1,
46 kDead = 2,
47 };
48
49 ✗ inline const char* RdmaServerHealthToString(RdmaServerHealth health) {
50
0/4
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 ✗ switch (health) {
51 ✗ case RdmaServerHealth::kServing:
52 ✗ return "serving";
53 ✗ case RdmaServerHealth::kSuspect:
54 ✗ return "suspect";
55 ✗ case RdmaServerHealth::kDead:
56 ✗ return "dead";
57 }
58 ✗ return "unknown";
59 }
60
61 // Pure staleness evaluation, shared by the monitor thread and unit tests. A
62 // record that has never been touched (default last_seen) is treated as serving
63 // so a freshly created record is never reported dead.
64 44 inline RdmaServerHealth EvaluateServerHealth(
65 std::chrono::steady_clock::time_point last_seen,
66 std::chrono::steady_clock::time_point now,
67 int stale_ms,
68 int dead_ms) {
69
3/4
✓ Branch 2 taken 44 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 42 times.
 44 if (last_seen == std::chrono::steady_clock::time_point{}) {
70 2 return RdmaServerHealth::kServing;
71 }
72 const auto age_ms =
73
1/2
✓ Branch 2 taken 42 times.
✗ Branch 3 not taken.
 42 std::chrono::duration_cast<std::chrono::milliseconds>(now - last_seen)
74 42 .count();
75
3/4
✓ Branch 0 taken 42 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 14 times.
✓ Branch 3 taken 28 times.
 42 if (dead_ms > 0 && age_ms >= dead_ms) {
76 14 return RdmaServerHealth::kDead;
77 }
78
3/4
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✓ Branch 3 taken 22 times.
 28 if (stale_ms > 0 && age_ms >= stale_ms) {
79 6 return RdmaServerHealth::kSuspect;
80 }
81 22 return RdmaServerHealth::kServing;
82 }
83
84 // Snapshot returned by the GetServerHealth RPC and the matching server-side
85 // accessor. `known=false` means shard 0 has no READY record for that shard, so
86 // the data path must stay on its own budgets instead of failing the request.
87 struct RdmaServerHealthSnapshot {
88 bool known = false;
89 RdmaServerHealth health = RdmaServerHealth::kServing;
90 std::int64_t last_seen_age_ms = -1;
91 };
92
93 // Slow-path health probe policy. A client that has been waiting longer than the
94 // probe threshold asks the control plane about its shard; the query is
95 // throttled so many waiters share at most one RPC per probe interval.
96 // True when the control plane positively judged the shard dead. An inconclusive
97 // query never fails a request: an unreachable control plane and a shard without
98 // a READY record both report `known=false`, and then the request stays on its
99 // data-path wait budgets instead of producing a false verdict.
100 10 inline bool IsShardUnreachable(const RdmaServerHealthSnapshot& snapshot) {
101
4/4
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 4 times.
 10 return snapshot.known && snapshot.health == RdmaServerHealth::kDead;
102 }
103
104 // Throttles slow-path probes: at most one in-flight query, and at most one
105 // query per `ttl_ms`. A caller that cannot get the gate reads the cached
106 // verdict instead of blocking (or skips its own probe when no verdict exists
107 // yet); the probe timestamp is taken before the RPC so a slow or failing
108 // control plane is throttled as well.
109 //
110 // The verdict is per shard, and a client owns one shard, so every lane can
111 // share it: without the cache each lane would have to burn its own probe slot
112 // and a dead shard could keep sibling lanes serving for N-1 extra intervals.
113 class RdmaSlowPathProbeGate {
114 public:
115 // Returns true when the caller owns the probe and must call EndProbe().
116 12 bool TryBeginProbe(std::chrono::steady_clock::time_point now, int ttl_ms) {
117
1/2
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
 12 std::lock_guard<std::mutex> guard(mu_);
118
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 8 times.
 12 if (probing_) {
119 4 return false;
120 }
121
5/6
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✓ Branch 5 taken 2 times.
✓ Branch 6 taken 4 times.
✓ Branch 7 taken 2 times.
 12 if (last_probe_ != std::chrono::steady_clock::time_point{} && ttl_ms > 0 &&
122
6/8
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 5 taken 4 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 2 times.
✓ Branch 9 taken 2 times.
✓ Branch 10 taken 6 times.
 12 now - last_probe_ < std::chrono::milliseconds(ttl_ms)) {
123 2 return false;
124 }
125 6 probing_ = true;
126 6 last_probe_ = now;
127 6 return true;
128 12 }
129
130 6 void EndProbe() {
131
1/2
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
 6 std::lock_guard<std::mutex> guard(mu_);
132 6 probing_ = false;
133 6 }
134
135 // Returns true when a verdict newer than `ttl_ms` is cached. A ttl of 0
136 // disables the freshness window.
137 12 bool TryReadVerdict(std::chrono::steady_clock::time_point now,
138 int ttl_ms,
139 RdmaServerHealthSnapshot* verdict) const {
140
1/2
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
 12 std::lock_guard<std::mutex> guard(mu_);
141
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 8 times.
 12 if (!has_verdict_) {
142 4 return false;
143 }
144
8/10
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 6 times.
✗ Branch 8 not taken.
✓ Branch 9 taken 2 times.
✓ Branch 10 taken 4 times.
✓ Branch 11 taken 2 times.
✓ Branch 12 taken 6 times.
 8 if (ttl_ms > 0 && now - verdict_at_ >= std::chrono::milliseconds(ttl_ms)) {
145 2 return false;
146 }
147 6 *verdict = verdict_;
148 6 return true;
149 12 }
150
151 2 void StoreVerdict(std::chrono::steady_clock::time_point now,
152 const RdmaServerHealthSnapshot& verdict) {
153
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 std::lock_guard<std::mutex> guard(mu_);
154 2 verdict_ = verdict;
155 2 verdict_at_ = now;
156 2 has_verdict_ = true;
157 2 }
158
159 // Drops the cached verdict. Called when a lane rebuild proves the shard is
160 // reachable again, so a cached `dead` cannot re-retire the fresh lane.
161 2 void InvalidateVerdict() {
162
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 std::lock_guard<std::mutex> guard(mu_);
163 2 has_verdict_ = false;
164 2 }
165
166 private:
167 mutable std::mutex mu_;
168 bool probing_ = false;
169 std::chrono::steady_clock::time_point last_probe_{};
170 RdmaServerHealthSnapshot verdict_{};
171 std::chrono::steady_clock::time_point verdict_at_{};
172 bool has_verdict_ = false;
173 };
174
175 std::string EncodeRawVerbsNodeMeta(const RawVerbsNodeMeta& meta);
176 RawVerbsNodeMeta DecodeRawVerbsNodeMeta(const std::string& payload);
177 void ValidateRawVerbsPeerMeta(const RawVerbsConfig& local,
178 int expected_node_id,
179 const RawVerbsNodeMeta& remote);
180
181 // Opaque per-client gRPC channel/stub cache; defined in control_plane.cc. Kept
182 // out of the header so this file does not pull in the generated stubs.
183 struct RdmaControlPlaneChannelCache;
184
185 class RdmaControlPlaneClient {
186 public:
187 explicit RdmaControlPlaneClient(RdmaControlPlaneEndpoint endpoint);
188 // Out of line: `channels_` is an incomplete type here.
189 ~RdmaControlPlaneClient();
190
191 void PublishMeta(int publisher_node_id,
192 int publisher_lane,
193 int receiver_node_id,
194 int receiver_lane,
195 const RawVerbsNodeMeta& meta) const;
196 RawVerbsNodeMeta
197 GetMeta(int publisher_node_id,
198 int publisher_lane,
199 int receiver_node_id,
200 int receiver_lane,
201 int timeout_ms = -1) const;
202 void PublishServerReady(int server_id) const;
203 void PublishServerDraining(int server_id) const;
204 // Lightweight liveness heartbeat; the server records it as `last_seen`.
205 void SendHeartbeat(int server_id) const;
206 // Queries shard health from the shard 0 control plane. Never throws: RPC
207 // failures (and shards without a READY record) come back as `known=false`, so
208 // the client slow path falls back to its wait budgets instead of failing.
209 RdmaServerHealthSnapshot
210 GetServerHealth(int server_id, int timeout_ms = -1) const;
211 void WaitServer(int server_id, int timeout_ms = -1) const;
212 void WaitServerReady(int num_servers, int timeout_ms = -1) const;
213
214 private:
215 RdmaControlPlaneEndpoint endpoint_;
216 std::unique_ptr<RdmaControlPlaneChannelCache> channels_;
217 };
218
219 class RdmaControlPlaneServer {
220 public:
221 explicit RdmaControlPlaneServer(RdmaControlPlaneEndpoint endpoint);
222 ~RdmaControlPlaneServer();
223
224 void Start();
225 void Stop();
226
227 // Fresh snapshot for the GetServerHealth RPC: `known=false` when there is no
228 // READY record, otherwise the health evaluated at query time plus the
229 // last-seen age.
230 RdmaServerHealthSnapshot ServerHealthSnapshot(int server_id) const;
231
232 private:
233 friend class RdmaControlPlaneService;
234
235 struct MetaKey {
236 std::string deployment_id;
237 std::uint64_t deployment_epoch = 0;
238 int publisher_node_id = 0;
239 int publisher_lane = 0;
240 int receiver_node_id = 0;
241 int receiver_lane = 0;
242
243 4 bool operator==(const MetaKey& other) const {
244 4 return deployment_id == other.deployment_id &&
245
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 deployment_epoch == other.deployment_epoch &&
246
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 publisher_node_id == other.publisher_node_id &&
247
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 publisher_lane == other.publisher_lane &&
248
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 12 receiver_node_id == other.receiver_node_id &&
249
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 8 receiver_lane == other.receiver_lane;
250 }
251 };
252
253 struct MetaKeyHash {
254 std::size_t operator()(const MetaKey& key) const;
255 };
256
257 struct ReadyKey {
258 std::string deployment_id;
259 std::uint64_t deployment_epoch = 0;
260 int server_id = 0;
261
262 66 bool operator==(const ReadyKey& other) const {
263 66 return deployment_id == other.deployment_id &&
264
2/4
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 66 times.
✗ Branch 3 not taken.
 132 deployment_epoch == other.deployment_epoch &&
265
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 132 server_id == other.server_id;
266 }
267 };
268
269 struct ReadyKeyHash {
270 std::size_t operator()(const ReadyKey& key) const;
271 };
272
273 struct ReadyRecord {
274 enum class State : std::uint8_t { kServing, kDraining };
275
276 std::uint32_t protocol_version = 0;
277 std::string configuration_digest;
278 std::string fabric_digest;
279 State state = State::kServing;
280
281 // `last_seen` is refreshed by PublishServerReady and by heartbeat Probes.
282 std::chrono::steady_clock::time_point last_seen{};
283 };
284
285 int ReadyCount(const std::string& deployment_id,
286 std::uint64_t deployment_epoch,
287 std::uint32_t protocol_version,
288 const std::string& configuration_digest,
289 const std::string& fabric_digest) const;
290 bool HasReadyContractMismatch(
291 const std::string& deployment_id,
292 std::uint64_t deployment_epoch,
293 std::uint32_t protocol_version,
294 const std::string& configuration_digest,
295 const std::string& fabric_digest) const;
296 bool HasDrainingServer(const std::string& deployment_id,
297 std::uint64_t deployment_epoch,
298 int num_servers) const;
299
300 RdmaControlPlaneEndpoint endpoint_;
301 mutable std::mutex mu_;
302 std::condition_variable cv_;
303 std::unordered_map<MetaKey, RawVerbsNodeMeta, MetaKeyHash> metadata_;
304 std::unordered_map<ReadyKey, ReadyRecord, ReadyKeyHash> ready_servers_;
305 std::atomic<bool> stop_requested_{false};
306 std::unique_ptr<RdmaControlPlaneService> service_;
307 std::unique_ptr<grpc::Server> server_;
308
309 // Records an accepted heartbeat. The caller has already validated the
310 // heartbeat identity against this endpoint.
311 void RecordHeartbeat(int server_id);
312 };
313
314 } // namespace petps
315