GH-50282: [C++][FlightRPC] Refactor GRPC server and transport classes

ci/docker/ubuntu-22.04-cpp.dockerfile

@@ -223,6 +223,8 @@ ENV absl_SOURCE=BUNDLED \
  google_cloud_cpp_storage_SOURCE=BUNDLED \
  opentelemetry_cpp_SOURCE=BUNDLED \
  ORC_SOURCE=BUNDLED \
+ Protobuf_SOURCE=BUNDLED \
+ gRPC_SOURCE=BUNDLED \
  PARQUET_BUILD_EXAMPLES=ON \
  PARQUET_BUILD_EXECUTABLES=ON \
  PATH=/usr/lib/ccache/:$PATH \

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

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());

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',

cpp/src/arrow/flight/server.cc

@@ -23,119 +23,30 @@
 
 #include "arrow/flight/server.h"
 
-#include <atomic>
-#include <cerrno>
 #include <chrono>
-#include <cstdint>
 #include <deque>
 #include <memory>
-#include <string_view>
-#include <thread>
 #include <utility>
 
 #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 <typename Fn>
- arrow::Result<std::shared_ptr<SelfPipe>> 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<SelfPipe> 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<internal::ServerTransport> 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<Impl*> 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<SelfPipe> self_pipe_;
-
- // Signal handling
- std::vector<int> signals_;
- std::vector<SignalHandler> old_signal_handlers_;
- std::atomic<int> 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<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) {
- ARROW_WARN_NOT_OK(instance->transport_->Shutdown(), "Error shutting down server");
- }
- }
+ internal::ServerSignalState signal_state_;
 };
 
-std::atomic<FlightServerBase::Impl*> 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<int> 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,

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<int32_t>(grpc_status.error_code()));