diff --git a/cpp/src/arrow/flight/CMakeLists.txt b/cpp/src/arrow/flight/CMakeLists.txt index 8974c9581f7c..87561b295aa5 100644 --- a/cpp/src/arrow/flight/CMakeLists.txt +++ b/cpp/src/arrow/flight/CMakeLists.txt @@ -186,9 +186,11 @@ set(ARROW_FLIGHT_SRCS server_tracing_middleware.cc transport.cc transport_server.cc + transport_server_internal.cc # Bundle the gRPC impl with libarrow_flight transport/grpc/grpc_client.cc transport/grpc/grpc_server.cc + transport/grpc/grpc_server_internal.cc transport/grpc/serialization_internal.cc transport/grpc/protocol_grpc_internal.cc transport/grpc/util_internal.cc diff --git a/cpp/src/arrow/flight/flight_test.cc b/cpp/src/arrow/flight/flight_test.cc index eb931011906f..352bd7c1dfea 100644 --- a/cpp/src/arrow/flight/flight_test.cc +++ b/cpp/src/arrow/flight/flight_test.cc @@ -1668,12 +1668,14 @@ TEST_F(TestCancel, DoExchange) { stop_source.RequestStop(Status::Cancelled("StopSource")); ASSERT_OK_AND_ASSIGN(auto do_exchange_result, client_->DoExchange(options, FlightDescriptor::Command(""))); + ASSERT_OK(do_exchange_result.writer->DoneWriting()); EXPECT_RAISES_WITH_MESSAGE_THAT(Cancelled, ::testing::HasSubstr("StopSource"), do_exchange_result.reader->ToTable()); ARROW_UNUSED(do_exchange_result.writer->Close()); ASSERT_OK_AND_ASSIGN(do_exchange_result, client_->DoExchange(FlightDescriptor::Command(""))); + ASSERT_OK(do_exchange_result.writer->DoneWriting()); EXPECT_RAISES_WITH_MESSAGE_THAT(Cancelled, ::testing::HasSubstr("StopSource"), do_exchange_result.reader->ToTable(options.stop_token)); ARROW_UNUSED(do_exchange_result.writer->Close()); diff --git a/cpp/src/arrow/flight/meson.build b/cpp/src/arrow/flight/meson.build index e66bc904d5fa..dcdacc42ffdc 100644 --- a/cpp/src/arrow/flight/meson.build +++ b/cpp/src/arrow/flight/meson.build @@ -88,8 +88,10 @@ arrow_flight_srcs = [ 'server_tracing_middleware.cc', 'transport.cc', 'transport_server.cc', + 'transport_server_internal.cc', 'transport/grpc/grpc_client.cc', 'transport/grpc/grpc_server.cc', + 'transport/grpc/grpc_server_internal.cc', 'transport/grpc/serialization_internal.cc', 'transport/grpc/protocol_grpc_internal.cc', 'transport/grpc/util_internal.cc', diff --git a/cpp/src/arrow/flight/server.cc b/cpp/src/arrow/flight/server.cc index 7807829b3b3b..68e307816cdf 100644 --- a/cpp/src/arrow/flight/server.cc +++ b/cpp/src/arrow/flight/server.cc @@ -23,119 +23,30 @@ #include "arrow/flight/server.h" -#include -#include #include -#include #include #include -#include -#include #include #include "arrow/device.h" #include "arrow/flight/transport.h" #include "arrow/flight/transport/grpc/grpc_server.h" #include "arrow/flight/transport_server.h" +#include "arrow/flight/transport_server_internal.h" #include "arrow/flight/types.h" #include "arrow/status.h" #include "arrow/type.h" -#include "arrow/util/io_util.h" -#include "arrow/util/logging.h" -#include "arrow/util/uri.h" namespace arrow { namespace flight { -namespace { -#if (ATOMIC_INT_LOCK_FREE != 2 || ATOMIC_POINTER_LOCK_FREE != 2) -# error "atomic ints and atomic pointers not always lock-free!" -#endif - -using ::arrow::internal::SelfPipe; -using ::arrow::internal::SetSignalHandler; -using ::arrow::internal::SignalHandler; - -/// RAII guard that manages a self-pipe and a thread that listens on -/// the self-pipe, shutting down the server when a signal handler -/// writes to the pipe. -class ServerSignalHandler { - public: - ARROW_DISALLOW_COPY_AND_ASSIGN(ServerSignalHandler); - ServerSignalHandler() = default; - - /// Create the pipe and handler thread. - /// - /// \return the fd of the write side of the pipe. - template - arrow::Result> Init(Fn handler) { - ARROW_ASSIGN_OR_RAISE(self_pipe_, SelfPipe::Make(/*signal_safe=*/true)); - handle_signals_ = std::thread(handler, self_pipe_); - return self_pipe_; - } - - Status Shutdown() { - RETURN_NOT_OK(self_pipe_->Shutdown()); - handle_signals_.join(); - return Status::OK(); - } - - ~ServerSignalHandler() { ARROW_CHECK_OK(Shutdown()); } - - private: - std::shared_ptr self_pipe_; - std::thread handle_signals_; -}; -} // namespace - /// Server implementation. Manages the lifecycle of the "real" server /// (ServerTransport) and contains struct FlightServerBase::Impl { std::unique_ptr transport_; - - // Signal handlers (on Windows) and the shutdown handler (other platforms) - // are executed in a separate thread, so getting the current thread instance - // wouldn't make sense. This means only a single instance can receive signals. - static std::atomic running_instance_; - // We'll use the self-pipe trick to notify a thread from the signal - // handler. The thread will then shut down the server. - std::shared_ptr self_pipe_; - - // Signal handling - std::vector signals_; - std::vector old_signal_handlers_; - std::atomic got_signal_; - - static void HandleSignal(int signum) { - auto instance = running_instance_.load(); - if (instance != nullptr) { - instance->DoHandleSignal(signum); - } - } - - void DoHandleSignal(int signum) { - got_signal_ = signum; - - // Send dummy payload over self-pipe - self_pipe_->Send(/*payload=*/0); - } - - static void WaitForSignals(std::shared_ptr self_pipe) { - // Wait for a signal handler to wake up the pipe - auto st = self_pipe->Wait().status(); - // Status::Invalid means the pipe was shutdown without any wakeup - if (!st.ok() && !st.IsInvalid()) { - ARROW_LOG(FATAL) << "Failed to wait on self-pipe: " << st.ToString(); - } - auto instance = running_instance_.load(); - if (instance != nullptr) { - ARROW_WARN_NOT_OK(instance->transport_->Shutdown(), "Error shutting down server"); - } - } + internal::ServerSignalState signal_state_; }; -std::atomic FlightServerBase::Impl::running_instance_; - FlightServerOptions::FlightServerOptions(const Location& location_) : location(location_), auth_handler(nullptr), @@ -159,66 +70,61 @@ Status FlightServerBase::Init(const FlightServerOptions& options) { ARROW_ASSIGN_OR_RAISE(impl_->transport_, internal::GetDefaultTransportRegistry()->MakeServer( scheme, this, options.memory_manager)); - return impl_->transport_->Init(options, *options.location.uri_); + ARROW_ASSIGN_OR_RAISE(auto uri, internal::ParseLocationUri(options.location)); + return impl_->transport_->Init(options, uri); } -int FlightServerBase::port() const { return location().uri_->port(); } +int FlightServerBase::port() const { return internal::PortFromLocation(location()); } Location FlightServerBase::location() const { return impl_->transport_->location(); } Status FlightServerBase::SetShutdownOnSignals(const std::vector sigs) { - impl_->signals_ = sigs; - impl_->old_signal_handlers_.clear(); - return Status::OK(); + return impl_->signal_state_.SetShutdownOnSignals(sigs); } Status FlightServerBase::Serve() { - if (!impl_->transport_) { - return Status::UnknownError("Server did not start properly"); - } - impl_->got_signal_ = 0; - impl_->old_signal_handlers_.clear(); - impl_->running_instance_ = impl_.get(); - - ServerSignalHandler signal_handler; - ARROW_ASSIGN_OR_RAISE(impl_->self_pipe_, signal_handler.Init(&Impl::WaitForSignals)); - // Override existing signal handlers with our own handler so as to stop the server. - for (size_t i = 0; i < impl_->signals_.size(); ++i) { - int signum = impl_->signals_[i]; - SignalHandler new_handler(&Impl::HandleSignal), old_handler; - ARROW_ASSIGN_OR_RAISE(old_handler, SetSignalHandler(signum, new_handler)); - impl_->old_signal_handlers_.push_back(std::move(old_handler)); - } - - RETURN_NOT_OK(impl_->transport_->Wait()); - impl_->running_instance_ = nullptr; - - // Restore signal handlers - for (size_t i = 0; i < impl_->signals_.size(); ++i) { - RETURN_NOT_OK( - SetSignalHandler(impl_->signals_[i], impl_->old_signal_handlers_[i]).status()); - } - return Status::OK(); + return impl_->signal_state_.Serve( + [this]() -> Status { + if (!impl_->transport_) { + return Status::UnknownError("Server did not start properly"); + } + return impl_->transport_->Wait(); + }, + [this](const std::chrono::system_clock::time_point* deadline) -> Status { + if (!impl_->transport_) { + return Status::Invalid("Shutdown() on uninitialized FlightServerBase"); + } + if (deadline) { + return impl_->transport_->Shutdown(*deadline); + } + return impl_->transport_->Shutdown(); + }, + "Server did not start properly", "Error shutting down server"); } -int FlightServerBase::GotSignal() const { return impl_->got_signal_; } +int FlightServerBase::GotSignal() const { return impl_->signal_state_.GotSignal(); } Status FlightServerBase::Shutdown(const std::chrono::system_clock::time_point* deadline) { - auto server = impl_->transport_.get(); - if (!server) { - return Status::Invalid("Shutdown() on uninitialized FlightServerBase"); - } - impl_->running_instance_ = nullptr; - if (deadline) { - return impl_->transport_->Shutdown(*deadline); - } - return impl_->transport_->Shutdown(); + return impl_->signal_state_.Shutdown( + [this](const std::chrono::system_clock::time_point* maybe_deadline) -> Status { + if (!impl_->transport_) { + return Status::Invalid("Shutdown() on uninitialized FlightServerBase"); + } + if (maybe_deadline) { + return impl_->transport_->Shutdown(*maybe_deadline); + } + return impl_->transport_->Shutdown(); + }, + deadline); } Status FlightServerBase::Wait() { - RETURN_NOT_OK(impl_->transport_->Wait()); - impl_->running_instance_ = nullptr; - return Status::OK(); + return impl_->signal_state_.Wait([this] { + if (!impl_->transport_) { + return Status::Invalid("Wait() on uninitialized FlightServerBase"); + } + return impl_->transport_->Wait(); + }); } Status FlightServerBase::ListFlights(const ServerCallContext& context, diff --git a/cpp/src/arrow/flight/server_tracing_middleware.cc b/cpp/src/arrow/flight/server_tracing_middleware.cc index a22e1a98ce49..5ce337e3cb43 100644 --- a/cpp/src/arrow/flight/server_tracing_middleware.cc +++ b/cpp/src/arrow/flight/server_tracing_middleware.cc @@ -95,7 +95,7 @@ class TracingServerMiddleware::Impl { : scope_(std::move(scope)), span_(std::move(span)) {} void CallCompleted(const Status& status) { if (!status.ok()) { - auto grpc_status = transport::grpc::ToGrpcStatus(status, /*ctx=*/nullptr); + auto grpc_status = transport::grpc::ToGrpcStatus(status); span_->SetStatus(otel::trace::StatusCode::kError, status.ToString()); span_->SetAttribute(SemanticConventions::kRpcGrpcStatusCode, static_cast(grpc_status.error_code())); diff --git a/cpp/src/arrow/flight/test_definitions.cc b/cpp/src/arrow/flight/test_definitions.cc index c6b8e2b422ca..ecf5c3cc03b3 100644 --- a/cpp/src/arrow/flight/test_definitions.cc +++ b/cpp/src/arrow/flight/test_definitions.cc @@ -314,11 +314,10 @@ void DataTest::TestOverflowServerBatch() { // DoExchange: check for overflow on large batch from server auto descr = FlightDescriptor::Command("large_batch"); ASSERT_OK_AND_ASSIGN(auto do_exchange_result, client_->DoExchange(descr)); - RecordBatchVector batches; + ASSERT_OK(do_exchange_result.writer->DoneWriting()); EXPECT_RAISES_WITH_MESSAGE_THAT( Invalid, ::testing::HasSubstr("Cannot send record batches exceeding 2GiB yet"), - do_exchange_result.reader->ToRecordBatches().Value(&batches)); - ARROW_UNUSED(do_exchange_result.writer->Close()); + do_exchange_result.writer->Close()); } } void DataTest::TestOverflowClientBatch() { diff --git a/cpp/src/arrow/flight/transport/grpc/customize_grpc.h b/cpp/src/arrow/flight/transport/grpc/customize_grpc.h index 7836f7c161d6..c55211cf4bfd 100644 --- a/cpp/src/arrow/flight/transport/grpc/customize_grpc.h +++ b/cpp/src/arrow/flight/transport/grpc/customize_grpc.h @@ -86,6 +86,16 @@ ::grpc::Status FlightDataSerialize(const arrow::flight::FlightPayload& msg, // protobuf without copying ARROW_FLIGHT_EXPORT ::grpc::Status FlightDataDeserialize( ::grpc::ByteBuffer* buffer, arrow::flight::internal::FlightData* out); + +ARROW_FLIGHT_EXPORT +bool IsRegisteredGrpcFlightDataMessage( + const arrow::flight::protocol::FlightData* message); + +ARROW_FLIGHT_EXPORT +void RegisterGrpcFlightDataMessage(const arrow::flight::protocol::FlightData* message); + +ARROW_FLIGHT_EXPORT +void UnregisterGrpcFlightDataMessage(const arrow::flight::protocol::FlightData* message); } // namespace grpc } // namespace transport } // namespace flight @@ -105,11 +115,22 @@ class SerializationTraits { // In the functions below, we cast back the Message argument to its real // type (see ReadPayload() and WritePayload() for the initial cast). static Status Serialize(const MessageType& msg, ByteBuffer* bb, bool* own_buffer) { + const auto* flight_data = + reinterpret_cast(&msg); + if (arrow::flight::transport::grpc::IsRegisteredGrpcFlightDataMessage(flight_data)) { + return GenericSerialize( + msg, bb, own_buffer); + } return arrow::flight::transport::grpc::FlightDataSerialize( *reinterpret_cast(&msg), bb, own_buffer); } static Status Deserialize(ByteBuffer* buffer, MessageType* msg) { + auto* flight_data = reinterpret_cast(msg); + if (arrow::flight::transport::grpc::IsRegisteredGrpcFlightDataMessage(flight_data)) { + return GenericDeserialize( + buffer, msg); + } return arrow::flight::transport::grpc::FlightDataDeserialize( buffer, reinterpret_cast(msg)); } diff --git a/cpp/src/arrow/flight/transport/grpc/grpc_server.cc b/cpp/src/arrow/flight/transport/grpc/grpc_server.cc index 03f22f8864d4..49af881da924 100644 --- a/cpp/src/arrow/flight/transport/grpc/grpc_server.cc +++ b/cpp/src/arrow/flight/transport/grpc/grpc_server.cc @@ -34,6 +34,7 @@ #include "arrow/flight/server.h" #include "arrow/flight/server_middleware.h" #include "arrow/flight/transport.h" +#include "arrow/flight/transport/grpc/grpc_server_internal.h" #include "arrow/flight/transport/grpc/serialization_internal.h" #include "arrow/flight/transport/grpc/util_internal.h" #include "arrow/flight/transport_server.h" @@ -112,78 +113,8 @@ class GrpcServerAuthSender : public ServerAuthSender { ::grpc::ServerReaderWriter* stream_; }; -class GrpcServerCallContext : public ServerCallContext { - public: - explicit GrpcServerCallContext(::grpc::ServerContext* context) - : context_(context), peer_(context_->peer()) { - for (const auto& entry : context->client_metadata()) { - incoming_headers_.insert( - {std::string_view(entry.first.data(), entry.first.length()), - std::string_view(entry.second.data(), entry.second.length())}); - } - } - - const std::string& peer_identity() const override { return peer_identity_; } - const std::string& peer() const override { return peer_; } - bool is_cancelled() const override { return context_->IsCancelled(); } - const CallHeaders& incoming_headers() const override { return incoming_headers_; } - - // Helper method that runs interceptors given the result of an RPC, - // then returns the final gRPC status to send to the client - ::grpc::Status FinishRequest(const ::grpc::Status& status) { - // Don't double-convert status - return the original one here - FinishRequest(FromGrpcStatus(status)); - return status; - } - - ::grpc::Status FinishRequest(const arrow::Status& status) { - for (const auto& instance : middleware_) { - instance->CallCompleted(status); - } - - // Set custom headers to map the exact Arrow status for clients - // who want it. - return ToGrpcStatus(status, context_); - } - - void AddHeader(const std::string& key, const std::string& value) const override { - context_->AddInitialMetadata(key, value); - } - - void AddTrailer(const std::string& key, const std::string& value) const override { - context_->AddTrailingMetadata(key, value); - } - - ServerMiddleware* GetMiddleware(const std::string& key) const override { - const auto& instance = middleware_map_.find(key); - if (instance == middleware_map_.end()) { - return nullptr; - } - return instance->second.get(); - } - - private: - friend class GrpcServiceHandler; - ServerContext* context_; - std::string peer_; - std::string peer_identity_; - std::vector> middleware_; - std::unordered_map> middleware_map_; - CallHeaders incoming_headers_; -}; - -class GrpcAddServerHeaders : public AddCallHeaders { - public: - explicit GrpcAddServerHeaders(::grpc::ServerContext* context) : context_(context) {} - ~GrpcAddServerHeaders() override = default; - - void AddHeader(const std::string& key, const std::string& value) override { - context_->AddInitialMetadata(key, value); - } - - private: - ::grpc::ServerContext* context_; -}; +using GrpcServerCallContext = transport::grpc::GrpcServerCallContext; +using GrpcContextHelper = transport::grpc::GrpcServerCallContextHelper; // A ServerDataStream for streaming data to the client. class GetDataStream : public internal::ServerDataStream { @@ -248,7 +179,7 @@ class GrpcServiceHandler final : public FlightService::Service { std::vector>> middleware, internal::ServerTransport* impl) - : auth_handler_(auth_handler), middleware_(middleware), impl_(impl) {} + : helper_(std::move(auth_handler), std::move(middleware)), impl_(impl) {} template ::grpc::Status WriteStream(Iterator* iterator, ServerWriter* writer) { @@ -290,83 +221,15 @@ class GrpcServiceHandler final : public FlightService::Service { return ::grpc::Status::OK; } - // Authenticate the client (if applicable) and construct the call context - ::grpc::Status CheckAuth(const FlightMethod& method, ServerContext* context, - GrpcServerCallContext& flight_context, - bool skip_headers = false) { - if (!auth_handler_) { - const auto auth_context = context->auth_context(); - if (auth_context && auth_context->IsPeerAuthenticated()) { - auto peer_identity = auth_context->GetPeerIdentity(); - flight_context.peer_identity_ = - peer_identity.empty() - ? "" - : std::string(peer_identity.front().begin(), peer_identity.front().end()); - } else { - flight_context.peer_identity_ = ""; - } - } else { - const auto client_metadata = context->client_metadata(); - const auto [auth_header, auth_header_end] = - client_metadata.equal_range(kGrpcAuthHeader); - std::string token; - if (auth_header == auth_header_end) { - token = ""; - } else { - token = std::string(auth_header->second.data(), auth_header->second.length()); - } - GRPC_RETURN_NOT_OK( - auth_handler_->IsValid(flight_context, token, &flight_context.peer_identity_)); - } - - return MakeCallContext(method, context, flight_context); - } - - // Authenticate the client (if applicable) and construct the call context - ::grpc::Status MakeCallContext(const FlightMethod& method, ServerContext* context, - GrpcServerCallContext& flight_context, - bool skip_headers = false) { - // Run server middleware - const CallInfo info{method}; - - for (const auto& factory : middleware_) { - std::shared_ptr instance; - Status result = factory.second->StartCall(info, flight_context, &instance); - if (!result.ok()) { - // Interceptor rejected call, end the request on all existing - // interceptors - return flight_context.FinishRequest(result); - } - if (instance != nullptr) { - flight_context.middleware_.push_back(instance); - flight_context.middleware_map_.insert({factory.first, instance}); - } - } - - // TODO factor this out after fixing all streaming and non-streaming handlers - if (!skip_headers) { - addMiddlewareHeaders(context, flight_context); - } - - return ::grpc::Status::OK; - } - - void addMiddlewareHeaders(ServerContext* context, - GrpcServerCallContext& flight_context) { - GrpcAddServerHeaders outgoing_headers(context); - for (const std::shared_ptr& instance : flight_context.middleware_) { - instance->SendingHeaders(&outgoing_headers); - } - } - ::grpc::Status Handshake( ServerContext* context, ::grpc::ServerReaderWriter* stream) { GrpcServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( - MakeCallContext(FlightMethod::Handshake, context, flight_context)); + helper_.MakeCallContext(FlightMethod::Handshake, context, &flight_context)); + helper_.AddMiddlewareHeaders(context, &flight_context); - if (!auth_handler_) { + if (!helper_.auth_handler()) { RETURN_WITH_MIDDLEWARE( flight_context, ::grpc::Status( @@ -375,7 +238,7 @@ class GrpcServiceHandler final : public FlightService::Service { } GrpcServerAuthSender outgoing{stream}; GrpcServerAuthReader incoming{stream}; - RETURN_WITH_MIDDLEWARE(flight_context, auth_handler_->Authenticate( + RETURN_WITH_MIDDLEWARE(flight_context, helper_.auth_handler()->Authenticate( flight_context, &outgoing, &incoming)); } @@ -383,7 +246,8 @@ class GrpcServiceHandler final : public FlightService::Service { ServerWriter* writer) { GrpcServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( - CheckAuth(FlightMethod::ListFlights, context, flight_context)); + helper_.CheckAuth(FlightMethod::ListFlights, context, &flight_context)); + helper_.AddMiddlewareHeaders(context, &flight_context); // Retrieve the listing from the implementation std::unique_ptr listing; @@ -407,7 +271,7 @@ class GrpcServiceHandler final : public FlightService::Service { pb::FlightInfo* response) { GrpcServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( - CheckAuth(FlightMethod::GetFlightInfo, context, flight_context)); + helper_.CheckAuth(FlightMethod::GetFlightInfo, context, &flight_context)); CHECK_ARG_NOT_NULL(flight_context, request, "FlightDescriptor cannot be null"); @@ -415,8 +279,8 @@ class GrpcServiceHandler final : public FlightService::Service { SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &descr)); std::unique_ptr info; + helper_.AddMiddlewareHeaders(context, &flight_context); auto res = impl_->base()->GetFlightInfo(flight_context, descr, &info); - addMiddlewareHeaders(context, flight_context); SERVICE_RETURN_NOT_OK(flight_context, res); if (!info) { @@ -434,13 +298,14 @@ class GrpcServiceHandler final : public FlightService::Service { pb::PollInfo* response) { GrpcServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( - CheckAuth(FlightMethod::PollFlightInfo, context, flight_context)); + helper_.CheckAuth(FlightMethod::PollFlightInfo, context, &flight_context)); CHECK_ARG_NOT_NULL(flight_context, request, "FlightDescriptor cannot be null"); FlightDescriptor descr; SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &descr)); + helper_.AddMiddlewareHeaders(context, &flight_context); std::unique_ptr info; SERVICE_RETURN_NOT_OK(flight_context, impl_->base()->PollFlightInfo(flight_context, descr, &info)); @@ -458,13 +323,15 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status GetSchema(ServerContext* context, const pb::FlightDescriptor* request, pb::SchemaResult* response) { GrpcServerCallContext flight_context(context); - GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::GetSchema, context, flight_context)); + GRPC_RETURN_NOT_GRPC_OK( + helper_.CheckAuth(FlightMethod::GetSchema, context, &flight_context)); CHECK_ARG_NOT_NULL(flight_context, request, "FlightDescriptor cannot be null"); FlightDescriptor descr; SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &descr)); + helper_.AddMiddlewareHeaders(context, &flight_context); std::unique_ptr result; SERVICE_RETURN_NOT_OK(flight_context, impl_->base()->GetSchema(flight_context, descr, &result)); @@ -482,7 +349,9 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status DoGet(ServerContext* context, const pb::Ticket* request, ServerWriter* writer) { GrpcServerCallContext flight_context(context); - GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::DoGet, context, flight_context)); + GRPC_RETURN_NOT_GRPC_OK( + helper_.CheckAuth(FlightMethod::DoGet, context, &flight_context)); + helper_.AddMiddlewareHeaders(context, &flight_context); CHECK_ARG_NOT_NULL(flight_context, request, "ticket cannot be null"); @@ -498,7 +367,9 @@ class GrpcServiceHandler final : public FlightService::Service { ServerContext* context, ::grpc::ServerReaderWriter* reader) { GrpcServerCallContext flight_context(context); - GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::DoPut, context, flight_context)); + GRPC_RETURN_NOT_GRPC_OK( + helper_.CheckAuth(FlightMethod::DoPut, context, &flight_context)); + helper_.AddMiddlewareHeaders(context, &flight_context); PutDataStream stream(reader); RETURN_WITH_MIDDLEWARE(flight_context, impl_->DoPut(flight_context, &stream)); @@ -508,7 +379,9 @@ class GrpcServiceHandler final : public FlightService::Service { ServerContext* context, ::grpc::ServerReaderWriter* stream) { GrpcServerCallContext flight_context(context); - GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::DoExchange, context, flight_context)); + GRPC_RETURN_NOT_GRPC_OK( + helper_.CheckAuth(FlightMethod::DoExchange, context, &flight_context)); + helper_.AddMiddlewareHeaders(context, &flight_context); ExchangeDataStream data_stream(stream); RETURN_WITH_MIDDLEWARE(flight_context, @@ -519,7 +392,8 @@ class GrpcServiceHandler final : public FlightService::Service { ServerWriter* writer) { GrpcServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( - CheckAuth(FlightMethod::ListActions, context, flight_context)); + helper_.CheckAuth(FlightMethod::ListActions, context, &flight_context)); + helper_.AddMiddlewareHeaders(context, &flight_context); // Retrieve the listing from the implementation std::vector types; SERVICE_RETURN_NOT_OK(flight_context, @@ -530,7 +404,9 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status DoAction(ServerContext* context, const pb::Action* request, ServerWriter* writer) { GrpcServerCallContext flight_context(context); - GRPC_RETURN_NOT_GRPC_OK(CheckAuth(FlightMethod::DoAction, context, flight_context)); + GRPC_RETURN_NOT_GRPC_OK( + helper_.CheckAuth(FlightMethod::DoAction, context, &flight_context)); + helper_.AddMiddlewareHeaders(context, &flight_context); CHECK_ARG_NOT_NULL(flight_context, request, "Action cannot be null"); Action action; SERVICE_RETURN_NOT_OK(flight_context, internal::FromProto(*request, &action)); @@ -561,9 +437,7 @@ class GrpcServiceHandler final : public FlightService::Service { } private: - std::shared_ptr auth_handler_; - std::vector>> - middleware_; + GrpcContextHelper helper_; internal::ServerTransport* impl_; }; @@ -583,66 +457,17 @@ class GrpcServerTransport : public internal::ServerTransport { new GrpcServiceHandler(options.auth_handler, options.middleware, this)); ::grpc::ServerBuilder builder; - // Allow uploading messages of any length - builder.SetMaxReceiveMessageSize(-1); - - const std::string scheme = uri.scheme(); int port = 0; - if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp || scheme == kSchemeGrpcTls) { - std::stringstream address; - address << arrow::util::UriEncodeHost(uri.host()) << ':' << uri.port_text(); - - std::shared_ptr<::grpc::ServerCredentials> creds; - if (scheme == kSchemeGrpcTls) { - ::grpc::SslServerCredentialsOptions ssl_options; - for (const auto& pair : options.tls_certificates) { - ssl_options.pem_key_cert_pairs.push_back({pair.pem_key, pair.pem_cert}); - } - if (options.verify_client) { - ssl_options.client_certificate_request = - GRPC_SSL_REQUEST_AND_REQUIRE_CLIENT_CERTIFICATE_AND_VERIFY; - } - if (!options.root_certificates.empty()) { - ssl_options.pem_root_certs = options.root_certificates; - } - creds = ::grpc::SslServerCredentials(ssl_options); - } else { - creds = ::grpc::InsecureServerCredentials(); - } - - builder.AddListeningPort(address.str(), creds, &port); - } else if (scheme == kSchemeGrpcUnix) { - std::stringstream address; - address << "unix:" << uri.path(); - builder.AddListeningPort(address.str(), ::grpc::InsecureServerCredentials()); - location_ = options.location; - } else { - return Status::NotImplemented("Scheme is not supported: " + scheme); - } + RETURN_NOT_OK(AddServerListeningPort(options, uri, &builder, &location_, &port)); builder.RegisterService(grpc_service_.get()); - - // Disable SO_REUSEPORT - it makes debugging/testing a pain as - // leftover processes can handle requests on accident - builder.AddChannelArgument(GRPC_ARG_ALLOW_REUSEPORT, 0); - - if (options.builder_hook) { - options.builder_hook(&builder); - } + ConfigureServerBuilderOptions(options, &builder); grpc_server_ = builder.BuildAndStart(); if (!grpc_server_) { return Status::UnknownError("Server did not start properly"); } - - if (scheme == kSchemeGrpcTls) { - ARROW_ASSIGN_OR_RAISE( - location_, Location::ForGrpcTls(arrow::util::UriEncodeHost(uri.host()), port)); - } else if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp) { - ARROW_ASSIGN_OR_RAISE( - location_, Location::ForGrpcTcp(arrow::util::UriEncodeHost(uri.host()), port)); - } - return Status::OK(); + return SetServerLocationFromUri(uri, port, &location_); } Status Shutdown() override { grpc_server_->Shutdown(); diff --git a/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.cc b/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.cc new file mode 100644 index 000000000000..b28b6f88d9e6 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.cc @@ -0,0 +1,101 @@ +// 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. + +#include "arrow/flight/transport/grpc/grpc_server_internal.h" + +#include + +#include + +#include "arrow/flight/transport/grpc/customize_grpc.h" + +namespace arrow::flight::transport::grpc { + +::grpc::Status FinishGrpcServerStatus(const Status& arrow_status, + ::grpc::ServerContext* context) { + return ToGrpcStatus(arrow_status, context); +} + +::grpc::Status FinishGrpcServerStatus(const Status& arrow_status, + ::grpc::CallbackServerContext* context) { + return ToGrpcStatus(arrow_status, context); +} + +Status AddServerListeningPort(const FlightServerOptions& options, + const arrow::util::Uri& uri, ::grpc::ServerBuilder* builder, + Location* location, int* port) { + const std::string scheme = uri.scheme(); + if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp || scheme == kSchemeGrpcTls) { + std::stringstream address; + address << arrow::util::UriEncodeHost(uri.host()) << ':' << uri.port_text(); + + std::shared_ptr<::grpc::ServerCredentials> creds; + if (scheme == kSchemeGrpcTls) { + ::grpc::SslServerCredentialsOptions ssl_options; + for (const auto& pair : options.tls_certificates) { + ssl_options.pem_key_cert_pairs.push_back({pair.pem_key, pair.pem_cert}); + } + if (options.verify_client) { + ssl_options.client_certificate_request = + GRPC_SSL_REQUEST_AND_REQUIRE_CLIENT_CERTIFICATE_AND_VERIFY; + } + if (!options.root_certificates.empty()) { + ssl_options.pem_root_certs = options.root_certificates; + } + creds = ::grpc::SslServerCredentials(ssl_options); + } else { + creds = ::grpc::InsecureServerCredentials(); + } + builder->AddListeningPort(address.str(), creds, port); + return Status::OK(); + } + if (scheme == kSchemeGrpcUnix) { + std::stringstream address; + address << "unix:" << uri.path(); + builder->AddListeningPort(address.str(), ::grpc::InsecureServerCredentials()); + *location = options.location; + return Status::OK(); + } + return Status::NotImplemented("Scheme is not supported: " + scheme); +} + +void ConfigureServerBuilderOptions(const FlightServerOptions& options, + ::grpc::ServerBuilder* builder) { + // Allow uploading messages of any length + builder->SetMaxReceiveMessageSize(-1); + // Disable SO_REUSEPORT - it makes debugging/testing a pain as + // leftover processes can handle requests on accident + builder->AddChannelArgument(GRPC_ARG_ALLOW_REUSEPORT, 0); + if (options.builder_hook) { + options.builder_hook(builder); + } +} + +Status SetServerLocationFromUri(const arrow::util::Uri& uri, int port, + Location* location) { + const std::string scheme = uri.scheme(); + if (scheme == kSchemeGrpcTls) { + ARROW_ASSIGN_OR_RAISE( + *location, Location::ForGrpcTls(arrow::util::UriEncodeHost(uri.host()), port)); + } else if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp) { + ARROW_ASSIGN_OR_RAISE( + *location, Location::ForGrpcTcp(arrow::util::UriEncodeHost(uri.host()), port)); + } + return Status::OK(); +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.h b/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.h new file mode 100644 index 000000000000..6d827402bc5d --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.h @@ -0,0 +1,213 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include "arrow/flight/server.h" +#include "arrow/flight/server_middleware.h" +#include "arrow/flight/transport/grpc/util_internal.h" +#include "arrow/util/uri.h" + +namespace arrow::flight::transport::grpc { + +using MiddlewareFactoryList = + std::vector>>; + +ARROW_FLIGHT_EXPORT +::grpc::Status FinishGrpcServerStatus(const Status& arrow_status, + ::grpc::ServerContext* context); + +ARROW_FLIGHT_EXPORT +::grpc::Status FinishGrpcServerStatus(const Status& arrow_status, + ::grpc::CallbackServerContext* context); + +template +class GrpcAddServerHeaders : public AddCallHeaders { + public: + explicit GrpcAddServerHeaders(GrpcContext* context) : context_(context) {} + ~GrpcAddServerHeaders() override = default; + + void AddHeader(const std::string& key, const std::string& value) override { + context_->AddInitialMetadata(key, value); + } + + private: + GrpcContext* context_; +}; + +template +class GrpcServerCallContext : public ServerCallContext { + public: + explicit GrpcServerCallContext(GrpcContext* context) + : context_(context), peer_(context->peer()) { + for (const auto& entry : context->client_metadata()) { + incoming_headers_.insert( + {std::string_view(entry.first.data(), entry.first.length()), + std::string_view(entry.second.data(), entry.second.length())}); + } + } + + const std::string& peer_identity() const override { return peer_identity_; } + const std::string& peer() const override { return peer_; } + bool is_cancelled() const override { return context_->IsCancelled(); } + const CallHeaders& incoming_headers() const override { return incoming_headers_; } + + // Helper method that runs interceptors given the result of an RPC, + // then returns the final gRPC status to send to the client + ::grpc::Status FinishRequest(const ::grpc::Status& status) { + // Don't double-convert status - return the original one here + FinishRequest(FromGrpcStatus(status)); + return status; + } + + ::grpc::Status FinishRequest(const Status& status) { + for (const auto& instance : middleware_) { + instance->CallCompleted(status); + } + return FinishGrpcServerStatus(status, context_); + } + + void AddHeader(const std::string& key, const std::string& value) const override { + context_->AddInitialMetadata(key, value); + } + + void AddTrailer(const std::string& key, const std::string& value) const override { + context_->AddTrailingMetadata(key, value); + } + + ServerMiddleware* GetMiddleware(const std::string& key) const override { + const auto& instance = middleware_map_.find(key); + if (instance == middleware_map_.end()) { + return nullptr; + } + return instance->second.get(); + } + + private: + template + friend class GrpcServerCallContextHelper; + + GrpcContext* context_; + std::string peer_; + std::string peer_identity_; + std::vector> middleware_; + std::unordered_map> middleware_map_; + CallHeaders incoming_headers_; +}; + +template +class GrpcServerCallContextHelper { + public: + GrpcServerCallContextHelper(std::shared_ptr auth_handler, + MiddlewareFactoryList middleware) + : auth_handler_(std::move(auth_handler)), middleware_(std::move(middleware)) {} + + // Authenticate the client (if applicable) and construct the call context + ::grpc::Status MakeCallContext(FlightMethod method, GrpcContext* context, + GrpcServerCallContext* flight_context, + bool skip_headers = false) const { + // Run server middleware + const CallInfo info{method}; + for (const auto& factory : middleware_) { + std::shared_ptr instance; + Status result = factory.second->StartCall(info, *flight_context, &instance); + if (!result.ok()) { + // Interceptor rejected call, end the request on all existing + // interceptors + return flight_context->FinishRequest(result); + } + if (instance != nullptr) { + flight_context->middleware_.push_back(instance); + flight_context->middleware_map_.insert({factory.first, instance}); + } + } + // TODO factor this out after fixing all streaming and non-streaming handlers + if (!skip_headers) { + AddMiddlewareHeaders(context, flight_context); + } + + return ::grpc::Status::OK; + } + + void AddMiddlewareHeaders(GrpcContext* context, + GrpcServerCallContext* flight_context) const { + GrpcAddServerHeaders outgoing_headers(context); + for (const auto& instance : flight_context->middleware_) { + instance->SendingHeaders(&outgoing_headers); + } + } + + // Authenticate the client (if applicable) and construct the call context + ::grpc::Status CheckAuth(FlightMethod method, GrpcContext* context, + GrpcServerCallContext* flight_context) const { + if (!auth_handler_) { + const auto auth_context = context->auth_context(); + if (auth_context && auth_context->IsPeerAuthenticated()) { + auto peer_identity = auth_context->GetPeerIdentity(); + flight_context->peer_identity_ = + peer_identity.empty() + ? "" + : std::string(peer_identity.front().begin(), peer_identity.front().end()); + } else { + flight_context->peer_identity_ = ""; + } + } else { + const auto client_metadata = context->client_metadata(); + const auto [auth_header, auth_header_end] = + client_metadata.equal_range(kGrpcAuthHeader); + std::string token; + if (auth_header != auth_header_end) { + token = std::string(auth_header->second.data(), auth_header->second.length()); + } + auto auth_status = + auth_handler_->IsValid(*flight_context, token, &flight_context->peer_identity_); + if (!auth_status.ok()) { + return flight_context->FinishRequest(auth_status); + } + } + return MakeCallContext(method, context, flight_context); + } + + const std::shared_ptr& auth_handler() const { return auth_handler_; } + + private: + std::shared_ptr auth_handler_; + MiddlewareFactoryList middleware_; +}; + +ARROW_FLIGHT_EXPORT +Status AddServerListeningPort(const FlightServerOptions& options, + const arrow::util::Uri& uri, ::grpc::ServerBuilder* builder, + Location* location, int* port); + +ARROW_FLIGHT_EXPORT +void ConfigureServerBuilderOptions(const FlightServerOptions& options, + ::grpc::ServerBuilder* builder); + +ARROW_FLIGHT_EXPORT +Status SetServerLocationFromUri(const arrow::util::Uri& uri, int port, + Location* location); + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/serialization_internal.cc b/cpp/src/arrow/flight/transport/grpc/serialization_internal.cc index ee13a8e2022f..d0914d016490 100644 --- a/cpp/src/arrow/flight/transport/grpc/serialization_internal.cc +++ b/cpp/src/arrow/flight/transport/grpc/serialization_internal.cc @@ -19,7 +19,9 @@ #include #include +#include #include +#include #include #include "arrow/flight/platform.h" @@ -56,6 +58,16 @@ using ::grpc::ByteBuffer; namespace { +std::mutex& FlightDataRegistryMutex() { + static std::mutex mutex; + return mutex; +} + +std::unordered_set& FlightDataRegistry() { + static auto* registry = new std::unordered_set(); + return *registry; +} + // Internal wrapper for gRPC ByteBuffer so its memory can be exposed to Arrow // consumers with zero-copy class GrpcBuffer : public MutableBuffer { @@ -151,6 +163,21 @@ arrow::Result<::grpc::Slice> SliceFromBuffer(const std::shared_ptr& buf) } // namespace +bool IsRegisteredGrpcFlightDataMessage(const pb::FlightData* message) { + std::lock_guard guard(FlightDataRegistryMutex()); + return FlightDataRegistry().find(message) != FlightDataRegistry().end(); +} + +void RegisterGrpcFlightDataMessage(const pb::FlightData* message) { + std::lock_guard guard(FlightDataRegistryMutex()); + FlightDataRegistry().insert(message); +} + +void UnregisterGrpcFlightDataMessage(const pb::FlightData* message) { + std::lock_guard guard(FlightDataRegistryMutex()); + FlightDataRegistry().erase(message); +} + ::grpc::Status FlightDataSerialize(const FlightPayload& msg, ByteBuffer* out, bool* own_buffer) { // Retrieve BufferVector from the transport-agnostic serialization. diff --git a/cpp/src/arrow/flight/transport/grpc/util_internal.cc b/cpp/src/arrow/flight/transport/grpc/util_internal.cc index 95b205d67808..78908b6825f9 100644 --- a/cpp/src/arrow/flight/transport/grpc/util_internal.cc +++ b/cpp/src/arrow/flight/transport/grpc/util_internal.cc @@ -312,7 +312,8 @@ static ::grpc::Status ToRawGrpcStatus(const Status& arrow_status) { /// Convert an Arrow status to a gRPC status, and add extra headers to /// the response to encode the original Arrow status. -::grpc::Status ToGrpcStatus(const Status& arrow_status, ::grpc::ServerContext* ctx) { +template +::grpc::Status ToGrpcStatusImpl(const Status& arrow_status, GrpcContext* ctx) { ::grpc::Status status = ToRawGrpcStatus(arrow_status); if (!status.ok() && ctx) { const std::string code = ToChars(static_cast(arrow_status.code())); @@ -331,6 +332,19 @@ ::grpc::Status ToGrpcStatus(const Status& arrow_status, ::grpc::ServerContext* c return status; } +::grpc::Status ToGrpcStatus(const Status& arrow_status, ::grpc::ServerContext* ctx) { + return ToGrpcStatusImpl(arrow_status, ctx); +} + +::grpc::Status ToGrpcStatus(const Status& arrow_status, + ::grpc::CallbackServerContext* ctx) { + ::grpc::Status status = ToGrpcStatusImpl(arrow_status, ctx); + if (status.error_code() == ::grpc::StatusCode::UNKNOWN && arrow_status.IsIOError()) { + status = ::grpc::Status(::grpc::StatusCode::INTERNAL, arrow_status.message()); + } + return status; +} + #if GRPC_CPP_VERSION_CHECK(1, 80, 0) Status FromAbslStatus(const ::absl::Status& absl_status) { switch (absl_status.code()) { diff --git a/cpp/src/arrow/flight/transport/grpc/util_internal.h b/cpp/src/arrow/flight/transport/grpc/util_internal.h index 7e306d574a58..d023a6d99074 100644 --- a/cpp/src/arrow/flight/transport/grpc/util_internal.h +++ b/cpp/src/arrow/flight/transport/grpc/util_internal.h @@ -17,28 +17,11 @@ #pragma once -#include - #include "arrow/flight/transport/grpc/protocol_grpc_internal.h" #include "arrow/flight/types.h" #include "arrow/flight/visibility.h" #include "arrow/util/macros.h" -// gRPC 1.51.0 or later defines GRPC_CPP_VERSION_MAJOR and so on. -#ifdef GRPC_CPP_VERSION_MAJOR -# define GRPC_CPP_VERSION_CHECK(major, minor, patch) \ - ((GRPC_CPP_VERSION_MAJOR > (major) || \ - (GRPC_CPP_VERSION_MAJOR == (major) && GRPC_CPP_VERSION_MINOR > (minor)) || \ - ((GRPC_CPP_VERSION_MAJOR == (major) && GRPC_CPP_VERSION_MINOR == (minor) && \ - GRPC_CPP_VERSION_PATCH >= (patch))))) -#else -# define GRPC_CPP_VERSION_CHECK(major, minor, patch) 0 -#endif - -#if GRPC_CPP_VERSION_CHECK(1, 80, 0) -# include -#endif - namespace grpc { class Status; @@ -51,6 +34,12 @@ class Status; namespace flight { +#define GRPC_CPP_VERSION_CHECK(major, minor, patch) \ + ((GRPC_CPP_VERSION_MAJOR > (major) || \ + (GRPC_CPP_VERSION_MAJOR == (major) && GRPC_CPP_VERSION_MINOR > (minor)) || \ + ((GRPC_CPP_VERSION_MAJOR == (major) && GRPC_CPP_VERSION_MINOR == (minor) && \ + GRPC_CPP_VERSION_PATCH >= (patch))))) + #define GRPC_RETURN_NOT_OK(expr) \ do { \ ::arrow::Status _s = (expr); \ @@ -107,6 +96,10 @@ ARROW_FLIGHT_EXPORT ::grpc::Status ToGrpcStatus(const Status& arrow_status, ::grpc::ServerContext* ctx = nullptr); +ARROW_FLIGHT_EXPORT +::grpc::Status ToGrpcStatus(const Status& arrow_status, + ::grpc::CallbackServerContext* ctx); + // gRPC 1.80.0 or later use absl::Status. #if GRPC_CPP_VERSION_CHECK(1, 80, 0) /// Convert an Abseil status to an Arrow status. diff --git a/cpp/src/arrow/flight/transport_server.cc b/cpp/src/arrow/flight/transport_server.cc index 197def7cabdc..df4e9fbabe26 100644 --- a/cpp/src/arrow/flight/transport_server.cc +++ b/cpp/src/arrow/flight/transport_server.cc @@ -17,16 +17,8 @@ #include "arrow/flight/transport_server.h" -#include - -#include "arrow/buffer.h" -#include "arrow/flight/serialization_internal.h" #include "arrow/flight/server.h" -#include "arrow/flight/types.h" -#include "arrow/ipc/reader.h" -#include "arrow/result.h" -#include "arrow/status.h" -#include "arrow/util/logging_internal.h" +#include "arrow/flight/transport_server_internal.h" namespace arrow { namespace flight { @@ -36,297 +28,17 @@ Status ServerDataStream::WritePutMetadata(const Buffer&) { return Status::NotImplemented("Writing put metadata for this stream"); } -namespace { -class TransportIpcMessageReader : public ipc::MessageReader { - public: - TransportIpcMessageReader( - std::shared_ptr peekable_reader, - std::shared_ptr memory_manager, - std::shared_ptr* app_metadata) - : peekable_reader_(peekable_reader), - memory_manager_(std::move(memory_manager)), - app_metadata_(app_metadata) {} - - ::arrow::Result> ReadNextMessage() override { - if (stream_finished_) return nullptr; - internal::FlightData* data; - peekable_reader_->Next(&data); - if (!data) { - stream_finished_ = true; - return nullptr; - } - if (data->body) { - ARROW_ASSIGN_OR_RAISE(data->body, Buffer::ViewOrCopy(data->body, memory_manager_)); - } - *app_metadata_ = std::move(data->app_metadata); - return data->OpenMessage(); - } - - protected: - std::shared_ptr peekable_reader_; - std::shared_ptr memory_manager_; - // A reference to TransportDataStream.app_metadata_. That class - // can't access the app metadata because when it Peek()s the stream, - // it may be looking at a dictionary batch, not the record - // batch. Updating it here ensures the reader is always updated with - // the last metadata message read. - std::shared_ptr* app_metadata_; - bool stream_finished_ = false; -}; - -/// \brief Adapt TransportDataStream to the FlightMessageReader -/// interface for DoPut. -class TransportMessageReader final : public FlightMessageReader { - public: - TransportMessageReader(ServerDataStream* stream, - std::shared_ptr memory_manager) - : peekable_reader_(new internal::PeekableFlightDataReader(stream)), - memory_manager_(std::move(memory_manager)) {} - - Status Init() { - // Peek the first message to get the descriptor. - internal::FlightData* data; - peekable_reader_->Peek(&data); - if (!data) { - return Status::IOError("Stream finished before first message sent"); - } - if (!data->descriptor) { - return Status::IOError("Descriptor missing on first message"); - } - descriptor_ = *data->descriptor; - // If there's a schema (=DoPut), also Open(). - if (data->metadata) { - return EnsureDataStarted(); - } - peekable_reader_->Next(&data); - return Status::OK(); - } - - const FlightDescriptor& descriptor() const override { return descriptor_; } - - arrow::Result> GetSchema() override { - RETURN_NOT_OK(EnsureDataStarted()); - return batch_reader_->schema(); - } - - arrow::Result Next() override { - FlightStreamChunk out; - internal::FlightData* data = nullptr; - peekable_reader_->Peek(&data); - if (!data) { - out.app_metadata = nullptr; - out.data = nullptr; - return out; - } - - if (!data->metadata) { - // Metadata-only (data->metadata is the IPC header) - out.app_metadata = data->app_metadata; - out.data = nullptr; - peekable_reader_->Next(&data); - return out; - } - - if (!batch_reader_) { - RETURN_NOT_OK(EnsureDataStarted()); - // re-peek here since EnsureDataStarted() advances the stream - return Next(); - } - RETURN_NOT_OK(batch_reader_->ReadNext(&out.data)); - out.app_metadata = std::move(app_metadata_); - return out; - } - - ipc::ReadStats stats() const override { - if (batch_reader_ == nullptr) { - return ipc::ReadStats{}; - } - return batch_reader_->stats(); - } - - private: - /// Ensure we are set up to read data. - Status EnsureDataStarted() { - if (!batch_reader_) { - // peek() until we find the first data message; discard metadata - if (!peekable_reader_->SkipToData()) { - return Status::IOError("Client never sent a data message"); - } - auto message_reader = - std::unique_ptr(new TransportIpcMessageReader( - peekable_reader_, memory_manager_, &app_metadata_)); - ARROW_ASSIGN_OR_RAISE( - batch_reader_, ipc::RecordBatchStreamReader::Open(std::move(message_reader))); - } - return Status::OK(); - } - - FlightDescriptor descriptor_; - std::shared_ptr peekable_reader_; - std::shared_ptr memory_manager_; - std::shared_ptr batch_reader_; - std::shared_ptr app_metadata_; -}; - -/// \brief An IpcPayloadWriter for ServerDataStream. -/// -/// To support app_metadata and reuse the existing IPC infrastructure, -/// this takes a pointer to a buffer to be combined with the IPC -/// payload when writing a Flight payload. -class TransportMessagePayloadWriter : public ipc::internal::IpcPayloadWriter { - public: - TransportMessagePayloadWriter(ServerDataStream* stream, - std::shared_ptr* app_metadata) - : stream_(stream), app_metadata_(app_metadata) {} - - Status Start() override { return Status::OK(); } - Status WritePayload(const ipc::IpcPayload& ipc_payload) override { - FlightPayload payload; - payload.ipc_message = ipc_payload; - - if (ipc_payload.type == ipc::MessageType::RECORD_BATCH && *app_metadata_) { - payload.app_metadata = std::move(*app_metadata_); - } - ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload)); - if (!success) { - return MakeFlightError( - FlightStatusCode::Internal, - "Could not write record batch to stream (client disconnect?)"); - } - return arrow::Status::OK(); - } - Status Close() override { - // Closing is handled one layer up in TransportMessageWriter::Close - return Status::OK(); - } - - private: - ServerDataStream* stream_; - std::shared_ptr* app_metadata_; -}; - -class TransportMessageWriter final : public FlightMessageWriter { - public: - explicit TransportMessageWriter(ServerDataStream* stream) - : stream_(stream), - app_metadata_(nullptr), - ipc_options_(::arrow::ipc::IpcWriteOptions::Defaults()) {} - - Status Begin(const std::shared_ptr& schema, - const ipc::IpcWriteOptions& options) override { - if (batch_writer_) { - return Status::Invalid("This writer has already been started."); - } - ipc_options_ = options; - std::unique_ptr payload_writer( - new TransportMessagePayloadWriter(stream_, &app_metadata_)); - - ARROW_ASSIGN_OR_RAISE(batch_writer_, - ipc::internal::OpenRecordBatchWriter(std::move(payload_writer), - schema, ipc_options_)); - return Status::OK(); - } - - Status WriteRecordBatch(const RecordBatch& batch) override { - return WriteWithMetadata(batch, nullptr); - } - - Status WriteMetadata(std::shared_ptr app_metadata) override { - FlightPayload payload{}; - payload.app_metadata = app_metadata; - ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload)); - if (!success) { - ARROW_RETURN_NOT_OK(Close()); - return MakeFlightError(FlightStatusCode::Internal, - "Could not write metadata to stream (client disconnect?)"); - } - return Status::OK(); - } - - Status WriteWithMetadata(const RecordBatch& batch, - std::shared_ptr app_metadata) override { - RETURN_NOT_OK(CheckStarted()); - app_metadata_ = app_metadata; - auto status = batch_writer_->WriteRecordBatch(batch); - if (!status.ok()) { - ARROW_RETURN_NOT_OK(Close()); - } - return status; - } - - Status Close() override { - if (batch_writer_) { - RETURN_NOT_OK(batch_writer_->Close()); - } - return Status::OK(); - } - - ipc::WriteStats stats() const override { - ARROW_CHECK_NE(batch_writer_, nullptr); - return batch_writer_->stats(); - } - - private: - Status CheckStarted() { - if (!batch_writer_) { - return Status::Invalid("This writer is not started. Call Begin() with a schema"); - } - return Status::OK(); - } - - ServerDataStream* stream_; - std::unique_ptr batch_writer_; - std::shared_ptr app_metadata_; - ::arrow::ipc::IpcWriteOptions ipc_options_; -}; - -/// \brief Adapt TransportDataStream to the FlightMetadataWriter -/// interface for DoPut. -class TransportMetadataWriter final : public FlightMetadataWriter { - public: - explicit TransportMetadataWriter(ServerDataStream* stream) : stream_(stream) {} - - Status WriteMetadata(const Buffer& buffer) override { - return stream_->WritePutMetadata(buffer); - } - - private: - ServerDataStream* stream_; -}; -} // namespace - Status ServerTransport::DoGet(const ServerCallContext& context, const Ticket& ticket, ServerDataStream* stream) { std::unique_ptr data_stream; RETURN_NOT_OK(base_->DoGet(context, ticket, &data_stream)); - - if (!data_stream) return Status::KeyError("No data in this flight"); - - // Write the schema as the first message in the stream - ARROW_ASSIGN_OR_RAISE(auto schema_payload, data_stream->GetSchemaPayload()); - ARROW_ASSIGN_OR_RAISE(auto success, stream->WriteData(schema_payload)); - // Connection terminated - if (!success) return Status::OK(); - - // Consume data stream and write out payloads - while (true) { - ARROW_ASSIGN_OR_RAISE(FlightPayload payload, data_stream->Next()); - // End of stream - if (payload.ipc_message.metadata == nullptr) break; - ARROW_ASSIGN_OR_RAISE(auto success, stream->WriteData(payload)); - // Connection terminated - if (!success) return Status::OK(); - } - RETURN_NOT_OK(stream->WritesDone()); - return data_stream->Close(); + return WriteDataStream(std::move(data_stream), stream); } Status ServerTransport::DoPut(const ServerCallContext& context, ServerDataStream* stream) { - std::unique_ptr reader( - new TransportMessageReader(stream, memory_manager_)); - std::unique_ptr writer(new TransportMetadataWriter(stream)); - RETURN_NOT_OK(reader->Init()); + ARROW_ASSIGN_OR_RAISE(auto reader, MakeMessageReader(stream)); + auto writer = MakeMetadataWriter(stream); RETURN_NOT_OK(base_->DoPut(context, std::move(reader), std::move(writer))); RETURN_NOT_OK(stream->WritesDone()); return Status::OK(); @@ -334,10 +46,8 @@ Status ServerTransport::DoPut(const ServerCallContext& context, Status ServerTransport::DoExchange(const ServerCallContext& context, ServerDataStream* stream) { - std::unique_ptr reader( - new TransportMessageReader(stream, memory_manager_)); - std::unique_ptr writer(new TransportMessageWriter(stream)); - RETURN_NOT_OK(reader->Init()); + ARROW_ASSIGN_OR_RAISE(auto reader, MakeMessageReader(stream)); + auto writer = MakeMessageWriter(stream); RETURN_NOT_OK(base_->DoExchange(context, std::move(reader), std::move(writer))); RETURN_NOT_OK(stream->WritesDone()); return Status::OK(); diff --git a/cpp/src/arrow/flight/transport_server.h b/cpp/src/arrow/flight/transport_server.h index 8e5fe3e710c1..6d59c63fa6b2 100644 --- a/cpp/src/arrow/flight/transport_server.h +++ b/cpp/src/arrow/flight/transport_server.h @@ -21,6 +21,7 @@ #include #include "arrow/flight/transport.h" +#include "arrow/flight/transport_server_internal.h" #include "arrow/flight/type_fwd.h" #include "arrow/flight/visibility.h" #include "arrow/type_fwd.h" @@ -54,10 +55,10 @@ class ARROW_FLIGHT_EXPORT ServerDataStream : public TransportDataStream { /// application RPC method handlers. /// /// Used by FlightServerBase to manage the server lifecycle. -class ARROW_FLIGHT_EXPORT ServerTransport { +class ARROW_FLIGHT_EXPORT ServerTransport : public ServerTransportBase { public: ServerTransport(FlightServerBase* base, std::shared_ptr memory_manager) - : base_(base), memory_manager_(std::move(memory_manager)) {} + : ServerTransportBase(std::move(memory_manager)), base_(base) {} virtual ~ServerTransport() = default; /// \name Server Lifecycle Methods @@ -125,7 +126,6 @@ class ARROW_FLIGHT_EXPORT ServerTransport { protected: FlightServerBase* base_; - std::shared_ptr memory_manager_; }; } // namespace internal diff --git a/cpp/src/arrow/flight/transport_server_internal.cc b/cpp/src/arrow/flight/transport_server_internal.cc new file mode 100644 index 000000000000..c243e6dfe404 --- /dev/null +++ b/cpp/src/arrow/flight/transport_server_internal.cc @@ -0,0 +1,384 @@ +// 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. + +#include "arrow/flight/transport_server_internal.h" + +#include "arrow/buffer.h" +#include "arrow/flight/serialization_internal.h" +#include "arrow/flight/transport_server.h" +#include "arrow/ipc/reader.h" + +namespace arrow::flight::internal { +namespace { + +class TransportIpcMessageReader : public ipc::MessageReader { + public: + TransportIpcMessageReader(std::shared_ptr peekable_reader, + std::shared_ptr memory_manager, + std::shared_ptr* app_metadata) + : peekable_reader_(std::move(peekable_reader)), + memory_manager_(std::move(memory_manager)), + app_metadata_(app_metadata) {} + + arrow::Result> ReadNextMessage() override { + if (stream_finished_) return nullptr; + FlightData* data; + peekable_reader_->Next(&data); + if (!data) { + stream_finished_ = true; + return nullptr; + } + if (data->body) { + ARROW_ASSIGN_OR_RAISE(data->body, Buffer::ViewOrCopy(data->body, memory_manager_)); + } + *app_metadata_ = std::move(data->app_metadata); + return data->OpenMessage(); + } + + private: + std::shared_ptr peekable_reader_; + std::shared_ptr memory_manager_; + // A reference to TransportDataStream.app_metadata_. That class + // can't access the app metadata because when it Peek()s the stream, + // it may be looking at a dictionary batch, not the record + // batch. Updating it here ensures the reader is always updated with + // the last metadata message read. + std::shared_ptr* app_metadata_; + bool stream_finished_ = false; +}; + +/// \brief Adapt TransportDataStream to the FlightMessageReader +/// interface for DoPut. +class TransportMessageReader final : public FlightMessageReader { + public: + TransportMessageReader(ServerDataStream* stream, + std::shared_ptr memory_manager) + : peekable_reader_(new PeekableFlightDataReader(stream)), + memory_manager_(std::move(memory_manager)) {} + + Status Init() { + // Peek the first message to get the descriptor. + FlightData* data; + peekable_reader_->Peek(&data); + if (!data) { + return Status::IOError("Stream finished before first message sent"); + } + if (!data->descriptor) { + return Status::IOError("Descriptor missing on first message"); + } + descriptor_ = *data->descriptor; + // If there's a schema (=DoPut), also Open(). + if (data->metadata) { + return EnsureDataStarted(); + } + peekable_reader_->Next(&data); + return Status::OK(); + } + + const FlightDescriptor& descriptor() const override { return descriptor_; } + + arrow::Result> GetSchema() override { + RETURN_NOT_OK(EnsureDataStarted()); + return batch_reader_->schema(); + } + + arrow::Result Next() override { + FlightStreamChunk out; + FlightData* data = nullptr; + peekable_reader_->Peek(&data); + if (!data) { + out.app_metadata = nullptr; + out.data = nullptr; + return out; + } + + if (!data->metadata) { + // Metadata-only (data->metadata is the IPC header) + out.app_metadata = data->app_metadata; + out.data = nullptr; + peekable_reader_->Next(&data); + return out; + } + + if (!batch_reader_) { + RETURN_NOT_OK(EnsureDataStarted()); + // re-peek here since EnsureDataStarted() advances the stream + return Next(); + } + RETURN_NOT_OK(batch_reader_->ReadNext(&out.data)); + out.app_metadata = std::move(app_metadata_); + return out; + } + + ipc::ReadStats stats() const override { + if (batch_reader_ == nullptr) { + return ipc::ReadStats{}; + } + return batch_reader_->stats(); + } + + private: + /// Ensure we are set up to read data. + Status EnsureDataStarted() { + if (!batch_reader_) { + // peek() until we find the first data message; discard metadata + if (!peekable_reader_->SkipToData()) { + return Status::IOError("Client never sent a data message"); + } + auto message_reader = + std::unique_ptr(new TransportIpcMessageReader( + peekable_reader_, memory_manager_, &app_metadata_)); + ARROW_ASSIGN_OR_RAISE( + batch_reader_, ipc::RecordBatchStreamReader::Open(std::move(message_reader))); + } + return Status::OK(); + } + + FlightDescriptor descriptor_; + std::shared_ptr peekable_reader_; + std::shared_ptr memory_manager_; + std::shared_ptr batch_reader_; + std::shared_ptr app_metadata_; +}; + +/// \brief An IpcPayloadWriter for ServerDataStream. +/// +/// To support app_metadata and reuse the existing IPC infrastructure, +/// this takes a pointer to a buffer to be combined with the IPC +/// payload when writing a Flight payload. +class TransportMessagePayloadWriter : public ipc::internal::IpcPayloadWriter { + public: + TransportMessagePayloadWriter(ServerDataStream* stream, + std::shared_ptr* app_metadata) + : stream_(stream), app_metadata_(app_metadata) {} + + Status Start() override { return Status::OK(); } + + Status WritePayload(const ipc::IpcPayload& ipc_payload) override { + FlightPayload payload; + payload.ipc_message = ipc_payload; + if (ipc_payload.type == ipc::MessageType::RECORD_BATCH && *app_metadata_) { + payload.app_metadata = std::move(*app_metadata_); + } + ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload)); + if (!success) { + return MakeFlightError( + FlightStatusCode::Internal, + "Could not write record batch to stream (client disconnect?)"); + } + return Status::OK(); + } + + Status Close() override { + // Closing is handled one layer up in TransportMessageWriter::Close + return Status::OK(); + } + + private: + ServerDataStream* stream_; + std::shared_ptr* app_metadata_; +}; + +class TransportMessageWriter final : public FlightMessageWriter { + public: + explicit TransportMessageWriter(ServerDataStream* stream) + : stream_(stream), + app_metadata_(nullptr), + ipc_options_(::arrow::ipc::IpcWriteOptions::Defaults()) {} + + Status Begin(const std::shared_ptr& schema, + const ipc::IpcWriteOptions& options) override { + if (batch_writer_) { + return Status::Invalid("This writer has already been started."); + } + ipc_options_ = options; + std::unique_ptr payload_writer( + new TransportMessagePayloadWriter(stream_, &app_metadata_)); + ARROW_ASSIGN_OR_RAISE(batch_writer_, + ipc::internal::OpenRecordBatchWriter(std::move(payload_writer), + schema, ipc_options_)); + return Status::OK(); + } + + Status WriteRecordBatch(const RecordBatch& batch) override { + return WriteWithMetadata(batch, nullptr); + } + + Status WriteMetadata(std::shared_ptr app_metadata) override { + FlightPayload payload{}; + payload.app_metadata = std::move(app_metadata); + ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload)); + if (!success) { + ARROW_RETURN_NOT_OK(Close()); + return MakeFlightError(FlightStatusCode::Internal, + "Could not write metadata to stream (client disconnect?)"); + } + return Status::OK(); + } + + Status WriteWithMetadata(const RecordBatch& batch, + std::shared_ptr app_metadata) override { + RETURN_NOT_OK(CheckStarted()); + app_metadata_ = std::move(app_metadata); + auto status = batch_writer_->WriteRecordBatch(batch); + if (!status.ok()) { + ARROW_RETURN_NOT_OK(Close()); + } + return status; + } + + Status Close() override { + if (batch_writer_) { + RETURN_NOT_OK(batch_writer_->Close()); + } + return Status::OK(); + } + + ipc::WriteStats stats() const override { + ARROW_CHECK_NE(batch_writer_, nullptr); + return batch_writer_->stats(); + } + + private: + Status CheckStarted() { + if (!batch_writer_) { + return Status::Invalid("This writer is not started. Call Begin() with a schema"); + } + return Status::OK(); + } + + ServerDataStream* stream_; + std::unique_ptr batch_writer_; + std::shared_ptr app_metadata_; + ::arrow::ipc::IpcWriteOptions ipc_options_; +}; + +/// \brief Adapt TransportDataStream to the FlightMetadataWriter +/// interface for DoPut. +class TransportMetadataWriter final : public FlightMetadataWriter { + public: + explicit TransportMetadataWriter(ServerDataStream* stream) : stream_(stream) {} + + Status WriteMetadata(const Buffer& buffer) override { + return stream_->WritePutMetadata(buffer); + } + + private: + ServerDataStream* stream_; +}; + +} // namespace + +std::atomic ServerSignalState::running_instance_; + +Status ServerSignalState::SetShutdownOnSignals(const std::vector& sigs) { + signals_ = sigs; + old_signal_handlers_.clear(); + return Status::OK(); +} + +int ServerSignalState::GotSignal() const { return got_signal_; } + +void ServerSignalState::HandleSignal(int signum) { + auto instance = running_instance_.load(); + if (instance != nullptr) { + instance->DoHandleSignal(signum); + } +} + +void ServerSignalState::DoHandleSignal(int signum) { + got_signal_ = signum; + // Send dummy payload over self-pipe + self_pipe_->Send(/*payload=*/0); +} + +void ServerSignalState::WaitForSignals( + std::shared_ptr<::arrow::internal::SelfPipe> self_pipe) { + // Wait for a signal handler to wake up the pipe + auto st = self_pipe->Wait().status(); + // Status::Invalid means the pipe was shutdown without any wakeup + if (!st.ok() && !st.IsInvalid()) { + ARROW_LOG(FATAL) << "Failed to wait on self-pipe: " << st.ToString(); + } + auto instance = running_instance_.load(); + if (instance != nullptr && instance->shutdown_) { + ARROW_WARN_NOT_OK(instance->shutdown_(), instance->shutdown_warning_); + } +} + +arrow::Result ParseLocationUri(const Location& location) { + return arrow::util::Uri::FromString(location.ToString()); +} + +int PortFromLocation(const Location& location) { + auto maybe_uri = ParseLocationUri(location); + if (!maybe_uri.ok()) { + return -1; + } + return maybe_uri.ValueUnsafe().port(); +} + +arrow::Result> +ServerTransportBase::MakeMessageReader(ServerDataStream* stream) const { + auto reader = std::unique_ptr( + new TransportMessageReader(stream, memory_manager_)); + RETURN_NOT_OK(reader->Init()); + return std::unique_ptr(std::move(reader)); +} + +std::unique_ptr ServerTransportBase::MakeMetadataWriter( + ServerDataStream* stream) const { + return std::unique_ptr(new TransportMetadataWriter(stream)); +} + +std::unique_ptr ServerTransportBase::MakeMessageWriter( + ServerDataStream* stream) const { + return std::unique_ptr(new TransportMessageWriter(stream)); +} + +Status ServerTransportBase::WriteDataStream(std::unique_ptr data_stream, + ServerDataStream* stream) const { + if (!data_stream) { + return Status::KeyError("No data in this flight"); + } + + // Write the schema as the first message in the stream + ARROW_ASSIGN_OR_RAISE(auto schema_payload, data_stream->GetSchemaPayload()); + ARROW_ASSIGN_OR_RAISE(auto success, stream->WriteData(schema_payload)); + // Connection terminated + if (!success) { + return Status::OK(); + } + + // Consume data stream and write out payloads + while (true) { + ARROW_ASSIGN_OR_RAISE(FlightPayload payload, data_stream->Next()); + // End of stream + if (payload.ipc_message.metadata == nullptr) { + break; + } + ARROW_ASSIGN_OR_RAISE(success, stream->WriteData(payload)); + // Connection terminated + if (!success) { + return Status::OK(); + } + } + RETURN_NOT_OK(stream->WritesDone()); + return data_stream->Close(); +} + +} // namespace arrow::flight::internal diff --git a/cpp/src/arrow/flight/transport_server_internal.h b/cpp/src/arrow/flight/transport_server_internal.h new file mode 100644 index 000000000000..5552b788737e --- /dev/null +++ b/cpp/src/arrow/flight/transport_server_internal.h @@ -0,0 +1,250 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/flight/platform.h" +#include "arrow/flight/server.h" +#include "arrow/flight/visibility.h" +#include "arrow/result.h" +#include "arrow/status.h" +#include "arrow/util/io_util.h" +#include "arrow/util/logging.h" +#include "arrow/util/uri.h" + +namespace arrow::flight::internal { + +class ServerDataStream; + +/// \brief Base class for server transport implementations. +/// +/// Provides shared helpers for constructing message readers/writers from a +/// transport-level data stream, and for writing a FlightDataStream to a +/// transport-level stream (used by DoGet/DoExchange). +class ARROW_FLIGHT_EXPORT ServerTransportBase { + public: + explicit ServerTransportBase(std::shared_ptr memory_manager) + : memory_manager_(std::move(memory_manager)) {} + virtual ~ServerTransportBase() = default; + + protected: + /// \brief Create a FlightMessageReader that reads from a transport-level stream. + /// + /// \param[in] stream The transport-specific data stream to read from. + arrow::Result> MakeMessageReader( + ServerDataStream* stream) const; + /// \brief Create a FlightMetadataWriter that writes to a transport-level stream. + /// + /// \param[in] stream The transport-specific data stream to write to. + std::unique_ptr MakeMetadataWriter( + ServerDataStream* stream) const; + /// \brief Create a FlightMessageWriter that writes to a transport-level stream. + /// + /// \param[in] stream The transport-specific data stream to write to. + std::unique_ptr MakeMessageWriter(ServerDataStream* stream) const; + /// \brief Write a FlightDataStream to a transport-level stream. + /// + /// Used by DoGet and DoExchange to stream results back to the client. + /// The schema is written first, followed by record batches until the + /// stream is exhausted. + /// + /// \param[in] data_stream The Arrow data stream to read from. + /// \param[in] stream The transport-specific data stream to write to. + Status WriteDataStream(std::unique_ptr data_stream, + ServerDataStream* stream) const; + + std::shared_ptr memory_manager_; +}; + +/// \brief Manages a background thread that waits for OS signals. +/// +/// Creates a signal-safe SelfPipe and spawns a thread that blocks on it. +/// When a signal is received, the signal handler writes to the SelfPipe +/// to wake the thread. The destructor ensures the thread is joined. +class ServerSignalHandler { + public: + ARROW_DISALLOW_COPY_AND_ASSIGN(ServerSignalHandler); + ServerSignalHandler() = default; + + /// \brief Initialize the signal handler and start the background thread. + /// + /// \param[in] handler A callable that will be invoked on the background + /// thread, receiving the SelfPipe as an argument. + /// \return The SelfPipe, which can be used to signal the handler thread. + template + arrow::Result> Init(Fn handler) { + ARROW_ASSIGN_OR_RAISE(self_pipe_, + ::arrow::internal::SelfPipe::Make(/*signal_safe=*/true)); + handle_signals_ = std::thread(handler, self_pipe_); + return self_pipe_; + } + + /// \brief Shut down the signal handler and join the background thread. + Status Shutdown() { + RETURN_NOT_OK(self_pipe_->Shutdown()); + handle_signals_.join(); + return Status::OK(); + } + + ~ServerSignalHandler() { ARROW_CHECK_OK(Shutdown()); } + + private: + std::shared_ptr<::arrow::internal::SelfPipe> self_pipe_; + std::thread handle_signals_; +}; + +/// \brief Manages signal-driven graceful shutdown for server transports. +/// +/// Installs signal handlers for the given set of signals. When a signal is +/// received, the registered shutdown callback is invoked. Uses a SelfPipe +/// to safely bridge the signal handler (which runs in a restricted context) +/// to a background thread that performs the actual shutdown. +class ARROW_FLIGHT_EXPORT ServerSignalState { + public: + ServerSignalState() = default; + + /// \brief Configure which signals should trigger shutdown. + /// + /// \param[in] sigs List of signal numbers (e.g. SIGINT, SIGTERM). + Status SetShutdownOnSignals(const std::vector& sigs); + /// \brief Return the last signal that was received, or 0 if none. + int GotSignal() const; + + /// \brief Run the server with signal-aware lifecycle management. + /// + /// Installs signal handlers, calls \p wait (which should block until the + /// server stops), and restores original signal handlers on return. + /// If a signal is received during \p wait, \p shutdown is called first. + /// + /// \param[in] wait A callable that blocks until the server stops. + /// \param[in] shutdown A callable invoked on signal receipt to stop the server. + /// \param[in] not_started_message Logged if shutdown is called before the server + /// starts. \param[in] shutdown_warning Logged if shutdown itself fails. + template + Status Serve(WaitFn&& wait, ShutdownFn&& shutdown, const char* not_started_message, + const char* shutdown_warning) { + got_signal_ = 0; + old_signal_handlers_.clear(); + running_instance_ = this; + + ServerSignalHandler signal_handler; + ARROW_ASSIGN_OR_RAISE(self_pipe_, + signal_handler.Init(&ServerSignalState::WaitForSignals)); + for (size_t i = 0; i < signals_.size(); ++i) { + int signum = signals_[i]; + ::arrow::internal::SignalHandler new_handler(&ServerSignalState::HandleSignal), + old_handler; + ARROW_ASSIGN_OR_RAISE(old_handler, + ::arrow::internal::SetSignalHandler(signum, new_handler)); + old_signal_handlers_.push_back(std::move(old_handler)); + } + + shutdown_ = [shutdown = std::forward(shutdown)]() mutable { + return shutdown(nullptr); + }; + shutdown_warning_ = shutdown_warning; + auto status = wait(); + if (!status.ok()) { + running_instance_ = nullptr; + shutdown_ = nullptr; + shutdown_warning_ = nullptr; + return status; + } + running_instance_ = nullptr; + shutdown_ = nullptr; + shutdown_warning_ = nullptr; + + for (size_t i = 0; i < signals_.size(); ++i) { + RETURN_NOT_OK( + ::arrow::internal::SetSignalHandler(signals_[i], old_signal_handlers_[i]) + .status()); + } + return Status::OK(); + } + + /// \brief Shut down the server through the given callback. + /// + /// Clears the running instance pointer, then invokes \p shutdown. + /// + /// \param[in] shutdown A callable that performs the transport-specific shutdown. + /// \param[in] deadline Optional deadline for the shutdown. + template + Status Shutdown(ShutdownFn&& shutdown, + const std::chrono::system_clock::time_point* deadline) { + running_instance_ = nullptr; + return shutdown(deadline); + } + + /// \brief Wait for the server to stop without initiating shutdown. + /// + /// \param[in] wait A callable that blocks until the server stops. + template + Status Wait(WaitFn&& wait) { + RETURN_NOT_OK(wait()); + running_instance_ = nullptr; + return Status::OK(); + } + + private: + /// \brief Static entry point for the OS signal handler. + static void HandleSignal(int signum); + + /// \brief Record the received signal and wake the background thread. + void DoHandleSignal(int signum); + + /// \brief Background thread that blocks on the SelfPipe, then invokes shutdown. + static void WaitForSignals(std::shared_ptr<::arrow::internal::SelfPipe> self_pipe); + + // Signal handlers (on Windows) and the shutdown handler (other platforms) + // are executed in a separate thread, so getting the current thread instance + // wouldn't make sense. This means only a single instance can receive signals. + static std::atomic running_instance_; + + // We'll use the self-pipe trick to notify a thread from the signal + // handler. The thread will then shut down the server. + std::shared_ptr<::arrow::internal::SelfPipe> self_pipe_; + + // Signal handling + std::vector signals_; + std::vector<::arrow::internal::SignalHandler> old_signal_handlers_; + std::atomic got_signal_{0}; + std::function shutdown_; + const char* shutdown_warning_ = nullptr; +}; + +/// \brief Parse a Flight Location into a URI. +/// +/// \param[in] location The Flight location to parse. +ARROW_FLIGHT_EXPORT +arrow::Result ParseLocationUri(const Location& location); + +/// \brief Extract the port number from a Flight Location. +/// +/// Returns -1 if the location cannot be parsed or has no port. +/// +/// \param[in] location The Flight location. +ARROW_FLIGHT_EXPORT +int PortFromLocation(const Location& location); + +} // namespace arrow::flight::internal