From 58583eed3ab9c103cb4c0b3a2f035a1a9c75b8f6 Mon Sep 17 00:00:00 2001
From: David Li <li.davidm96@gmail.com>
Date: Thu, 23 Dec 2021 10:11:11 -0500
Subject: [PATCH] [C++][FlightRPC] Fix race between signal handler and shutdown

---
 cpp/src/arrow/flight/flight_test.cc | 18 ++++++++++++++++++
 cpp/src/arrow/flight/server.cc      |  3 ++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/cpp/src/arrow/flight/flight_test.cc b/cpp/src/arrow/flight/flight_test.cc
index 56ca468a043..99579b954a4 100644
--- a/cpp/src/arrow/flight/flight_test.cc
+++ b/cpp/src/arrow/flight/flight_test.cc
@@ -356,6 +356,24 @@ TEST(TestFlight, BuilderHook) {
   ASSERT_OK(server->Shutdown());
 }
 
+TEST(TestFlight, ServeShutdown) {
+  // Regression test for ARROW-15181
+  constexpr int kIterations = 10;
+  for (int i = 0; i < kIterations; i++) {
+    Location location;
+    std::unique_ptr<FlightServerBase> server = ExampleTestServer();
+
+    ASSERT_OK(Location::ForGrpcTcp("localhost", 0, &location));
+    FlightServerOptions options(location);
+    ASSERT_OK(server->Init(options));
+    ASSERT_GT(server->port(), 0);
+    std::thread t([&]() { ASSERT_OK(server->Serve()); });
+    ASSERT_OK(server->Shutdown());
+    ASSERT_OK(server->Wait());
+    t.join();
+  }
+}
+
 // ----------------------------------------------------------------------
 // Client tests
 
diff --git a/cpp/src/arrow/flight/server.cc b/cpp/src/arrow/flight/server.cc
index b52c1624657..988bd690d51 100644
--- a/cpp/src/arrow/flight/server.cc
+++ b/cpp/src/arrow/flight/server.cc
@@ -826,9 +826,9 @@ class ServerSignalHandler {
         errno != EWOULDBLOCK && errno != EINTR) {
       return arrow::internal::IOErrorFromErrno(errno, "Could not unblock signal thread");
     }
+    handle_signals_.join();
     RETURN_NOT_OK(arrow::internal::FileClose(self_pipe_.rfd));
     RETURN_NOT_OK(arrow::internal::FileClose(self_pipe_.wfd));
-    handle_signals_.join();
     self_pipe_.rfd = 0;
     self_pipe_.wfd = 0;
     return Status::OK();
@@ -1012,6 +1012,7 @@ Status FlightServerBase::Shutdown() {
   if (!server) {
     return Status::Invalid("Shutdown() on uninitialized FlightServerBase");
   }
+  impl_->running_instance_ = nullptr;
   impl_->server_->Shutdown();
   return Status::OK();
 }