chenBright commented on code in PR #3144:
URL: https://github.com/apache/brpc/pull/3144#discussion_r4505565366
##########
src/brpc/policy/baidu_rpc_protocol.cpp:
##########
@@ -823,7 +829,21 @@ void ProcessRpcRequest(InputMessageBase* msg_base) {
butil::IOBuf req_buf;
int body_without_attachment_size = req_size -
meta.attachment_size();
- msg->payload.cutn(&req_buf, body_without_attachment_size);
+#if BRPC_WITH_GDR
+ bool is_gpu_memory = msg->payload.is_gpu_memory();
+ if (is_gpu_memory) {
+ if (!FillReqBufGpu(&req_buf, msg.get(),
body_without_attachment_size)) {
Review Comment:
Only the message body is copied back to host memory here; attachments remain
in GPU memory. When a request has a non-empty attachment and both CRC32C and
`checksum_with_attachment` are enabled, `DeserializeRpcMessage()` below passes
the GPU memory address directly to CPU CRC32C through
`IOBufAsZeroCopyInputStream`, causing an invalid memory access. The four-byte
checksum check does not prevent this path, and the same issue applies on the
response side.
##########
src/brpc/rdma_transport.cpp:
##########
@@ -30,11 +30,24 @@ DECLARE_bool(usercode_in_pthread);
extern SocketVarsCollector *g_vars;
-void RdmaTransport::Init(Socket *socket, const SocketOptions &options) {
+void RdmaTransport::DoInit(Socket *socket, const SocketOptions &options, bool
use_gdr) {
CHECK(_rdma_ep == nullptr);
- if (options.socket_mode == SOCKET_MODE_RDMA) {
- _rdma_ep = new rdma::RdmaEndpoint(socket);
- _rdma_state = RDMA_UNKNOWN;
+ // gdr mode is a special mode of rdma mode.
+ // both rdma mode and gdr mode need init rdma::RdmaEndpoint.
+ if (options.socket_mode == SOCKET_MODE_RDMA ||
+ options.socket_mode == SOCKET_MODE_GDR) {
+ _rdma_ep = new(std::nothrow)rdma::RdmaEndpoint(socket, use_gdr);
+ if (!_rdma_ep) {
+ const int saved_errno = errno;
+ PLOG(ERROR) << "Fail to create RdmaEndpoint";
+ socket->SetFailed(
+ saved_errno, "Fail to create RdmaEndpoint: %s",
berror(saved_errno));
+ // No endpoint: later code paths must not assume RDMA is active
+ // and dereference _rdma_ep. Mark the transport as not using RDMA.
+ _rdma_state = RDMA_OFF;
Review Comment:
When endpoint allocation fails, `Socket::OnCreated()` has not yet called
`HoldHCRelatedRef()`. For a new Socket with only its creation reference,
`SetFailed()` here synchronously releases the last reference and returns the
object to the resource pool, but `DoInit()` and `OnCreated()` then continue
initializing it. Setting `RDMA_OFF` only prevents access to a null endpoint; it
does not resolve the lifetime issue.
##########
src/butil/config.h.in:
##########
@@ -26,6 +26,12 @@
#endif
#define BRPC_WITH_FLATBUFFERS @WITH_FLATBUFFERS_VAL@
+// Defined to 1 when brpc is built with GPU Direct RDMA support.
+#ifdef BRPC_WITH_GDR
+#undef BRPC_WITH_GDR
+#endif
+#cmakedefine BRPC_WITH_GDR @WITH_GDR_VAL@
Review Comment:
`#cmakedefine BRPC_WITH_GDR` checks the CMake variable of the same name, but
only `WITH_GDR_VAL` and `-DBRPC_WITH_GDR=1` are currently set. As a result,
when `WITH_GDR=ON` is enabled, this generated header first `#undef`s the macro
and leaves it undefined, so the GPU implementation being compiled cannot see
the GPU interfaces in subsequent headers. Please generate a numeric macro
directly from `WITH_GDR_VAL`, or also set the variable used by `#cmakedefine`.
##########
src/brpc/channel.cpp:
##########
@@ -145,6 +145,8 @@ static ChannelSignature ComputeChannelSignature(const
ChannelOptions& opt) {
buf.append("|rdma");
} else if (opt.socket_mode == SOCKET_MODE_URMA) {
buf.append("|urma");
+ } else if (opt.socket_mode == SOCKET_MODE_GDR) {
+ buf.append("|gdr");
Review Comment:
With the default configuration, `ComputeChannelSignature()` returns a zero
signature at the start and never reaches `|gdr` here. If a GDR Channel is
created for an address first, followed by a regular `baidu_std` Channel for the
same address, `SocketMap::Insert()` reuses the former's Socket. The regular
channel will therefore also receive attachments in GPU memory, and the existing
`response_attachment().copy_to()` will read the device address directly.
##########
src/brpc/rdma/rdma_endpoint.cpp:
##########
@@ -1133,26 +1148,94 @@ int RdmaEndpoint::DoPostRecv(void* block, size_t
block_size) {
return 0;
}
+int RdmaEndpoint::DoPostRecvGDR(void* block, size_t block_size, uint32_t lkey)
{
+ ibv_recv_wr wr;
+ memset(&wr, 0, sizeof(wr));
+ ibv_sge sge;
+ sge.addr = (uint64_t)block;
+ sge.length = block_size;
+ sge.lkey = lkey;
+ wr.num_sge = 1;
+ wr.sg_list = &sge;
+ //LOG(INFO) << "POST recv: addr=0x" << std::hex << sge.addr
+ // << std::dec << " length=0x" << sge.length
+ // << " lkey=0x" << sge.lkey;
+ //LOG(INFO) << block << " " << _device_allocator->get_lkey();
+ ibv_recv_wr* bad = NULL;
+ int err = ibv_post_recv(_resource->qp, &wr, &bad);
+ if (err != 0) {
+ LOG(WARNING) << "Fail to ibv_post_recv: " << berror(err);
+ return -1;
+ }
+ return 0;
+}
+
int RdmaEndpoint::PostRecv(uint32_t num, bool zerocopy) {
// We do the post repeatedly from the _rbuf[_rq_received].
while (num > 0) {
+#if BRPC_WITH_GDR
+ uint32_t lkey = 0;
+#endif // if BRPC_WITH_GDR
if (zerocopy) {
_rbuf[_rq_received].clear();
- butil::IOBufAsZeroCopyOutputStream os(&_rbuf[_rq_received],
- g_rdma_recv_block_size + IOBUF_BLOCK_HEADER_LEN);
- int size = 0;
- if (!os.Next(&_rbuf_data[_rq_received], &size)) {
- // Memory is not enough for preparing a block
- PLOG(WARNING) << "Fail to allocate rbuf";
- return -1;
- } else {
- CHECK_EQ(static_cast<uint32_t>(size), g_rdma_recv_block_size);
+
+#if BRPC_WITH_GDR
+ if (_use_gdr) {
+ butil::gdr::BlockPoolAllocator* device_allocator =
butil::gdr::BlockPoolAllocators::GetInstance()->get_gpu_allocator();
+ void* device_ptr =
device_allocator->AllocateRaw(g_gdr_recv_block_size);
+ if (device_ptr == nullptr) {
+ // GPU pool exhausted; abort this post, the caller will
retry.
+ PLOG(WARNING) << "Fail to allocate gpu rbuf";
+ return -1;
+ }
+ auto deleter = [device_allocator](void* data) {
device_allocator->DeallocateRaw(data); };
+ lkey = device_allocator->get_lkey(device_ptr);
+ // Tag the block as GPU memory so that IOBuf::is_gpu_memory()
+ // works deterministically. The lkey is carried via `data_meta'
+ // so the recv path can locate the correct memory region key.
+ if (_rbuf[_rq_received].append_user_data_gpu(
+ device_ptr, g_gdr_recv_block_size, deleter, lkey) !=
0) {
+ // Failed to attach the GPU block to the IOBuf (e.g. the
+ // IOBuf metadata block allocation failed). Release the
+ // device block back to the pool and abort this post; the
+ // caller will retry. Publishing _rbuf_data without a
+ // successful append would leave the posted block ownerless
+ // and silently drop the receive on completion.
+ PLOG(WARNING) << "Fail to append gpu rbuf to IOBuf";
+ device_allocator->DeallocateRaw(device_ptr);
+ return -1;
+ }
+ _rbuf_data[_rq_received] = device_ptr;
+ } else
+#endif // if BRPC_WITH_GDR
+ {
+ butil::IOBufAsZeroCopyOutputStream os(&_rbuf[_rq_received],
+ g_rdma_recv_block_size + IOBUF_BLOCK_HEADER_LEN);
+ int size = 0;
+ if (!os.Next(&_rbuf_data[_rq_received], &size)) {
+ // Memory is not enough for preparing a block
+ PLOG(WARNING) << "Fail to allocate rbuf";
+ return -1;
+ } else {
+ CHECK_EQ(static_cast<uint32_t>(size),
g_rdma_recv_block_size);
+ }
}
}
- if (DoPostRecv(_rbuf_data[_rq_received], g_rdma_recv_block_size) < 0) {
- _rbuf[_rq_received].clear();
- return -1;
+#if BRPC_WITH_GDR
+ if (_use_gdr) {
+ if (DoPostRecvGDR(_rbuf_data[_rq_received], g_gdr_recv_block_size,
lkey) < 0) {
Review Comment:
When `rdma_recv_zerocopy=false`, a zero-byte ACK does not enter the GDR
`zerocopy=true` branch above and instead leads to `PostRecv(1, false)`. This
reuses the GPU memory address, but `lkey` remains at its initial value of `0`;
the next data receive may therefore cause a protection error and disconnect the
connection.
##########
src/butil/iobuf.cpp:
##########
@@ -768,6 +771,46 @@ size_t IOBuf::cutn(IOBuf* out, size_t n) {
return saved_n;
}
+#if BRPC_WITH_GDR
+size_t IOBuf::cutn_from_gpu(IOBuf* out, size_t n) {
+ if (n == 0) {
+ return 0;
+ }
+
+ butil::gdr::BlockPoolAllocator* host_allocator =
butil::gdr::BlockPoolAllocators::GetInstance()->get_cpu_allocator();
+ bool alloc_from_host_alloc = (n <= host_allocator->get_block_size());
+ void* mem = NULL;
+ if (alloc_from_host_alloc) {
+ mem = host_allocator->AllocateRaw(n);
+ } else {
+ mem = malloc(n);
+ }
+
+ if (mem == NULL) {
+ return 0;
+ }
+ size_t saved_n = copy_from_gpu(mem, n, 0, false);
+ if (saved_n > 0) {
+ if (alloc_from_host_alloc) {
+ auto deleter = [host_allocator](void* data) {
host_allocator->DeallocateRaw(data); };
+ out->append_user_data(mem, saved_n, deleter);
+ } else {
+ auto deleter = [](void* data) { free(data); };
+ out->append_user_data(mem, saved_n, deleter);
+ }
+ pop_front(saved_n);
Review Comment:
`append_user_data()` can return `-1` if block metadata allocation fails, but
this code still calls `pop_front(saved_n)` and returns a successful length. The
destination receives no data, while the source data has already been consumed
and the temporary host memory leaks. The protocol layer's length check also
incorrectly reports success.
##########
src/brpc/policy/baidu_rpc_protocol_gpu.cpp:
##########
@@ -0,0 +1,294 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <google/protobuf/descriptor.h> // MethodDescriptor
+#include <google/protobuf/message.h> // Message
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include <google/protobuf/io/coded_stream.h>
+#include <google/protobuf/text_format.h>
+
+#include "butil/gpu/gpu_block_pool.h"
+#include "butil/iobuf.h" // butil::IOBuf
+#include "butil/logging.h" // LOG()
+#include "butil/memory/scope_guard.h"
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/strings/string_util.h"
+
+#include "json2pb/json_to_pb.h"
+#include "json2pb/pb_to_json.h"
+#include "brpc/controller.h" // Controller
+#include "brpc/socket.h" // Socket
+#include "brpc/server.h" // Server
+#include "brpc/span.h"
+#include "brpc/compress.h" // ParseFromCompressedData
+#include "brpc/checksum.h"
+#include "brpc/stream_impl.h"
+#include "brpc/rpc_dump.h" // SampledRequest
+#include "brpc/rpc_pb_message_factory.h"
+#include "brpc/policy/baidu_rpc_meta.pb.h" // RpcRequestMeta
+#include "brpc/policy/baidu_rpc_protocol.h"
+#include "brpc/policy/most_common_message.h"
+#include "brpc/policy/streaming_rpc_protocol.h"
+#include "brpc/details/usercode_backup_pool.h"
+#include "brpc/details/controller_private_accessor.h"
+#include "brpc/details/server_private_accessor.h"
+
+namespace brpc {
+namespace policy {
+
+// Notes:
+// 1. 12-byte header [PRPC][body_size][meta_size]
+// 2. body_size and meta_size are in network byte order
+// 3. Use service->full_name() + method_name to specify the method to call
+// 4. `attachment_size' is set iff request/response has attachment
+// 5. Not supported: chunk_info
+
+// Pack header into `buf'
+
+const int header_size = 12;
+// if we recv data into gpu, the header/meta/body will be copied to cpu and
processed.
+// in to to limit the count of d2h, we will prefetch some bytes from gpu to
cpu.
+// if header_size + meta_size + body_size(without attachment) is less than the
+// prefetch size, then one d2h is enough for one rpc.
+// The prefetch size is derived from the CPU block pool's block size so that
the
+// allocation never exceeds the allocator's single-block limit and stays
effective
+// if the block size changes.
+
+static const int kDefaultPrefetchD2HSize = 512;
+
+static int GetPrefetchD2HSize(butil::gdr::BlockPoolAllocator* host_allocator) {
+ const int64_t block_size = host_allocator->get_block_size();
+ return block_size > 0 ? static_cast<int>(block_size)
+ : kDefaultPrefetchD2HSize;
+}
+
+ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
+ bool /*read_eof*/, const void*) {
+
+ char header_buf[12];
+ size_t n = 0;
+ uint32_t body_size;
+ uint32_t meta_size;
+ ParseError pe = PARSE_OK;
+
+ void* prefetch_d2h_data = NULL;
+ bool is_gpu_memory = source->is_gpu_memory();
+ if (!is_gpu_memory) {
+ // A non-GPU buffer here means the message arrived through a non-GDR
+ // path; reject it instead of aborting the whole process.
+ LOG(ERROR) << "RpcMessage is not in gpu!!!";
+ return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
+ }
+ butil::gdr::BlockPoolAllocator* host_allocator =
butil::gdr::BlockPoolAllocators::GetInstance()->get_cpu_allocator();
+ const int prefetch_d2h_size = GetPrefetchD2HSize(host_allocator);
+ prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size);
+ if (prefetch_d2h_data == NULL) {
+ // Pool exhausted or single-block limit hit: ask the caller to retry
+ // instead of crashing the whole server.
+ LOG(ERROR) << "alloc host data for prefetch d2h failed, size="
+ << prefetch_d2h_size;
+ return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
+ }
+
+ // n is the bytes we real frefetch, n maybe less than prefetch_d2h_size;
+ n = source->copy_from_gpu(prefetch_d2h_data, prefetch_d2h_size);
+ size_t copy_size = n > 12 ? 12 : n;
+ memcpy(header_buf, prefetch_d2h_data, copy_size);
+
+ do {
+ // NOTE: unlike the CPU ParseRpcMessage, malformed GPU frames must NOT
+ // return PARSE_ERROR_TRY_OTHERS: InputMessengerProcessor would then
+ // dispatch the same GPU-resident IOBuf to non-GDR parsers whose
+ // copy_to cannot dereference device memory and would crash. A GDR
+ // connection only speaks baidu_std, so treat this as terminal.
+ if (n >= 4) {
+ void* dummy = header_buf;
+ if (*(const uint32_t*)dummy != *(const uint32_t*)"PRPC") {
+ pe = PARSE_ERROR_ABSOLUTELY_WRONG;
+ break;
+ }
+ } else {
+ if (memcmp(header_buf, "PRPC", n) != 0) {
+ pe = PARSE_ERROR_ABSOLUTELY_WRONG;
+ break;
+ }
+ }
+ if (n < sizeof(header_buf)) {
+ pe = PARSE_ERROR_NOT_ENOUGH_DATA;
+ break;
+ }
+ butil::RawUnpacker(header_buf +
4).unpack32(body_size).unpack32(meta_size);
+ if (body_size > FLAGS_max_body_size) {
+ // We need this log to report the body_size to give users some
clues
+ // which is not printed in InputMessenger.
+ LOG(ERROR) << "body_size=" << body_size << " from "
+ << socket->remote_side() << " is too large";
+ pe = PARSE_ERROR_TOO_BIG_DATA;
+ break;
+ } else if (source->length() < sizeof(header_buf) + body_size) {
+ pe = PARSE_ERROR_NOT_ENOUGH_DATA;
+ break;
+ }
+ if (meta_size > body_size) {
+ LOG(ERROR) << "meta_size=" << meta_size << " is bigger than
body_size="
+ << body_size;
+ // Pop the message
+ source->pop_front(sizeof(header_buf) + body_size);
+ pe = PARSE_ERROR_ABSOLUTELY_WRONG;
+ break;
+ }
+ } while (0);
+
+ if (pe != PARSE_OK) {
+ host_allocator->DeallocateRaw(prefetch_d2h_data);
+ return MakeParseError(pe);
+ }
+
+ source->pop_front(sizeof(header_buf));
+ MostCommonMessage* msg = MostCommonMessage::Get();
+
+ if (header_size + meta_size <= n) {
+ auto deleter = [host_allocator, prefetch_d2h_data](void* data) {
host_allocator->DeallocateRaw(prefetch_d2h_data); };
+ // `n` is the number of bytes actually prefetched from GPU. It is
stored
+ // in the meta's `data_meta' so that FillReqBufGpu/FillResBufGpu can
+ // decide whether the prefetched host buffer also covers the body (and
+ // thus avoid a second D2H copy). This block is host memory
(prefetched),
+ // NOT GPU memory, so we deliberately use append_user_data_with_meta
+ // (not append_user_data_gpu) to keep is_gpu_memory() returning false.
+ if (msg->meta.append_user_data_with_meta(
+ (char*)prefetch_d2h_data + header_size, meta_size, deleter, n)
!= 0) {
+ // append failed (internal malloc): the prefetch buffer is NOT
+ // owned by anyone; release it back to the pool instead of leaking.
+ LOG(ERROR) << "Fail to append prefetched meta to IOBuf";
+ host_allocator->DeallocateRaw(prefetch_d2h_data);
+ return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
Review Comment:
Before returning here, the code has already called `pop_front(12)` and
obtained `msg`. Returning `NOT_ENOUGH_DATA` at this point makes the caller wait
for new input, but the next parse still starts at the `PRPC` frame header while
the source buffer now starts at the metadata. Meanwhile, `msg` is not returned
to the message pool.
##########
src/brpc/policy/baidu_rpc_protocol_gpu.cpp:
##########
@@ -0,0 +1,294 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <google/protobuf/descriptor.h> // MethodDescriptor
+#include <google/protobuf/message.h> // Message
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include <google/protobuf/io/coded_stream.h>
+#include <google/protobuf/text_format.h>
+
+#include "butil/gpu/gpu_block_pool.h"
+#include "butil/iobuf.h" // butil::IOBuf
+#include "butil/logging.h" // LOG()
+#include "butil/memory/scope_guard.h"
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/strings/string_util.h"
+
+#include "json2pb/json_to_pb.h"
+#include "json2pb/pb_to_json.h"
+#include "brpc/controller.h" // Controller
+#include "brpc/socket.h" // Socket
+#include "brpc/server.h" // Server
+#include "brpc/span.h"
+#include "brpc/compress.h" // ParseFromCompressedData
+#include "brpc/checksum.h"
+#include "brpc/stream_impl.h"
+#include "brpc/rpc_dump.h" // SampledRequest
+#include "brpc/rpc_pb_message_factory.h"
+#include "brpc/policy/baidu_rpc_meta.pb.h" // RpcRequestMeta
+#include "brpc/policy/baidu_rpc_protocol.h"
+#include "brpc/policy/most_common_message.h"
+#include "brpc/policy/streaming_rpc_protocol.h"
+#include "brpc/details/usercode_backup_pool.h"
+#include "brpc/details/controller_private_accessor.h"
+#include "brpc/details/server_private_accessor.h"
+
+namespace brpc {
+namespace policy {
+
+// Notes:
+// 1. 12-byte header [PRPC][body_size][meta_size]
+// 2. body_size and meta_size are in network byte order
+// 3. Use service->full_name() + method_name to specify the method to call
+// 4. `attachment_size' is set iff request/response has attachment
+// 5. Not supported: chunk_info
+
+// Pack header into `buf'
+
+const int header_size = 12;
+// if we recv data into gpu, the header/meta/body will be copied to cpu and
processed.
+// in to to limit the count of d2h, we will prefetch some bytes from gpu to
cpu.
+// if header_size + meta_size + body_size(without attachment) is less than the
+// prefetch size, then one d2h is enough for one rpc.
+// The prefetch size is derived from the CPU block pool's block size so that
the
+// allocation never exceeds the allocator's single-block limit and stays
effective
+// if the block size changes.
+
+static const int kDefaultPrefetchD2HSize = 512;
+
+static int GetPrefetchD2HSize(butil::gdr::BlockPoolAllocator* host_allocator) {
+ const int64_t block_size = host_allocator->get_block_size();
+ return block_size > 0 ? static_cast<int>(block_size)
+ : kDefaultPrefetchD2HSize;
+}
+
+ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
+ bool /*read_eof*/, const void*) {
+
+ char header_buf[12];
+ size_t n = 0;
+ uint32_t body_size;
+ uint32_t meta_size;
+ ParseError pe = PARSE_OK;
+
+ void* prefetch_d2h_data = NULL;
+ bool is_gpu_memory = source->is_gpu_memory();
+ if (!is_gpu_memory) {
+ // A non-GPU buffer here means the message arrived through a non-GDR
+ // path; reject it instead of aborting the whole process.
+ LOG(ERROR) << "RpcMessage is not in gpu!!!";
+ return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
+ }
+ butil::gdr::BlockPoolAllocator* host_allocator =
butil::gdr::BlockPoolAllocators::GetInstance()->get_cpu_allocator();
+ const int prefetch_d2h_size = GetPrefetchD2HSize(host_allocator);
+ prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size);
+ if (prefetch_d2h_data == NULL) {
+ // Pool exhausted or single-block limit hit: ask the caller to retry
+ // instead of crashing the whole server.
+ LOG(ERROR) << "alloc host data for prefetch d2h failed, size="
+ << prefetch_d2h_size;
+ return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
+ }
+
+ // n is the bytes we real frefetch, n maybe less than prefetch_d2h_size;
+ n = source->copy_from_gpu(prefetch_d2h_data, prefetch_d2h_size);
+ size_t copy_size = n > 12 ? 12 : n;
+ memcpy(header_buf, prefetch_d2h_data, copy_size);
+
+ do {
+ // NOTE: unlike the CPU ParseRpcMessage, malformed GPU frames must NOT
+ // return PARSE_ERROR_TRY_OTHERS: InputMessengerProcessor would then
+ // dispatch the same GPU-resident IOBuf to non-GDR parsers whose
+ // copy_to cannot dereference device memory and would crash. A GDR
+ // connection only speaks baidu_std, so treat this as terminal.
+ if (n >= 4) {
+ void* dummy = header_buf;
+ if (*(const uint32_t*)dummy != *(const uint32_t*)"PRPC") {
+ pe = PARSE_ERROR_ABSOLUTELY_WRONG;
+ break;
+ }
+ } else {
+ if (memcmp(header_buf, "PRPC", n) != 0) {
+ pe = PARSE_ERROR_ABSOLUTELY_WRONG;
+ break;
+ }
+ }
+ if (n < sizeof(header_buf)) {
+ pe = PARSE_ERROR_NOT_ENOUGH_DATA;
+ break;
+ }
+ butil::RawUnpacker(header_buf +
4).unpack32(body_size).unpack32(meta_size);
+ if (body_size > FLAGS_max_body_size) {
+ // We need this log to report the body_size to give users some
clues
+ // which is not printed in InputMessenger.
+ LOG(ERROR) << "body_size=" << body_size << " from "
+ << socket->remote_side() << " is too large";
+ pe = PARSE_ERROR_TOO_BIG_DATA;
+ break;
+ } else if (source->length() < sizeof(header_buf) + body_size) {
+ pe = PARSE_ERROR_NOT_ENOUGH_DATA;
+ break;
+ }
+ if (meta_size > body_size) {
+ LOG(ERROR) << "meta_size=" << meta_size << " is bigger than
body_size="
+ << body_size;
+ // Pop the message
+ source->pop_front(sizeof(header_buf) + body_size);
+ pe = PARSE_ERROR_ABSOLUTELY_WRONG;
+ break;
+ }
+ } while (0);
+
+ if (pe != PARSE_OK) {
+ host_allocator->DeallocateRaw(prefetch_d2h_data);
+ return MakeParseError(pe);
+ }
+
+ source->pop_front(sizeof(header_buf));
+ MostCommonMessage* msg = MostCommonMessage::Get();
+
+ if (header_size + meta_size <= n) {
+ auto deleter = [host_allocator, prefetch_d2h_data](void* data) {
host_allocator->DeallocateRaw(prefetch_d2h_data); };
+ // `n` is the number of bytes actually prefetched from GPU. It is
stored
+ // in the meta's `data_meta' so that FillReqBufGpu/FillResBufGpu can
+ // decide whether the prefetched host buffer also covers the body (and
+ // thus avoid a second D2H copy). This block is host memory
(prefetched),
+ // NOT GPU memory, so we deliberately use append_user_data_with_meta
+ // (not append_user_data_gpu) to keep is_gpu_memory() returning false.
+ if (msg->meta.append_user_data_with_meta(
+ (char*)prefetch_d2h_data + header_size, meta_size, deleter, n)
!= 0) {
+ // append failed (internal malloc): the prefetch buffer is NOT
+ // owned by anyone; release it back to the pool instead of leaking.
+ LOG(ERROR) << "Fail to append prefetched meta to IOBuf";
+ host_allocator->DeallocateRaw(prefetch_d2h_data);
+ return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
+ }
+ source->pop_front(meta_size);
+ } else {
+ host_allocator->DeallocateRaw(prefetch_d2h_data);
+ if (source->cutn_from_gpu(&msg->meta, meta_size) != meta_size) {
+ // Host pool exhausted or the GPU copy failed: meta is NOT cut out
+ // of `source'. Continuing would make the following cutn treat the
+ // meta bytes as payload and silently corrupt the message.
+ LOG(ERROR) << "Fail to cut meta_size=" << meta_size << " from gpu";
+ return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
+ }
+ }
+ source->cutn(&msg->payload, body_size - meta_size);
+ return MakeMessage(msg);
+}
+
+
+// Returns true on success. Returns false (nothing cut out of msg->payload)
+// when the GPU -> host copy fails; the caller should fail the request.
+bool FillReqBufGpu(butil::IOBuf* req_buf, MostCommonMessage* msg, int
body_without_attachment_size) {
+ int meta_size = msg->meta.size();
+ if (!msg->payload.is_gpu_memory()) {
+ // Should not happen on a GDR connection; do not crash the server.
+ LOG(ERROR) << "message is not on gpu!!!";
+ return false;
+ }
+ int64_t real_prefetch_d2h_size = msg->meta.get_first_data_meta();
+ if (header_size + meta_size + body_without_attachment_size <=
real_prefetch_d2h_size) {
+ void* data = msg->meta.get_first_data_ptr();
+ if (data == nullptr) {
+ LOG(ERROR) << "illegal prefetched data";
+ return false;
+ }
+ req_buf->append((char*)data + meta_size, body_without_attachment_size);
+ msg->payload.pop_front(body_without_attachment_size);
+ } else {
+ if (msg->payload.cutn_from_gpu(req_buf, body_without_attachment_size)
+ != (size_t)body_without_attachment_size) {
+ LOG(ERROR) << "Fail to cut " << body_without_attachment_size
+ << " bytes from gpu";
+ return false;
+ }
+ }
+ return true;
+}
+
+// Returns true on success; on failure cntl is SetFailed and the caller must
+// stop processing this response (res_buf_ptr is left untouched).
+bool FillResBufGpu(butil::IOBuf* res_buf, MostCommonMessage* msg, const
RpcMeta& meta,
+ butil::IOBuf** res_buf_ptr, Controller* cntl) {
+ const int res_size = msg->payload.length();
+ int meta_size = msg->meta.size();
+ if (!msg->payload.is_gpu_memory()) {
+ // Should not happen on a GDR connection; fail this response instead
+ // of crashing the whole server.
+ cntl->SetFailed(ERESPONSE, "response payload is not on gpu memory");
+ return false;
+ }
+ if (meta.has_attachment_size()) {
+ if (meta.attachment_size() > res_size) {
+ cntl->SetFailed(
+ ERESPONSE, "attachment_size=%d is larger than
response_size=%d",
+ meta.attachment_size(), res_size);
+ return false;
+ }
+ int body_without_attachment_size = res_size - meta.attachment_size();
Review Comment:
This omits the `attachment_size < 0` check present in the CPU response path.
If the current payload contains only `A` and the attachment length is `-1`, the
message body length is calculated as `2`. When the prefetch also includes the
next frame, the fast path copies out `AP`, where `P` comes from the next
frame's `PRPC`. The `SerializedResponse` path passes it directly to the caller
as a successful response.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
