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

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

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

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<arrow::flight::protocol::FlightData> {
   // 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<const arrow::flight::protocol::FlightData*>(&msg);
+    if (arrow::flight::transport::grpc::IsRegisteredGrpcFlightDataMessage(flight_data)) {
+      return GenericSerialize<ProtoBufferWriter, arrow::flight::protocol::FlightData>(
+          msg, bb, own_buffer);
+    }
     return arrow::flight::transport::grpc::FlightDataSerialize(
         *reinterpret_cast<const arrow::flight::FlightPayload*>(&msg), bb, own_buffer);
   }
 
   static Status Deserialize(ByteBuffer* buffer, MessageType* msg) {
+    auto* flight_data = reinterpret_cast<arrow::flight::protocol::FlightData*>(msg);
+    if (arrow::flight::transport::grpc::IsRegisteredGrpcFlightDataMessage(flight_data)) {
+      return GenericDeserialize<ProtoBufferReader, arrow::flight::protocol::FlightData>(
+          buffer, msg);
+    }
     return arrow::flight::transport::grpc::FlightDataDeserialize(
         buffer, reinterpret_cast<arrow::flight::internal::FlightData*>(msg));
   }