ps/rdma/rc_transport.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include <atomic> | ||
| 4 | #include <cstddef> | ||
| 5 | #include <cstdint> | ||
| 6 | #include <memory> | ||
| 7 | #include <mutex> | ||
| 8 | #include <stdexcept> | ||
| 9 | #include <string> | ||
| 10 | #include <unordered_map> | ||
| 11 | #include <utility> | ||
| 12 | #include <vector> | ||
| 13 | |||
| 14 | #include "base/array.h" | ||
| 15 | #include "ps/rdma/raw_verbs_transport.h" | ||
| 16 | #include "ps/rdma/rdma_deployment.h" | ||
| 17 | #include "ps/rdma/rdma_protocol.h" | ||
| 18 | |||
| 19 | namespace petps { | ||
| 20 | |||
| 21 | struct RcTransportConfig { | ||
| 22 | int node_id = -1; // Dense deployment node id for this process. | ||
| 23 | int num_servers = 0; // Number of server nodes in the deployment. | ||
| 24 | int num_os_clients = 0; // Number of client processes in the deployment. | ||
| 25 | int shard_id = 0; // Logical shard served by this transport. | ||
| 26 | int client_id = -1; // Logical client id for response slot selection. | ||
| 27 | int num_clients = 1; // Total client count expected by the server. | ||
| 28 | int qps_per_client_per_shard = 32; // Number of lanes per client per shard. | ||
| 29 | int slots_per_qp = 1; // Logical slots multiplexed on each lane. | ||
| 30 | std::size_t request_slot_bytes = 1 << 20; // Bytes per server request slot. | ||
| 31 | std::size_t response_slot_bytes = 1 << 20; // Bytes per client response slot. | ||
| 32 | std::string control_plane_host = "127.0.0.1"; | ||
| 33 | int control_plane_port = 25100; | ||
| 34 | int control_plane_timeout_ms = 30000; | ||
| 35 | std::string namespace_token = | ||
| 36 | "default"; // Shared-memory namespace for this run. | ||
| 37 | std::string deployment_id; | ||
| 38 | std::uint64_t deployment_epoch = 0; | ||
| 39 | std::uint32_t protocol_version = recstore::kRdmaDeploymentProtocolVersion; | ||
| 40 | std::string configuration_digest; | ||
| 41 | std::string fabric_digest; | ||
| 42 | recstore::ResolvedRdmaFabric fabric; | ||
| 43 | }; | ||
| 44 | |||
| 45 | struct RcClientQpView { | ||
| 46 | int qp_index = 0; // Lane index local to the client. | ||
| 47 | int slot_in_qp = 0; // Logical slot index within the lane. | ||
| 48 | int slot_index = 0; // Global request slot index on the server. | ||
| 49 | void* request_slot = nullptr; // Base address of the request slot. | ||
| 50 | RequestDescriptor* descriptor = nullptr; | ||
| 51 | char* payload = nullptr; // Request payload region after descriptor. | ||
| 52 | CommitWord* commit = nullptr; // Commit word at end of request slot. | ||
| 53 | void* response_slot = nullptr; // Base address of the client response slot. | ||
| 54 | char* response_payload = nullptr; // Response payload region before status. | ||
| 55 | StatusWord* status = nullptr; // Status word at end of response slot. | ||
| 56 | }; | ||
| 57 | |||
| 58 | // Raised when a tracked completion (a submit-commit or response-status CQE) did | ||
| 59 | // not arrive within the wait budget: the slot's memory may still be referenced | ||
| 60 | // by the peer, so the caller must retire the lane instead of reusing the slot. | ||
| 61 | class RcTrackedCompletionTimeout : public std::runtime_error { | ||
| 62 | public: | ||
| 63 | ✗ | explicit RcTrackedCompletionTimeout(const std::string& message) | |
| 64 | ✗ | : std::runtime_error(message) {} | |
| 65 | }; | ||
| 66 | |||
| 67 | // Raised when a response payload must not be posted because its lane was | ||
| 68 | // rebuilt after the request was accepted. The client slot now belongs to a new | ||
| 69 | // generation, so writing it would corrupt that generation's response; the | ||
| 70 | // caller must drop the response instead of completing it. | ||
| 71 | class RcLaneRebuiltDuringResponse : public std::runtime_error { | ||
| 72 | public: | ||
| 73 | ✗ | explicit RcLaneRebuiltDuringResponse(const std::string& message) | |
| 74 | ✗ | : std::runtime_error(message) {} | |
| 75 | }; | ||
| 76 | |||
| 77 | class RcShardClientTransport { | ||
| 78 | public: | ||
| 79 | explicit RcShardClientTransport(const RcTransportConfig& config); | ||
| 80 | ~RcShardClientTransport(); | ||
| 81 | |||
| 82 | RcShardClientTransport(const RcShardClientTransport&) = delete; | ||
| 83 | RcShardClientTransport& operator=(const RcShardClientTransport&) = delete; | ||
| 84 | |||
| 85 | RcClientQpView OpenQp(int qp_index); | ||
| 86 | RcClientQpView OpenSlot(int qp_index, int slot_in_qp); | ||
| 87 | void SubmitRequest(const RcClientQpView& view, | ||
| 88 | const RequestDescriptor& descriptor, | ||
| 89 | const void* payload, | ||
| 90 | std::size_t payload_bytes); | ||
| 91 | void ClearRequestSlot(const RcClientQpView& view); | ||
| 92 | // Drains at most one completion from a lane CQ without blocking. Successful | ||
| 93 | // completions are folded into the lane's tracked-completion bookkeeping so a | ||
| 94 | // later SubmitRequest() drain still finds them; a completion error is copied | ||
| 95 | // into `error` and reported as true so the caller can classify and retire the | ||
| 96 | // lane. Callers must serialize with the client submit path (PetPSClient holds | ||
| 97 | // its mutex around both). | ||
| 98 | bool HarvestLaneCompletion(int qp_index, RawVerbsCompletionError* error); | ||
| 99 | // Rebuilds one lane QP after a peer-side failure and completes the bilateral | ||
| 100 | // handshake: publish the bumped lane generation, then wait until the peer | ||
| 101 | // rebuilds and republishes too. Returns true once the peer acknowledged | ||
| 102 | // within --rdma_lane_rebuild_wait_ms. Callers must serialize this with the | ||
| 103 | // submit path and the CQ harvest (PetPSClient holds its mutex). Requires a | ||
| 104 | // single OS client per shard (num_os_clients == 1): the server side cannot | ||
| 105 | // reset a QP that shares its CQ, so a multi-client deployment returns false | ||
| 106 | // and keeps the terminal-dirty behaviour. | ||
| 107 | bool RebuildLane(int qp_index); | ||
| 108 | std::uint64_t LaneGeneration(int qp_index) const; | ||
| 109 | std::size_t request_slot_bytes() const { return config_.request_slot_bytes; } | ||
| 110 | std::size_t response_slot_bytes() const { | ||
| 111 | return config_.response_slot_bytes; | ||
| 112 | } | ||
| 113 | const RcTransportConfig& config() const { return config_; } | ||
| 114 | |||
| 115 | private: | ||
| 116 | struct Lane { | ||
| 117 | std::unique_ptr<RawVerbsTransport> | ||
| 118 | verbs; // RC QP and registered memory for this lane. | ||
| 119 | void* lane_base = | ||
| 120 | nullptr; // Local registered region for all slots in this lane. | ||
| 121 | std::vector<std::uint8_t> submit_completion_pending; | ||
| 122 | std::vector<std::uint8_t> submit_completion_ready; | ||
| 123 | // Number of times this lane's QP was rebuilt; the value published in the | ||
| 124 | // lane metadata so the server can observe the bump. | ||
| 125 | std::uint64_t generation = 0; | ||
| 126 | }; | ||
| 127 | |||
| 128 | Lane& LaneAt(int qp_index); | ||
| 129 | const Lane& LaneAt(int qp_index) const; | ||
| 130 | |||
| 131 | RcTransportConfig config_; | ||
| 132 | int server_node_id_ = 0; // Global node id of the target shard server. | ||
| 133 | std::vector<Lane> lanes_; | ||
| 134 | }; | ||
| 135 | |||
| 136 | class RcShardServerTransport { | ||
| 137 | public: | ||
| 138 | explicit RcShardServerTransport(const RcTransportConfig& config); | ||
| 139 | ~RcShardServerTransport(); | ||
| 140 | |||
| 141 | RcShardServerTransport(const RcShardServerTransport&) = delete; | ||
| 142 | RcShardServerTransport& operator=(const RcShardServerTransport&) = delete; | ||
| 143 | |||
| 144 | int TotalSlots() const; | ||
| 145 | void RegisterLocalMemoryRegion(void* base, std::size_t bytes); | ||
| 146 | int SlotIndex(int client_id, int qp_index, int slot_in_qp) const; | ||
| 147 | void DecodeSlotIndex( | ||
| 148 | int slot_index, int* client_id, int* qp_index, int* slot_in_qp) const; | ||
| 149 | void* RequestSlot(int slot_index) const; | ||
| 150 | RequestDescriptor* RequestDescriptorAt(int slot_index) const; | ||
| 151 | char* RequestPayloadAt(int slot_index) const; | ||
| 152 | CommitWord* RequestCommitAt(int slot_index) const; | ||
| 153 | |||
| 154 | struct ResponseView { | ||
| 155 | void* slot = nullptr; // Base address of the client response slot. | ||
| 156 | char* payload = nullptr; // Response payload region. | ||
| 157 | StatusWord* status = nullptr; // Final completion word for this response. | ||
| 158 | }; | ||
| 159 | |||
| 160 | ResponseView OpenClientResponse(int client_id, int qp_index, int slot_in_qp); | ||
| 161 | void CompleteResponse(int client_id, | ||
| 162 | int qp_index, | ||
| 163 | int slot_in_qp, | ||
| 164 | const ResponseView& response, | ||
| 165 | std::uint64_t seq); | ||
| 166 | // Current response epoch of a lane. It changes on every server-side lane | ||
| 167 | // rebuild, so a caller that captured it before posting can tell whether the | ||
| 168 | // slot it is about to write still belongs to the same client generation. | ||
| 169 | std::uint64_t ResponseEpoch(int client_id, int qp_index) const; | ||
| 170 | // Number of response payload posts refused because the lane was rebuilt after | ||
| 171 | // the request was accepted (RcLaneRebuiltDuringResponse). A non-zero value is | ||
| 172 | // direct evidence that the rebuild fence fired instead of merely that no bad | ||
| 173 | // response reached the client. | ||
| 174 | ✗ | std::uint64_t ResponseRejectedEpochCount() const { | |
| 175 | ✗ | return response_rejected_epoch_count_.load(std::memory_order_relaxed); | |
| 176 | } | ||
| 177 | // Posts response payload rows from `sges`. Throws RcLaneRebuiltDuringResponse | ||
| 178 | // when the lane was rebuilt after `expected_epoch` was captured: the payload | ||
| 179 | // would land in a slot the client already released (and whose pool may have | ||
| 180 | // handed it to a new request). Get-payload workers pass the epoch they | ||
| 181 | // captured when the request was accepted. | ||
| 182 | void WriteResponsePayloadSg( | ||
| 183 | int client_id, | ||
| 184 | int qp_index, | ||
| 185 | int slot_in_qp, | ||
| 186 | std::uint64_t expected_epoch, | ||
| 187 | base::ConstArray<RawVerbsSge> sges, | ||
| 188 | std::uint64_t response_offset, | ||
| 189 | std::uint64_t bytes); | ||
| 190 | void CompleteResponseStatusOnly( | ||
| 191 | int client_id, | ||
| 192 | int qp_index, | ||
| 193 | int slot_in_qp, | ||
| 194 | const ResponseView& response, | ||
| 195 | std::uint64_t seq); | ||
| 196 | // For each client served by `qp_index`, rebuild the server-side QP when the | ||
| 197 | // client published a new lane generation (i.e. the client rebuilt its side). | ||
| 198 | // Must be called from a thread that owns those lanes (the polling thread that | ||
| 199 | // scans `qp_index`). Rebuilding requires the lane CQ to carry a single client | ||
| 200 | // QP (num_os_clients == 1); otherwise it is skipped, because a shared CQ | ||
| 201 | // cannot be drained safely. Returns the number of lanes rebuilt. | ||
| 202 | int PollClientRebuildRequests(int qp_index); | ||
| 203 | ✗ | const RcTransportConfig& config() const { return config_; } | |
| 204 | |||
| 205 | private: | ||
| 206 | struct Lane { | ||
| 207 | std::unique_ptr<RawVerbsTransport> | ||
| 208 | verbs; // RC QP and registered memory for this lane. | ||
| 209 | void* request_slots = nullptr; // Local registered slots for all clients. | ||
| 210 | std::vector<void*> response_staging; // Per-client-per-slot registered | ||
| 211 | // response staging slots. | ||
| 212 | std::vector<std::uint8_t> response_completion_pending; | ||
| 213 | std::vector<std::uint8_t> response_completion_ready; | ||
| 214 | // Last client lane generation observed per logical client id. A change | ||
| 215 | // means the client rebuilt its side, so this side must rebuild too. | ||
| 216 | std::vector<std::uint64_t> observed_client_generation; | ||
| 217 | // Serializes QP reset/rebuild against payload posts from the get-payload | ||
| 218 | // workers. A post either completes before the rebuild (its WR is flushed by | ||
| 219 | // the reset) or observes the bumped epoch afterwards and refuses. | ||
| 220 | std::mutex rebuild_mu; | ||
| 221 | // Bumped on every server-side rebuild. Payload posts carry the value they | ||
| 222 | // captured when the request was accepted. | ||
| 223 | std::atomic<std::uint64_t> response_epoch{0}; | ||
| 224 | }; | ||
| 225 | |||
| 226 | Lane& LaneAt(int client_id, int qp_index); | ||
| 227 | const Lane& LaneAt(int client_id, int qp_index) const; | ||
| 228 | bool RebuildClientLaneIfStale(int client_id, int qp_index); | ||
| 229 | |||
| 230 | RcTransportConfig config_; // Transport sizing and namespace config. | ||
| 231 | std::vector<Lane> lanes_; // One verbs RC lane per qp_index. | ||
| 232 | // Fence observability: incremented by WriteResponsePayloadSg() when it | ||
| 233 | // refuses to post into a rebuilt lane. | ||
| 234 | std::atomic<std::uint64_t> response_rejected_epoch_count_{0}; | ||
| 235 | }; | ||
| 236 | |||
| 237 | } // namespace petps | ||
| 238 |