Branch data Line data Source code
1 : : /*
2 : : * Copyright (C) 2026 by Thun Lu. All rights reserved.
3 : : * Author: Thun Lu <thun.lu@zohomail.cn>
4 : : * Repo: https://github.com/thun-res/vlink
5 : : * _ __ __ _ __
6 : : * | | / / / / (_) ____ / /__
7 : : * | | / / / / / / / __ \ / //_/
8 : : * | |/ / / /___ / / / / / / / ,<
9 : : * |___/ /_____/ /_/ /_/ /_/ /_/|_|
10 : : *
11 : : * Licensed under the Apache License, Version 2.0 (the "License");
12 : : * you may not use this file except in compliance with the License.
13 : : * You may obtain a copy of the License at
14 : : *
15 : : * http://www.apache.org/licenses/LICENSE-2.0
16 : : *
17 : : * Unless required by applicable law or agreed to in writing, software
18 : : * distributed under the License is distributed on an "AS IS" BASIS,
19 : : * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20 : : * See the License for the specific language governing permissions and
21 : : * limitations under the License.
22 : : */
23 : :
24 : : #include "./zerocopy/fast_buffer.h"
25 : :
26 : : #include <atomic>
27 : : #include <cstring>
28 : : #include <memory>
29 : : #include <new>
30 : : #include <utility>
31 : :
32 : : #include "./zerocopy/fast_buffer_manager.h"
33 : :
34 : : namespace vlink {
35 : :
36 : : namespace zerocopy {
37 : :
38 : : // FastBufferWireMetadata
39 : : struct FastBufferWireMetadata final {
40 : : Header header;
41 : : uint64_t size{0};
42 : : uint64_t protocol{0};
43 : : FastBuffer::Storage storage{FastBuffer::kStorageShared};
44 : : FastBuffer::MemoryType memory_type{FastBuffer::kMemoryDefault};
45 : : uint8_t reserved1[6]{};
46 : : uint64_t metadata_size{0};
47 : : uint64_t reserved[2]{};
48 : : };
49 : :
50 : : // FastBuffer::MetadataBuffer
51 : : struct FastBuffer::MetadataBuffer final {
52 : : Bytes bytes;
53 : : std::atomic<size_t> references{1};
54 : : };
55 : :
56 : : static constexpr size_t kFastBufferPrefixSize = sizeof(uint32_t) * 2 + sizeof(FastBufferWireMetadata);
57 : : static constexpr size_t kFastBufferEnvelopeSize = kFastBufferPrefixSize + sizeof(uint32_t);
58 : :
59 : : // FastBuffer
60 : 51 : FastBuffer::FastBuffer() {
61 : 51 : (void)FastBufferManager::get();
62 : :
63 : : static_assert(sizeof(void*) != 8 || sizeof(FastBuffer) == 96, "Sizeof must be 96 bytes on 64-bit targets.");
64 : : static_assert(sizeof(FastBufferWireMetadata) == 88, "Sizeof must be 88 bytes.");
65 : 51 : }
66 : :
67 : 51 : FastBuffer::~FastBuffer() noexcept { clear(); }
68 : :
69 : 1 : FastBuffer::FastBuffer(const FastBuffer& target) noexcept : FastBuffer() { (void)deep_copy(target); }
70 : :
71 : 1 : FastBuffer::FastBuffer(FastBuffer&& target) noexcept : FastBuffer() { move_copy(target); }
72 : :
73 : 1 : FastBuffer& FastBuffer::operator=(const FastBuffer& target) noexcept {
74 [ + - ]: 1 : if VLIKELY (this != &target) {
75 : 1 : (void)deep_copy(target);
76 : : }
77 : :
78 : 1 : return *this;
79 : : }
80 : :
81 : 1 : FastBuffer& FastBuffer::operator=(FastBuffer&& target) noexcept {
82 [ + - ]: 1 : if VLIKELY (this != &target) {
83 : 1 : move_copy(target);
84 : : }
85 : :
86 : 1 : return *this;
87 : : }
88 : :
89 : 31 : bool FastBuffer::operator<<(const Bytes& bytes) noexcept {
90 : 31 : auto& manager = FastBufferManager::get();
91 : 31 : auto* plugin = manager.get_interface();
92 : 31 : Metadata metadata;
93 : :
94 [ + - + + : 31 : if VUNLIKELY (!plugin || !read_metadata(bytes, metadata)) {
+ + ]
95 : 12 : return false;
96 : : }
97 : :
98 [ + + + + : 19 : if VUNLIKELY (metadata.storage == kStorageShared && metadata.protocol != plugin->get_protocol_id()) {
+ + ]
99 : 1 : return false;
100 : : }
101 : :
102 : 18 : std::unique_ptr<MetadataBuffer> copied_metadata(copy_metadata(metadata.buffer));
103 : :
104 [ + + - + : 18 : if VUNLIKELY (!metadata.buffer.empty() && !copied_metadata) {
- + ]
105 : 0 : return false;
106 : : }
107 : :
108 : 18 : const auto payload = Bytes::shallow_copy(bytes.data() + kFastBufferPrefixSize + metadata.buffer.size(),
109 : 36 : bytes.size() - kFastBufferEnvelopeSize - metadata.buffer.size());
110 : 18 : FastBuffer::Buffer buffer;
111 : :
112 [ + + ]: 18 : if (metadata.storage == kStorageShared) {
113 [ + + ]: 15 : if VUNLIKELY (!manager.import_handle(payload, buffer)) {
114 : 6 : return false;
115 : : }
116 : : } else {
117 [ - + ]: 3 : if VUNLIKELY (!manager.create(static_cast<size_t>(metadata.size), {}, buffer)) {
118 : 0 : return false;
119 : : }
120 : :
121 [ - + ]: 3 : if VUNLIKELY (!plugin->copy_from_host(payload, manager.native_buffer(buffer))) {
122 : 0 : manager.release(buffer);
123 : 0 : return false;
124 : : }
125 : : }
126 : :
127 [ + - - + : 12 : if VUNLIKELY (!buffer.handle || buffer.size != metadata.size) {
- + ]
128 : 0 : manager.release(buffer);
129 : 0 : return false;
130 : : }
131 : :
132 : 12 : clear();
133 : 12 : metadata_ = copied_metadata.release();
134 : 12 : buffer_ = buffer;
135 : 12 : header = metadata.header;
136 : 12 : reserved_buf_ = metadata.reserved[0];
137 : 12 : reserved_buf2_ = metadata.reserved[1];
138 : :
139 : 12 : return true;
140 : 31 : }
141 : :
142 : 19 : bool FastBuffer::operator>>(Bytes& bytes) const noexcept {
143 [ + + ]: 19 : if VUNLIKELY (!is_valid()) {
144 : 1 : return false;
145 : : }
146 : :
147 : 18 : auto& manager = FastBufferManager::get();
148 : 18 : auto* plugin = manager.get_interface();
149 : 18 : size_t payload_size = manager.get_handle_size(buffer_);
150 [ + + ]: 18 : const Storage storage = payload_size == 0 ? kStorageHost : kStorageShared;
151 : :
152 [ + + ]: 18 : if (storage == kStorageHost) {
153 : 7 : payload_size = buffer_.size;
154 : : }
155 : :
156 : 18 : const Bytes& metadata_bytes = get_metadata();
157 : :
158 : 18 : const size_t wire_size = kFastBufferEnvelopeSize + metadata_bytes.size() + payload_size;
159 : :
160 [ + + ]: 18 : if (bytes.size() != wire_size) {
161 [ + + ]: 14 : if VUNLIKELY (bytes.is_loaned()) {
162 : 1 : return false;
163 : : }
164 : :
165 : 13 : bytes = Bytes::create(wire_size);
166 : :
167 [ - + ]: 13 : if VUNLIKELY (bytes.empty()) {
168 : 0 : return false;
169 : : }
170 : : }
171 : :
172 : 17 : auto payload = Bytes::shallow_copy(bytes.data() + kFastBufferPrefixSize + metadata_bytes.size(), payload_size);
173 [ + + ]: 17 : const bool copied = storage == kStorageHost ? plugin->copy_to_host(manager.native_buffer(buffer_), payload)
174 : 11 : : manager.export_handle(buffer_, payload);
175 : :
176 [ - + ]: 17 : if VUNLIKELY (!copied) {
177 : 0 : return false;
178 : : }
179 : :
180 : 17 : FastBufferWireMetadata metadata;
181 : 17 : metadata.header = header;
182 : 17 : metadata.size = buffer_.size;
183 [ + + ]: 17 : metadata.protocol = storage == kStorageShared ? plugin->get_protocol_id() : 0;
184 : 17 : metadata.storage = storage;
185 : 17 : metadata.memory_type = buffer_.memory_type;
186 : 17 : metadata.metadata_size = metadata_bytes.size();
187 : 17 : metadata.reserved[0] = reserved_buf_;
188 : 17 : metadata.reserved[1] = reserved_buf2_;
189 : :
190 : 17 : std::memcpy(bytes.data(), &kMagicNumberBegin, sizeof(kMagicNumberBegin));
191 : 17 : std::memcpy(bytes.data() + sizeof(kMagicNumberBegin), &kWireVersion, sizeof(kWireVersion));
192 : 17 : std::memcpy(bytes.data() + sizeof(uint32_t) * 2, &metadata, sizeof(metadata));
193 : :
194 [ + + ]: 17 : if (!metadata_bytes.empty()) {
195 : 5 : std::memcpy(bytes.data() + kFastBufferPrefixSize, metadata_bytes.data(), metadata_bytes.size());
196 : : }
197 : :
198 : 17 : std::memcpy(bytes.data() + wire_size - sizeof(kMagicNumberEnd), &kMagicNumberEnd, sizeof(kMagicNumberEnd));
199 : :
200 : 17 : return true;
201 : 17 : }
202 : :
203 : 7 : bool FastBuffer::check_valid(const Bytes& bytes) noexcept {
204 : 7 : Metadata metadata;
205 : 7 : return read_metadata(bytes, metadata);
206 : 7 : }
207 : :
208 : 42 : bool FastBuffer::read_metadata(const Bytes& bytes, Metadata& metadata) noexcept {
209 [ + + ]: 42 : if VUNLIKELY (bytes.size() <= kFastBufferEnvelopeSize) {
210 : 5 : return false;
211 : : }
212 : :
213 : 37 : uint32_t magic = 0;
214 : 37 : uint32_t version = 0;
215 : 37 : std::memcpy(&magic, bytes.data(), sizeof(magic));
216 : 37 : std::memcpy(&version, bytes.data() + sizeof(magic), sizeof(version));
217 : :
218 [ + + + + : 37 : if VUNLIKELY (magic != kMagicNumberBegin || version_major(version) != version_major(kWireVersion)) {
+ + ]
219 : 2 : return false;
220 : : }
221 : :
222 : 35 : std::memcpy(&magic, bytes.end() - sizeof(magic), sizeof(magic));
223 : :
224 [ + + ]: 35 : if VUNLIKELY (magic != kMagicNumberEnd) {
225 : 2 : return false;
226 : : }
227 : :
228 : 33 : FastBufferWireMetadata decoded;
229 : 33 : std::memcpy(&decoded, bytes.data() + sizeof(uint32_t) * 2, sizeof(decoded));
230 : :
231 [ + + + + : 33 : if VUNLIKELY (decoded.size == 0 || decoded.metadata_size >= bytes.size() - kFastBufferEnvelopeSize) {
+ + ]
232 : 3 : return false;
233 : : }
234 : :
235 [ + + ]: 30 : if (decoded.storage == kStorageHost) {
236 [ + + - + : 11 : if VUNLIKELY (decoded.protocol != 0 ||
+ + ]
237 : : decoded.size != bytes.size() - kFastBufferEnvelopeSize - decoded.metadata_size) {
238 : 1 : return false;
239 : : }
240 [ + + + + : 19 : } else if VUNLIKELY (decoded.storage != kStorageShared || decoded.protocol == 0) {
+ + ]
241 : 2 : return false;
242 : : }
243 : :
244 : 27 : metadata.header = decoded.header;
245 : 27 : metadata.size = decoded.size;
246 : 27 : metadata.protocol = decoded.protocol;
247 : 27 : metadata.storage = decoded.storage;
248 : 27 : metadata.memory_type = decoded.memory_type;
249 : 27 : std::memcpy(metadata.reserved, decoded.reserved, sizeof(decoded.reserved));
250 : : metadata.buffer =
251 : 27 : Bytes::shallow_copy(bytes.data() + kFastBufferPrefixSize, static_cast<size_t>(decoded.metadata_size));
252 : 27 : return true;
253 : : }
254 : :
255 : 5 : size_t FastBuffer::get_serialized_size() const noexcept {
256 [ + + ]: 5 : if VUNLIKELY (!is_valid()) {
257 : 1 : return 0;
258 : : }
259 : :
260 : 4 : size_t payload_size = FastBufferManager::get().get_handle_size(buffer_);
261 : :
262 [ + + ]: 4 : if (payload_size == 0) {
263 : 3 : payload_size = buffer_.size;
264 : : }
265 : :
266 : 4 : const size_t metadata_size = get_metadata().size();
267 : 4 : return kFastBufferEnvelopeSize + metadata_size + payload_size;
268 : : }
269 : :
270 [ + + + - ]: 132 : bool FastBuffer::is_valid() const noexcept { return buffer_.handle != nullptr && buffer_.size != 0; }
271 : :
272 : 21 : bool FastBuffer::create(size_t size, int32_t device) noexcept {
273 : 21 : FastBuffer::Config config;
274 : 21 : config.device = device;
275 : 21 : return create(size, config);
276 : : }
277 : :
278 : 32 : bool FastBuffer::create(size_t size, const FastBuffer::Config& config) noexcept {
279 : 32 : auto& manager = FastBufferManager::get();
280 : 32 : auto* plugin = manager.get_interface();
281 : :
282 [ + + + + : 32 : if VUNLIKELY (!plugin || size == 0) {
+ + ]
283 : 4 : return false;
284 : : }
285 : :
286 : 28 : FastBuffer::Buffer buffer;
287 : :
288 [ + + ]: 28 : if VUNLIKELY (!manager.create(size, config, buffer)) {
289 : 5 : return false;
290 : : }
291 : :
292 [ + + ]: 23 : if (buffer_.handle) {
293 : 4 : manager.release(buffer_);
294 : : }
295 : :
296 : 23 : buffer_ = buffer;
297 : 23 : return true;
298 : : }
299 : :
300 : 15 : bool FastBuffer::shallow_copy(const FastBuffer& target) noexcept {
301 [ + + ]: 15 : if VUNLIKELY (this == &target) {
302 : 1 : return false;
303 : : }
304 : :
305 : 14 : FastBuffer::Buffer buffer;
306 : :
307 [ + + ]: 14 : if (target.is_valid()) {
308 : 13 : FastBufferManager::get().retain(target.buffer_, buffer);
309 : : }
310 : :
311 : 14 : clear();
312 : 14 : buffer_ = buffer;
313 : 14 : header = target.header;
314 : 14 : reserved_buf_ = target.reserved_buf_;
315 : 14 : reserved_buf2_ = target.reserved_buf2_;
316 : :
317 : 14 : metadata_ = target.metadata_;
318 : :
319 [ + + ]: 14 : if (metadata_) {
320 : 1 : metadata_->references.fetch_add(1, std::memory_order_relaxed);
321 : : }
322 : :
323 : 14 : return true;
324 : : }
325 : :
326 : 4 : bool FastBuffer::deep_copy(const FastBuffer& target) noexcept {
327 [ + + ]: 4 : if VUNLIKELY (this == &target) {
328 : 1 : return false;
329 : : }
330 : :
331 : 3 : FastBuffer::Buffer buffer;
332 : 3 : std::unique_ptr<MetadataBuffer> copied_metadata(copy_metadata(target.get_metadata()));
333 : :
334 [ + + - + : 3 : if VUNLIKELY (!target.get_metadata().empty() && !copied_metadata) {
- + ]
335 : 0 : return false;
336 : : }
337 : :
338 [ + + ]: 3 : if (target.is_valid()) {
339 : 2 : auto& manager = FastBufferManager::get();
340 : 2 : auto* plugin = manager.get_interface();
341 : 2 : FastBuffer::Config config;
342 : 2 : config.device = target.device();
343 : 2 : config.memory_type = target.memory_type();
344 : :
345 [ - + ]: 2 : if VUNLIKELY (!manager.create(target.size(), config, buffer)) {
346 : 0 : return false;
347 : : }
348 : :
349 [ - + ]: 2 : if VUNLIKELY (!plugin->copy(manager.native_buffer(target.buffer_), manager.native_buffer(buffer))) {
350 : 0 : manager.release(buffer);
351 : 0 : return false;
352 : : }
353 : : }
354 : :
355 : 3 : clear();
356 : 3 : metadata_ = copied_metadata.release();
357 : 3 : buffer_ = buffer;
358 : 3 : header = target.header;
359 : 3 : reserved_buf_ = target.reserved_buf_;
360 : 3 : reserved_buf2_ = target.reserved_buf2_;
361 : :
362 : 3 : return true;
363 : 3 : }
364 : :
365 : 3 : bool FastBuffer::move_copy(FastBuffer& target) noexcept {
366 [ + + ]: 3 : if VUNLIKELY (this == &target) {
367 : 1 : return false;
368 : : }
369 : :
370 : 2 : clear();
371 : 2 : buffer_ = target.buffer_;
372 : 2 : header = target.header;
373 : 2 : reserved_buf_ = target.reserved_buf_;
374 : 2 : reserved_buf2_ = target.reserved_buf2_;
375 : :
376 : 2 : metadata_ = std::exchange(target.metadata_, nullptr);
377 : 2 : target.buffer_ = {};
378 : 2 : target.header = {};
379 : 2 : target.reserved_buf_ = 0;
380 : 2 : target.reserved_buf2_ = 0;
381 : :
382 : 2 : return true;
383 : : }
384 : :
385 : 125 : void FastBuffer::clear() noexcept {
386 [ + + ]: 125 : if (buffer_.handle) {
387 : 46 : FastBufferManager::get().release(buffer_);
388 : : }
389 : :
390 : 125 : clear_metadata();
391 : 125 : header = {};
392 : 125 : reserved_buf_ = 0;
393 : 125 : reserved_buf2_ = 0;
394 : 125 : }
395 : :
396 : 5 : bool FastBuffer::synchronize() const noexcept {
397 : 5 : auto& manager = FastBufferManager::get();
398 [ + + + - ]: 5 : return is_valid() && manager.get_interface()->synchronize(manager.native_buffer(buffer_));
399 : : }
400 : :
401 : 29 : FastBuffer::MetadataBuffer* FastBuffer::copy_metadata(const Bytes& bytes) noexcept {
402 [ + + ]: 29 : if (bytes.empty()) {
403 : 16 : return nullptr;
404 : : }
405 : :
406 [ + - ]: 13 : auto metadata = std::unique_ptr<MetadataBuffer>(new (std::nothrow) MetadataBuffer);
407 : :
408 [ - + ]: 13 : if VUNLIKELY (!metadata) {
409 : 0 : return nullptr;
410 : : }
411 : :
412 : 13 : metadata->bytes = Bytes::deep_copy(bytes.data(), bytes.size());
413 [ - + ]: 13 : return metadata->bytes.empty() ? nullptr : metadata.release();
414 : 13 : }
415 : :
416 : 135 : void FastBuffer::clear_metadata() noexcept {
417 [ + + + + : 150 : if (metadata_ && metadata_->references.fetch_sub(1, std::memory_order_acq_rel) == 1) {
+ + ]
418 [ + - ]: 14 : delete metadata_;
419 : : }
420 : :
421 : 135 : metadata_ = nullptr;
422 : 135 : }
423 : :
424 : 8 : bool FastBuffer::set_metadata(const Bytes& bytes) noexcept {
425 : 8 : MetadataBuffer* metadata = copy_metadata(bytes);
426 : :
427 [ + + - + : 8 : if VUNLIKELY (!bytes.empty() && !metadata) {
- + ]
428 : 0 : return false;
429 : : }
430 : :
431 : 8 : clear_metadata();
432 : 8 : metadata_ = metadata;
433 : 8 : return true;
434 : : }
435 : :
436 : 4 : bool FastBuffer::set_metadata(Bytes&& bytes) noexcept {
437 [ + + + + : 4 : if (bytes.empty() || !bytes.is_owner()) {
+ + ]
438 : 2 : return set_metadata(static_cast<const Bytes&>(bytes));
439 : : }
440 : :
441 [ + - ]: 2 : auto* metadata = new (std::nothrow) MetadataBuffer;
442 : :
443 [ - + ]: 2 : if VUNLIKELY (!metadata) {
444 : 0 : return false;
445 : : }
446 : :
447 : 2 : metadata->bytes = std::move(bytes);
448 : 2 : clear_metadata();
449 : 2 : metadata_ = metadata;
450 : 2 : return true;
451 : : }
452 : :
453 : 50 : const Bytes& FastBuffer::get_metadata() const noexcept {
454 [ + + + - ]: 50 : static const Bytes kEmpty;
455 [ + + ]: 50 : return metadata_ ? metadata_->bytes : kEmpty;
456 : : }
457 : :
458 : 8 : bool FastBuffer::copy_to_host(Bytes& bytes) const noexcept {
459 : 8 : auto& manager = FastBufferManager::get();
460 [ + + ]: 8 : if VUNLIKELY (!is_valid()) {
461 : 1 : return false;
462 : : }
463 : :
464 [ + + ]: 7 : if (bytes.size() != buffer_.size) {
465 [ + + ]: 6 : if VUNLIKELY (bytes.is_loaned()) {
466 : 1 : return false;
467 : : }
468 : :
469 : 5 : bytes = Bytes::create(buffer_.size);
470 : :
471 [ - + ]: 5 : if VUNLIKELY (bytes.empty()) {
472 : 0 : return false;
473 : : }
474 : : }
475 : :
476 : 6 : return manager.get_interface()->copy_to_host(manager.native_buffer(buffer_), bytes);
477 : : }
478 : :
479 : 5 : bool FastBuffer::copy_from_host(const Bytes& bytes) noexcept {
480 : 5 : auto& manager = FastBufferManager::get();
481 [ + + + + : 5 : if VUNLIKELY (!is_valid() || bytes.size() != buffer_.size) {
+ + ]
482 : 2 : return false;
483 : : }
484 : :
485 : 3 : return manager.get_interface()->copy_from_host(bytes, manager.native_buffer(buffer_));
486 : : }
487 : :
488 : 33 : uintptr_t FastBuffer::address() const noexcept { return reinterpret_cast<uintptr_t>(buffer_.data); }
489 : :
490 : 35 : size_t FastBuffer::size() const noexcept { return buffer_.size; }
491 : :
492 : 3 : int32_t FastBuffer::device() const noexcept { return buffer_.device; }
493 : :
494 : 8 : FastBuffer::MemoryType FastBuffer::memory_type() const noexcept { return buffer_.memory_type; }
495 : :
496 : 5 : bool FastBuffer::is_owner() const noexcept { return buffer_.handle != nullptr; }
497 : :
498 : 25 : uint8_t* FastBuffer::map(FastBuffer::Access access) const noexcept {
499 : 25 : auto& manager = FastBufferManager::get();
500 [ + + ]: 25 : return is_valid() ? manager.get_interface()->map(manager.native_buffer(buffer_), access) : nullptr;
501 : : }
502 : :
503 : 25 : bool FastBuffer::unmap(FastBuffer::Access access) const noexcept {
504 : 25 : auto& manager = FastBufferManager::get();
505 [ + + + + ]: 25 : return is_valid() && manager.get_interface()->unmap(manager.native_buffer(buffer_), access);
506 : : }
507 : :
508 : 6 : const void* FastBuffer::native_handle() const noexcept {
509 : 6 : auto& manager = FastBufferManager::get();
510 [ + + ]: 6 : return is_valid() ? manager.get_interface()->native_handle(manager.native_buffer(buffer_)) : nullptr;
511 : : }
512 : :
513 : 2 : bool FastBuffer::import_native(const void* handle) noexcept {
514 : 2 : auto& manager = FastBufferManager::get();
515 : 2 : auto* plugin = manager.get_interface();
516 : :
517 [ + - + + : 2 : if VUNLIKELY (!plugin || !handle) {
+ + ]
518 : 1 : return false;
519 : : }
520 : :
521 : 1 : FastBuffer::Buffer buffer;
522 : :
523 [ + - ]: 1 : if VUNLIKELY (!manager.import_native(handle, buffer)) {
524 : 1 : return false;
525 : : }
526 : :
527 : 0 : clear();
528 : 0 : buffer_ = buffer;
529 : :
530 : 0 : return true;
531 : : }
532 : :
533 : : } // namespace zerocopy
534 : :
535 : : } // namespace vlink
|
