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_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

@@ -546,6 +546,75 @@ class CountingServerMiddlewareFactory : public ServerMiddlewareFactory {
   std::atomic<int> failed_;
 };
 
+// A server middleware that only sends a header once the RPC handler has
+// flagged it, to check that headers are gathered after the handler runs.
+class HandlerStateServerMiddleware : public ServerMiddleware {
+ public:
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+    if (handler_ran) {
+      outgoing_headers->AddHeader("x-handler-ran", "true");
+    }
+  }
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "HandlerStateServerMiddleware"; }
+
+  bool handler_ran = false;
+};
+
+class HandlerStateServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, const ServerCallContext& context,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    *middleware = std::make_shared<HandlerStateServerMiddleware>();
+    return Status::OK();
+  }
+};
+
+class HandlerStateTestServer : public FlightServerBase {
+  Status GetFlightInfo(const ServerCallContext& context, const FlightDescriptor& request,
+                       std::unique_ptr<FlightInfo>* info) override {
+    auto middleware = static_cast<HandlerStateServerMiddleware*>(
+        context.GetMiddleware("handler_state"));
+    if (middleware == nullptr) {
+      return Status::Invalid("Could not find middleware");
+    }
+    middleware->handler_ran = true;
+    ARROW_ASSIGN_OR_RAISE(auto flight_info,
+                          FlightInfo::Make(*arrow::schema({}), request, {}, -1, -1));
+    *info = std::make_unique<FlightInfo>(std::move(flight_info));
+    return Status::OK();
+  }
+};
+
+// A client middleware that records the values of the x-handler-ran header.
+class HandlerStateClientMiddleware : public ClientMiddleware {
+ public:
+  explicit HandlerStateClientMiddleware(std::vector<std::string>* values)
+      : values_(values) {}
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {}
+  void ReceivedHeaders(const CallHeaders& incoming_headers) override {
+    auto [begin, end] = incoming_headers.equal_range("x-handler-ran");
+    for (auto it = begin; it != end; ++it) {
+      values_->emplace_back(it->second);
+    }
+  }
+  void CallCompleted(const Status& status) override {}
+
+ private:
+  std::vector<std::string>* values_;
+};
+
+class HandlerStateClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  void StartCall(const CallInfo& info,
+                 std::unique_ptr<ClientMiddleware>* middleware) override {
+    *middleware = std::make_unique<HandlerStateClientMiddleware>(&values_);
+  }
+
+  std::vector<std::string> values_;
+};
+
 // The current span ID, used to emulate OpenTracing style distributed
 // tracing. Only used for communication between application code and
 // client middleware.
@@ -831,6 +900,34 @@ class TestCountingServerMiddleware : public ::testing::Test {
   std::unique_ptr<FlightServerBase> server_;
 };
 
+class TestHandlerStateServerMiddleware : public ::testing::Test {
+ public:
+  void SetUp() {
+    client_middleware_ = std::make_shared<HandlerStateClientMiddlewareFactory>();
+    ASSERT_OK(MakeServer<HandlerStateTestServer>(
+        &server_, &client_,
+        [](FlightServerOptions* options) {
+          options->middleware.push_back(
+              {"handler_state", std::make_shared<HandlerStateServerMiddlewareFactory>()});
+          return Status::OK();
+        },
+        [&](FlightClientOptions* options) {
+          options->middleware.push_back(client_middleware_);
+          return Status::OK();
+        }));
+  }
+
+  void TearDown() {
+    ASSERT_OK(client_->Close());
+    ASSERT_OK(server_->Shutdown());
+  }
+
+ protected:
+  std::shared_ptr<HandlerStateClientMiddlewareFactory> client_middleware_;
+  std::unique_ptr<FlightClient> client_;
+  std::unique_ptr<FlightServerBase> server_;
+};
+
 // Setup for this test is 2 servers
 // 1. Client makes request to server A with a request ID set
 // 2. server A extracts the request ID and makes a request to server B
@@ -1428,6 +1525,11 @@ TEST_F(TestCountingServerMiddleware, Count) {
   ASSERT_EQ(1, request_counter_->failed_);
 }
 
+TEST_F(TestHandlerStateServerMiddleware, GetFlightInfo) {
+  ASSERT_OK(client_->GetFlightInfo(FlightDescriptor::Command("")));
+  ASSERT_THAT(client_middleware_->values_, ::testing::ElementsAre("true"));
+}
+
 TEST_F(TestPropagatingMiddleware, Propagate) {
   Action action;
   std::unique_ptr<Result> result;

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