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

ci/docker/ubuntu-22.04-cpp.dockerfile

@@ -223,8 +223,11 @@ 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 \
     PYTHON=python3 \
     simdjson_SOURCE=BUNDLED \
-    xsimd_SOURCE=BUNDLED
+    xsimd_SOURCE=BUNDLED \
+    nlohmann_json_SOURCE=BUNDLED

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_internals_test.cc

@@ -18,6 +18,10 @@
 // ----------------------------------------------------------------------
 // Tests for Flight which don't actually spin up a client/server
 
+#include <atomic>
+#include <chrono>
+#include <csignal>
+#include <thread>
 #include <type_traits>
 
 #include <gmock/gmock.h>
@@ -37,6 +41,7 @@
 #include "arrow/record_batch.h"
 #include "arrow/status.h"
 #include "arrow/testing/gtest_util.h"
+#include "arrow/util/io_util.h"
 #include "arrow/util/string.h"
 
 // Include after Flight headers
@@ -861,5 +866,61 @@ TEST(TransportErrorHandling, ReconstructStatus) {
 
 // TODO: test TransportStatusDetail
 
+#if !defined(_WIN32) && defined(ARROW_ENABLE_THREADING)
+
+// Serve() must not leave the process signal handlers overridden when it fails
+TEST(TestServerSignalState, ServeBeforeInitRestoresSignalHandlers) {
+  FlightServerBase server;
+  ASSERT_OK(server.SetShutdownOnSignals({SIGUSR1}));
+  ASSERT_OK_AND_ASSIGN(auto before, ::arrow::internal::GetSignalHandler(SIGUSR1));
+
+  ASSERT_RAISES(UnknownError, server.Serve());
+
+  ASSERT_OK_AND_ASSIGN(auto after, ::arrow::internal::GetSignalHandler(SIGUSR1));
+  EXPECT_EQ(before.callback(), after.callback());
+}
+
+// Shutdown() on a server that was never initialized must not clear the
+// process-wide signal target of another, serving instance
+TEST(TestServerSignalState, ShutdownUninitializedKeepsOtherServerSignalable) {
+  FlightServerBase serving_server;
+  ASSERT_OK_AND_ASSIGN(auto location, Location::ForGrpcTcp("localhost", 0));
+  FlightServerOptions options(location);
+  ASSERT_OK(serving_server.Init(options));
+  ASSERT_OK(serving_server.SetShutdownOnSignals({SIGUSR1}));
+
+  std::atomic<bool> returned{false};
+  std::thread serving_thread([&] {
+    auto status = serving_server.Serve();
+    ARROW_UNUSED(status);
+    returned = true;
+  });
+
+  // Wait for Serve() to install the process signal handler.
+  ASSERT_OK_AND_ASSIGN(auto original, ::arrow::internal::GetSignalHandler(SIGUSR1));
+  for (int i = 0; i < 1000 && !returned.load(); ++i) {
+    ASSERT_OK_AND_ASSIGN(auto current, ::arrow::internal::GetSignalHandler(SIGUSR1));
+    if (current.callback() != original.callback()) {
+      break;
+    }
+    std::this_thread::sleep_for(std::chrono::milliseconds(1));
+  }
+
+  FlightServerBase uninitialized;
+  ASSERT_RAISES(Invalid, uninitialized.Shutdown());
+
+  ::raise(SIGUSR1);
+  for (int i = 0; i < 500 && !returned.load(); ++i) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(10));
+  }
+  EXPECT_TRUE(returned.load());
+  if (!returned.load()) {
+    ASSERT_OK(serving_server.Shutdown());
+  }
+  serving_thread.join();
+}
+
+#endif
+
 }  // namespace flight
 }  // namespace arrow

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,59 @@ 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, options.location.uri());
+  return impl_->transport_->Init(options, uri);
 }
 
-int FlightServerBase::port() const { return location().uri_->port(); }
+int FlightServerBase::port() const { return location().port(); }
 
 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() {
+  // Check before entering the signal helper: a failing wait() callback would
+  // otherwise leave the process signal handlers overridden, and clearing the
+  // signal target on the error path would disable signals for another server.
   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 { return impl_->transport_->Wait(); },
+      [this](const std::chrono::system_clock::time_point* deadline) -> Status {
+        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) {
+  if (!impl_->transport_) {
     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 (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.h

@@ -195,9 +195,8 @@ class ARROW_FLIGHT_EXPORT FlightServerBase {
   Status Init(const FlightServerOptions& options);
 
   /// \brief Get the port that the Flight server is listening on.
-  /// This method must only be called after Init().  Will return a
-  /// non-positive value if no port exists (e.g. when listening on a
-  /// domain socket).
+  /// This method must only be called after Init().
+  /// \return The port number. -1 if the port is not available.
   int port() const;
 
   /// \brief Get the address that the Flight server is listening on.

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