GCC Code Coverage Report

src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
80 / 102
78.4%
Branches:
50 / 91
54.9%
Functions:
14 / 18
77.8%
ps/rdma/raw_verbs_transport.h
Line Branch Exec Source
1 #pragma once
2
3 #include <cstddef>
4 #include <cstdint>
5 #include <memory>
6 #include <stdexcept>
7 #include <string>
8 #include <vector>
9
10 #include <infiniband/verbs.h>
11
12 #include "base/array.h"
13 #include "ps/rdma/global_address.h"
14 #include "ps/rdma/rdma_deployment.h"
15
16 namespace petps {
17
18 class RdmaControlPlaneClient;
19
20 struct RawVerbsConfig {
21 int global_id = 0;
22 int local_lane = 0;
23 int remote_lane = 0;
24 int only_node_id = -1; // Optional single peer node filter.
25 int num_servers = 1;
26 int num_clients = 1;
27 std::string device_name;
28 std::uint8_t port_num = 0;
29 recstore::RdmaFabricMode fabric_mode = recstore::RdmaFabricMode::kIb;
30 int gid_index = -1;
31 std::uint8_t hop_limit = 0;
32 std::uint8_t traffic_class = 0;
33 std::uint32_t flow_label = 0;
34 std::string deployment_id;
35 std::uint64_t deployment_epoch = 0;
36 std::uint32_t protocol_version = recstore::kRdmaDeploymentProtocolVersion;
37 std::string configuration_digest;
38 std::string fabric_digest;
39 std::uint32_t max_inline_data = 0;
40 bool connect_to_servers = true;
41 bool connect_to_clients = true;
42 std::size_t local_region_bytes = 128 * 1024 * 1024;
43 std::uint64_t local_base_addr = 0;
44 // Global offset of local_base_addr inside this lane's logical address
45 // window. A lane may allocate only its own slice while still publishing
46 // globally-addressed offsets to peers, which subtracts this origin again.
47 std::uint64_t region_origin_offset = 0;
48 std::string control_plane_host = "127.0.0.1";
49 int control_plane_port = 25100;
50 int control_plane_timeout_ms = 30000;
51 std::uint64_t allocation_start_offset = 0;
52 std::uint64_t reserved_region_offset = 0;
53 std::uint64_t reserved_region_bytes = 0;
54 };
55
56 2 inline std::string RawVerbsMetaKey(
57 int publisher_node_id,
58 int publisher_lane,
59 int receiver_node_id,
60 int receiver_lane) {
61
4/8
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
 4 return "raw-verbs-meta-" + std::to_string(publisher_node_id) + "-lane-" +
62
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
 8 std::to_string(publisher_lane) + "-to-" +
63
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 8 std::to_string(receiver_node_id) + "-lane-" +
64
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 6 std::to_string(receiver_lane);
65 }
66
67 inline bool
68 16 ShouldRawVerbsConnectToNode(const RawVerbsConfig& config, int node_id) {
69
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 12 times.
 16 if (node_id == config.global_id) {
70 4 return false;
71 }
72
2/4
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 12 if (node_id < 0 || node_id >= config.num_servers + config.num_clients) {
73 ✗ return false;
74 }
75
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 12 if (config.only_node_id >= 0 && node_id != config.only_node_id) {
76 ✗ return false;
77 }
78 12 const bool is_server = node_id < config.num_servers;
79
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 6 times.
 12 return is_server ? config.connect_to_servers : config.connect_to_clients;
80 }
81
82 struct RawVerbsReservedRegion {
83 std::uint64_t offset = 0;
84 std::uint64_t bytes = 0;
85 };
86
87 class RawVerbsRegionAllocator {
88 public:
89 6 explicit RawVerbsRegionAllocator(std::uint64_t limit_bytes,
90 std::uint64_t allocation_start_offset = 0)
91 6 : limit_bytes_(limit_bytes), allocation_offset_(allocation_start_offset) {
92
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (allocation_start_offset > limit_bytes) {
93
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ throw std::runtime_error(
94 ✗ "raw verbs allocation start outside local memory");
95 }
96 6 }
97
98 6 void SetReservedRegion(RawVerbsReservedRegion reserved) {
99
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (reserved.bytes != 0) {
100 6 const std::uint64_t reserved_end = reserved.offset + reserved.bytes;
101
3/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 4 times.
 6 if (reserved_end < reserved.offset || reserved_end > limit_bytes_) {
102
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 throw std::runtime_error(
103 4 "raw verbs reserved region outside local memory");
104 }
105 }
106 4 reserved_ = reserved;
107 4 const std::uint64_t reserved_end = reserved.offset + reserved.bytes;
108
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (reserved.bytes != 0 && allocation_offset_ >= reserved.offset &&
109
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ allocation_offset_ < reserved_end) {
110 ✗ allocation_offset_ = reserved_end;
111 }
112 4 }
113
114 12 std::uint64_t Allocate(std::size_t bytes) {
115 12 const std::uint64_t aligned = Align(bytes);
116 12 std::uint64_t offset = allocation_offset_;
117 12 const std::uint64_t reserved_begin = reserved_.offset;
118 12 const std::uint64_t reserved_end = reserved_.offset + reserved_.bytes;
119
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (reserved_.bytes != 0) {
120
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 10 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 12 if (offset >= reserved_begin && offset < reserved_end) {
121 2 offset = reserved_end;
122
2/4
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
 10 } else if (offset < reserved_begin && offset + aligned > reserved_begin) {
123 ✗ offset = reserved_end;
124 }
125 }
126
2/4
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 12 if (offset + aligned < offset || offset + aligned > limit_bytes_) {
127
0/2
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 ✗ throw std::runtime_error("raw verbs registered region exhausted");
128 }
129 12 allocation_offset_ = offset + aligned;
130 12 return offset;
131 }
132
133 2 std::uint64_t Checkpoint() const { return allocation_offset_; }
134
135 2 void Restore(std::uint64_t checkpoint) {
136
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (checkpoint > limit_bytes_) {
137
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ throw std::runtime_error(
138 ✗ "raw verbs allocation checkpoint outside local memory");
139 }
140 2 allocation_offset_ = checkpoint;
141 2 }
142
143 6 std::uint64_t current_offset() const { return allocation_offset_; }
144
145 private:
146 12 static std::uint64_t Align(std::size_t bytes) {
147 12 return (static_cast<std::uint64_t>(bytes) + 63) & ~std::uint64_t{63};
148 }
149
150 std::uint64_t limit_bytes_ = 0;
151 std::uint64_t allocation_offset_ = 0;
152 RawVerbsReservedRegion reserved_{};
153 };
154
155 class RawVerbsRegionAllocatorScope {
156 public:
157 2 explicit RawVerbsRegionAllocatorScope(RawVerbsRegionAllocator* allocator)
158 2 : allocator_(allocator),
159
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 checkpoint_(allocator != nullptr ? allocator->Checkpoint() : 0) {}
160
161 2 ~RawVerbsRegionAllocatorScope() {
162
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (allocator_ != nullptr) {
163 2 allocator_->Restore(checkpoint_);
164 }
165 2 }
166
167 RawVerbsRegionAllocatorScope(const RawVerbsRegionAllocatorScope&) = delete;
168 RawVerbsRegionAllocatorScope&
169 operator=(const RawVerbsRegionAllocatorScope&) = delete;
170
171 private:
172 RawVerbsRegionAllocator* allocator_ = nullptr;
173 std::uint64_t checkpoint_ = 0;
174 };
175
176 struct RawVerbsRemoteMemory {
177 std::uint16_t node_id = 0;
178 std::uint64_t base_addr = 0;
179 std::uint32_t rkey = 0;
180 };
181
182 struct RawVerbsCompletion {
183 std::uint64_t wr_id = 0;
184 std::uint32_t imm_data = 0;
185 bool has_imm = false;
186 ibv_wc_opcode opcode = IBV_WC_SEND;
187 };
188
189 // Classification of a non-successful CQ completion (device-layer failure
190 // semantics). Only the peer-side classes say the remote node stopped
191 // answering; local classes mean the fault is ours and a QP rebuild would not
192 // help.
193 enum class RawVerbsErrorClass : std::uint8_t {
194 kPeerUnreachable = 0, // Retry / RNR retry exhausted: no ACK from the peer.
195 kPeerRejected = 1, // The peer answered but rejected the operation.
196 kLocalError = 2, // Local length/protection/QP-operation fault.
197 kFlushed = 3, // The WR was flushed from a QP that was already dead.
198 kOther = 4,
199 };
200
201 10 inline const char* RawVerbsErrorClassToString(RawVerbsErrorClass error_class) {
202
5/6
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 10 switch (error_class) {
203 2 case RawVerbsErrorClass::kPeerUnreachable:
204 2 return "peer_unreachable";
205 2 case RawVerbsErrorClass::kPeerRejected:
206 2 return "peer_rejected";
207 2 case RawVerbsErrorClass::kLocalError:
208 2 return "local_error";
209 2 case RawVerbsErrorClass::kFlushed:
210 2 return "flushed";
211 2 case RawVerbsErrorClass::kOther:
212 2 return "other";
213 }
214 ✗ return "other";
215 }
216
217 inline RawVerbsErrorClass
218 26 ClassifyRawVerbsCompletionStatus(ibv_wc_status status) {
219
5/5
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 10 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 4 times.
 26 switch (status) {
220 4 case IBV_WC_RETRY_EXC_ERR:
221 case IBV_WC_RNR_RETRY_EXC_ERR:
222 4 return RawVerbsErrorClass::kPeerUnreachable;
223 6 case IBV_WC_REM_INV_REQ_ERR:
224 case IBV_WC_REM_ACCESS_ERR:
225 case IBV_WC_REM_OP_ERR:
226 6 return RawVerbsErrorClass::kPeerRejected;
227 10 case IBV_WC_LOC_LEN_ERR:
228 case IBV_WC_LOC_QP_OP_ERR:
229 case IBV_WC_LOC_PROT_ERR:
230 case IBV_WC_LOC_ACCESS_ERR:
231 case IBV_WC_LOC_RDD_VIOL_ERR:
232 10 return RawVerbsErrorClass::kLocalError;
233 2 case IBV_WC_WR_FLUSH_ERR:
234 2 return RawVerbsErrorClass::kFlushed;
235 4 default:
236 4 return RawVerbsErrorClass::kOther;
237 }
238 }
239
240 // True when the completion error means the shard cannot be reached on this
241 // lane: the peer stopped answering, or the lane was already flushed dead.
242 26 inline bool IsRawVerbsUnreachableError(RawVerbsErrorClass error_class) {
243
4/4
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 20 times.
 26 return error_class == RawVerbsErrorClass::kPeerUnreachable ||
244 26 error_class == RawVerbsErrorClass::kFlushed;
245 }
246
247 struct RawVerbsCompletionError {
248 ibv_wc_status status = IBV_WC_SUCCESS;
249 std::uint64_t wr_id = 0;
250 ibv_wc_opcode opcode = IBV_WC_SEND;
251 RawVerbsErrorClass error_class = RawVerbsErrorClass::kOther;
252 };
253
254 // Human-readable form, e.g. "raw verbs CQ error: retry exceeded
255 // class=peer_unreachable wr_id=0 opcode=0".
256 std::string
257 DescribeRawVerbsCompletionError(const RawVerbsCompletionError& error);
258
259 class RawVerbsCompletionException : public std::runtime_error {
260 public:
261 explicit RawVerbsCompletionException(const RawVerbsCompletionError& error,
262 const std::string& context = "");
263 ✗ const RawVerbsCompletionError& error() const { return error_; }
264
265 private:
266 RawVerbsCompletionError error_;
267 };
268
269 enum class RawVerbsPollOutcome : std::uint8_t {
270 kEmpty, // No completion available.
271 kCompletion, // A successful completion was returned.
272 kError, // A completion error was returned instead of thrown.
273 };
274
275 struct RawVerbsSge {
276 const void* data = nullptr;
277 std::size_t bytes = 0;
278 };
279
280 inline constexpr int kRawVerbsPollBatchSize = 16;
281
282 class RawVerbsCompletionBatchCursor {
283 public:
284 ✗ bool HasCachedCompletion() const { return current_ < size_; }
285
286 ✗ void Reset(ibv_wc* entries, int size) {
287 ✗ entries_ = entries;
288 ✗ size_ = size;
289 ✗ current_ = 0;
290 ✗ }
291
292 ✗ ibv_wc* TakeCachedCompletion() {
293
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ if (!HasCachedCompletion()) {
294 ✗ return nullptr;
295 }
296 ✗ return entries_ + current_++;
297 }
298
299 private:
300 ibv_wc* entries_ = nullptr;
301 int size_ = 0;
302 int current_ = 0;
303 };
304
305 struct RawVerbsNodeMeta {
306 std::string deployment_id;
307 std::string configuration_digest;
308 std::string fabric_digest;
309 std::uint16_t node_id = 0;
310 std::int32_t logical_id = -1;
311 std::uint16_t lid = 0;
312 std::uint32_t qpn = 0;
313 std::uint32_t psn = 3185;
314 std::uint32_t rkey = 0;
315 std::uint64_t base_addr = 0;
316 std::uint8_t gid[16] = {};
317 std::uint32_t protocol_version = recstore::kRdmaDeploymentProtocolVersion;
318 std::uint64_t deployment_epoch = 0;
319 std::uint8_t node_role = 0;
320 std::uint8_t port_num = 0;
321 std::uint8_t link_layer = 0;
322 std::int32_t gid_index = -1;
323 std::uint8_t active_mtu = 0;
324 std::uint8_t fabric_mode = 0;
325 // This lane's rebuild generation. Bumped before the QP is republished so the
326 // peer can observe "this side rebuilt" and re-pair the QP.
327 std::uint64_t lane_generation = 0;
328 };
329
330 class RawVerbsTransport {
331 public:
332 explicit RawVerbsTransport(const RawVerbsConfig& config);
333 ~RawVerbsTransport();
334
335 RawVerbsTransport(const RawVerbsTransport&) = delete;
336 RawVerbsTransport& operator=(const RawVerbsTransport&) = delete;
337
338 void RegisterMemoryRegion(void* base, std::size_t bytes);
339 void* AllocateRegistered(std::size_t bytes);
340 std::uint64_t SaveAllocationState() const;
341 void RestoreAllocationState(std::uint64_t checkpoint);
342 GlobalAddress LocalAddress(void* ptr) const;
343 void* LocalPointer(GlobalAddress address) const;
344
345 void Publish();
346 void Connect();
347 void PublishAndConnect();
348 RawVerbsNodeMeta LocalMeta() const;
349
350 void Write(const void* local,
351 GlobalAddress remote,
352 std::size_t bytes,
353 std::uint64_t wr_id,
354 bool signaled);
355 void WriteSg(base::ConstArray<RawVerbsSge> sges,
356 GlobalAddress remote,
357 std::uint64_t wr_id,
358 bool signaled);
359 void WriteWithImm(const void* local,
360 GlobalAddress remote,
361 std::size_t bytes,
362 std::uint32_t imm_data,
363 std::uint64_t wr_id,
364 bool signaled);
365 void Read(void* local,
366 GlobalAddress remote,
367 std::size_t bytes,
368 std::uint64_t wr_id,
369 bool signaled);
370 void SendDoorbell(
371 std::uint16_t node_id, std::uint32_t imm_data, std::uint64_t wr_id);
372 bool Poll(RawVerbsCompletion* completion, int timeout_ms);
373 // Non-blocking variant of Poll() for wait loops that must notice a broken QP
374 // without waiting for the next drain: a completion error is returned through
375 // `error` (which must not be null) instead of being thrown.
376 RawVerbsPollOutcome
377 TryPoll(RawVerbsCompletion* completion, RawVerbsCompletionError* error);
378 // Rebuilds the QP for one peer node after a device-layer failure, reusing the
379 // PD / MR / registered memory and the peer metadata captured by Connect():
380 // reset -> init -> RTR -> RTS. The reset flushes every outstanding WR on the
381 // QP, so those completions are dropped from the CQ afterwards and the receive
382 // queue is re-armed because reset flushes it too. Only peer-side failures
383 // (retry exhaustion / flush) are worth rebuilding; callers must decide that
384 // from RawVerbsErrorClass. The QP is a shared resource, so a rebuild is only
385 // half of recovery: the peer must rebuild its side too, which the caller
386 // coordinates through the control plane. Must not run concurrently with
387 // submits or polls on the same lane, and only on a QP whose CQ carries no
388 // other poster, because the post-reset drain would discard that peer's
389 // completions (see RcShardServerTransport::RebuildClientLaneIfStale).
390 void RebuildQp(std::uint16_t node_id);
391 // Lane generation handshake. The lane generation is the local side's rebuild
392 // counter toward `node_id`; `AdvanceLaneGeneration` bumps it and republishes
393 // the metadata so the peer can observe the change. `LaneGeneration` reads
394 // that counter back; the peer's counter is read from the meta returned by
395 // RefreshPeerMeta().
396 std::uint64_t AdvanceLaneGeneration(std::uint16_t node_id);
397 std::uint64_t LaneGeneration(std::uint16_t node_id) const;
398 // Re-reads `node_id`'s published metadata (GetMeta), validates it, and stores
399 // it for a subsequent RebuildQp(). Used to learn a peer's lane generation and
400 // to pick up addressing changes if the peer process restarted. `timeout_ms`
401 // overrides the configured control-plane deadline; callers on a data-plane
402 // thread pass a short bound so a sick control plane cannot stall it.
403 RawVerbsNodeMeta RefreshPeerMeta(std::uint16_t node_id, int timeout_ms = -1);
404 std::uint32_t max_inline_data(std::uint16_t node_id) const;
405
406 private:
407 struct Impl;
408 ibv_mr* FindLocalMr(const void* ptr, std::size_t bytes) const;
409 // Lazily opens (and caches) the one control-plane client this transport owns.
410 // Publish()/Connect()/RefreshPeerMeta() share it so a lane uses a single
411 // channel instead of a fresh connection per metadata RPC.
412 RdmaControlPlaneClient& ControlPlane();
413 std::unique_ptr<Impl> impl_;
414 };
415
416 class RawVerbsTransportAllocationScope {
417 public:
418 explicit RawVerbsTransportAllocationScope(RawVerbsTransport* transport)
419 : transport_(transport),
420 checkpoint_(
421 transport != nullptr ? transport->SaveAllocationState() : 0) {}
422
423 ~RawVerbsTransportAllocationScope() {
424 if (transport_ != nullptr) {
425 transport_->RestoreAllocationState(checkpoint_);
426 }
427 }
428
429 RawVerbsTransportAllocationScope(const RawVerbsTransportAllocationScope&) =
430 delete;
431 RawVerbsTransportAllocationScope&
432 operator=(const RawVerbsTransportAllocationScope&) = delete;
433
434 private:
435 RawVerbsTransport* transport_ = nullptr;
436 std::uint64_t checkpoint_ = 0;
437 };
438
439 } // namespace petps
440