diff --git a/src/runtime/rpc/rpc_event_impl.cc b/src/runtime/rpc/rpc_event_impl.cc index 3a7606f2b6b6..3f4782693d8a 100644 --- a/src/runtime/rpc/rpc_event_impl.cc +++ b/src/runtime/rpc/rpc_event_impl.cc @@ -32,7 +32,11 @@ namespace runtime { PackedFunc CreateEventDrivenServer(PackedFunc fsend, std::string name, std::string remote_key) { - std::unique_ptr ch(new CallbackChannel(fsend)); + static PackedFunc frecv([](TVMArgs args, TVMRetValue* rv) { + LOG(FATAL) << "Do not allow explicit receive"; + return 0; + }); + std::unique_ptr ch(new CallbackChannel(fsend, frecv)); std::shared_ptr sess = RPCSession::Create(std::move(ch), name, remote_key); return PackedFunc([sess](TVMArgs args, TVMRetValue* rv) { diff --git a/src/runtime/rpc/rpc_session.h b/src/runtime/rpc/rpc_session.h index 257f395f51d2..d982f68bcb6e 100644 --- a/src/runtime/rpc/rpc_session.h +++ b/src/runtime/rpc/rpc_session.h @@ -253,19 +253,15 @@ class RPCSession { std::string remote_key_; }; +/*! + * \brief RPC channel which callback + * frontend (Python/Java/etc.)'s send & recv function + */ class CallbackChannel final : public RPCChannel { public: explicit CallbackChannel(PackedFunc fsend, PackedFunc frecv) : fsend_(std::move(fsend)), frecv_(std::move(frecv)) {} - explicit CallbackChannel(PackedFunc fsend) - : fsend_(std::move(fsend)) { - PackedFunc frecv([](TVMArgs args, TVMRetValue* rv) { - LOG(FATAL) << "Do not allow explicit receive"; - return 0; - }); - frecv_ = frecv; - } ~CallbackChannel() {} /*! * \brief Send data over to the channel.