ARROW-4687: [Python] Stop Flight server on incoming signals

cpp/src/arrow/flight/perf-server.cc

@@ -192,12 +192,11 @@ void Shutdown(int signal) {
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
 
-  // SIGTERM shuts down the server
-  signal(SIGTERM, Shutdown);
-
   g_server.reset(new arrow::flight::FlightPerfServer);
 
-  // TODO(wesm): How can we tell if the server failed to start for some reason?
-  g_server->Run(FLAGS_port);
+  ARROW_CHECK_OK(g_server->Init(FLAGS_port));
+  // Exit with a clean error code (0) on SIGTERM
+  ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
+  ARROW_CHECK_OK(g_server->Serve());
   return 0;
 }

cpp/src/arrow/flight/server.cc

@@ -18,6 +18,7 @@
 #include "arrow/flight/server.h"
 #include "arrow/flight/protocol-internal.h"
 
+#include <signal.h>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -294,30 +295,97 @@ class FlightServiceImpl : public FlightService::Service {
   FlightServerBase* server_;
 };
 
-struct FlightServerBase::FlightServerBaseImpl {
-  std::unique_ptr<grpc::Server> server;
+struct FlightServerBase::Impl {
+  std::string address_;
+  std::unique_ptr<FlightServiceImpl> service_;
+  std::unique_ptr<grpc::Server> server_;
+
+  // Signal handling
+  std::vector<int> signals_;
+  std::vector<struct sigaction> old_signal_handlers_;
+  volatile int got_signal_;
+  static volatile thread_local Impl* running_instance_;
+
+  static void HandleSignal(int signum);
+
+  void DoHandleSignal(int signum) {
+    got_signal_ = signum;
+    server_->Shutdown();
+  }
 };
 
-FlightServerBase::FlightServerBase() { impl_.reset(new FlightServerBaseImpl); }
+volatile thread_local FlightServerBase::Impl* FlightServerBase::Impl::running_instance_ =
+    nullptr;
+
+void FlightServerBase::Impl::HandleSignal(int signum) {
+  auto instance = const_cast<Impl*>(running_instance_);  // remove volatile
+  if (instance != nullptr) {
+    instance->DoHandleSignal(signum);
+  }
+}
+
+FlightServerBase::FlightServerBase() { impl_.reset(new Impl); }
 
 FlightServerBase::~FlightServerBase() {}
 
-void FlightServerBase::Run(int port) {
-  std::string address = "localhost:" + std::to_string(port);
+Status FlightServerBase::Init(int port) {
+  impl_->address_ = "localhost:" + std::to_string(port);
+  impl_->service_.reset(new FlightServiceImpl(this));
 
-  FlightServiceImpl service(this);
   grpc::ServerBuilder builder;
-  builder.AddListeningPort(address, grpc::InsecureServerCredentials());
-  builder.RegisterService(&service);
+  builder.AddListeningPort(impl_->address_, grpc::InsecureServerCredentials());
+  builder.RegisterService(impl_->service_.get());
 
-  impl_->server = builder.BuildAndStart();
-  std::cout << "Server listening on " << address << std::endl;
-  impl_->server->Wait();
+  impl_->server_ = builder.BuildAndStart();
+  return Status::OK();
 }
 
+Status FlightServerBase::SetShutdownOnSignals(const std::vector<int> sigs) {
+  impl_->signals_ = sigs;
+  impl_->old_signal_handlers_.clear();
+  return Status::OK();
+}
+
+Status FlightServerBase::Serve() {
+  impl_->got_signal_ = 0;
+  impl_->running_instance_ = impl_.get();
+
+  // Setup signal handlers
+  impl_->old_signal_handlers_.clear();
+  for (size_t i = 0; i < impl_->signals_.size(); ++i) {
+    int signum = impl_->signals_[i];
+    // Override with our own handler so as to stop the server.
+    struct sigaction sa, old_handler;
+    sa.sa_handler = &Impl::HandleSignal;
+    sa.sa_flags = 0;
+    sigemptyset(&sa.sa_mask);
+    int ret = sigaction(signum, &sa, &old_handler);
+    if (ret != 0) {
+      return Status::IOError("sigaction call failed");
+    }
+    impl_->old_signal_handlers_.push_back(old_handler);
+  }
+
+  // TODO(wesm): How can we tell if the server failed to start for some reason?
+  impl_->server_->Wait();
+  impl_->running_instance_ = nullptr;
+
+  // Restore signal handlers
+  for (size_t i = 0; i < impl_->signals_.size(); ++i) {
+    int ret = sigaction(impl_->signals_[i], &impl_->old_signal_handlers_[i], nullptr);
+    if (ret != 0) {
+      return Status::IOError("sigaction call failed");
+    }
+  }
+
+  return Status::OK();
+}
+
+int FlightServerBase::GotSignal() const { return impl_->got_signal_; }
+
 void FlightServerBase::Shutdown() {
-  DCHECK(impl_->server);
-  impl_->server->Shutdown();
+  DCHECK(impl_->server_);
+  impl_->server_->Shutdown();
 }
 
 Status FlightServerBase::ListFlights(const Criteria* criteria,
@@ -327,7 +395,6 @@ Status FlightServerBase::ListFlights(const Criteria* criteria,
 
 Status FlightServerBase::GetFlightInfo(const FlightDescriptor& request,
                                        std::unique_ptr<FlightInfo>* info) {
-  std::cout << "GetFlightInfo" << std::endl;
   return Status::NotImplemented("NYI");
 }
 

cpp/src/arrow/flight/server.h

@@ -79,13 +79,31 @@ class ARROW_EXPORT FlightServerBase {
   FlightServerBase();
   virtual ~FlightServerBase();
 
-  /// \brief Run an insecure server on localhost at the indicated port. Block
-  /// until server is shut down or otherwise terminates
-  /// \param[in] port the port to bind to
-  void Run(int port);
+  // Lifecycle methods.
+
+  /// \brief Initialize an insecure TCP server listening on localhost
+  /// at the given port.
+  /// This method must be called before any other method.
+  Status Init(int port);
+
+  /// \brief Set the server to stop when receiving any of the given signal
+  /// numbers.
+  /// This method must be called before Serve().
+  Status SetShutdownOnSignals(const std::vector<int> sigs);
+
+  /// \brief Start serving.
+  /// This method blocks until either Shutdown() is called or one of the signals
+  /// registered in SetShutdownOnSignals() is received.
+  Status Serve();
+
+  /// \brief Query whether Serve() was interrupted by a signal.
+  /// This method must be called after Serve() has returned.
+  ///
+  /// \return int the signal number that interrupted Serve(), if any, otherwise 0
+  int GotSignal() const;
 
   /// \brief Shut down the server. Can be called from signal handler or another
-  /// thread while Run blocks
+  /// thread while Serve() blocks.
   ///
   /// TODO(wesm): Shutdown with deadline
   void Shutdown();
@@ -132,8 +150,8 @@ class ARROW_EXPORT FlightServerBase {
   virtual Status ListActions(std::vector<ActionType>* actions);
 
  private:
-  struct FlightServerBaseImpl;
-  std::unique_ptr<FlightServerBaseImpl> impl_;
+  struct Impl;
+  std::unique_ptr<Impl> impl_;
 };
 
 }  // namespace flight

cpp/src/arrow/flight/test-integration-server.cc

@@ -117,20 +117,16 @@ class FlightIntegrationTestServer : public FlightServerBase {
 
 std::unique_ptr<arrow::flight::FlightIntegrationTestServer> g_server;
 
-void Shutdown(int signal) {
-  if (g_server != nullptr) {
-    g_server->Shutdown();
-  }
-}
-
 int main(int argc, char** argv) {
   gflags::SetUsageMessage("Integration testing server for Flight.");
   gflags::ParseCommandLineFlags(&argc, &argv, true);
 
-  // SIGTERM shuts down the server
-  signal(SIGTERM, Shutdown);
-
   g_server.reset(new arrow::flight::FlightIntegrationTestServer);
-  g_server->Run(FLAGS_port);
+  ARROW_CHECK_OK(g_server->Init(FLAGS_port));
+  // Exit with a clean error code (0) on SIGTERM
+  ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
+
+  std::cout << "Server listening on localhost:" << FLAGS_port << std::endl;
+  ARROW_CHECK_OK(g_server->Serve());
   return 0;
 }

cpp/src/arrow/flight/test-server.cc

@@ -27,6 +27,7 @@
 
 #include "arrow/io/test-common.h"
 #include "arrow/record_batch.h"
+#include "arrow/util/logging.h"
 
 #include "arrow/flight/server.h"
 #include "arrow/flight/test-util.h"
@@ -121,21 +122,15 @@ class FlightTestServer : public FlightServerBase {
 
 std::unique_ptr<arrow::flight::FlightTestServer> g_server;
 
-void Shutdown(int signal) {
-  if (g_server != nullptr) {
-    g_server->Shutdown();
-  }
-}
-
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
 
-  // SIGTERM shuts down the server
-  signal(SIGTERM, Shutdown);
-
   g_server.reset(new arrow::flight::FlightTestServer);
+  ARROW_CHECK_OK(g_server->Init(FLAGS_port));
+  // Exit with a clean error code (0) on SIGTERM
+  ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
 
-  // TODO(wesm): How can we tell if the server failed to start for some reason?
-  g_server->Run(FLAGS_port);
+  std::cout << "Server listening on localhost:" << FLAGS_port << std::endl;
+  ARROW_CHECK_OK(g_server->Serve());
   return 0;
 }

cpp/src/arrow/python/flight.cc

@@ -15,10 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#include <signal.h>
 #include <utility>
 
 #include "arrow/flight/internal.h"
 #include "arrow/python/flight.h"
+#include "arrow/util/logging.h"
 
 namespace arrow {
 namespace py {
@@ -71,6 +73,37 @@ Status PyFlightServer::ListActions(std::vector<arrow::flight::ActionType>* actio
   return CheckPyError();
 }
 
+Status PyFlightServer::ServeWithSignals() {
+  // Respect the current Python settings, i.e. only interrupt the server if there is
+  // an active signal handler for SIGINT and SIGTERM.
+  std::vector<int> signals;
+  for (const int signum : {SIGINT, SIGTERM}) {
+    struct sigaction handler;
+    int ret = sigaction(signum, nullptr, &handler);
+    if (ret != 0) {
+      return Status::IOError("sigaction call failed");
+    }
+    if (handler.sa_handler != SIG_DFL && handler.sa_handler != SIG_IGN) {
+      signals.push_back(signum);
+    }
+  }
+  RETURN_NOT_OK(SetShutdownOnSignals(signals));
+
+  // Serve until we got told to shutdown or a signal interrupted us
+  RETURN_NOT_OK(Serve());
+  int signum = GotSignal();
+  if (signum != 0) {
+    // Issue the signal again with Python's signal handlers restored
+    PyAcquireGIL lock;
+    raise(signum);
+    // XXX Ideally we would loop and serve again if no exception was raised.
+    // Unfortunately, gRPC will return immediately if Serve() is called again.
+    ARROW_UNUSED(PyErr_CheckSignals());
+  }
+
+  return Status::OK();
+}
+
 PyFlightResultStream::PyFlightResultStream(PyObject* generator,
                                            PyFlightResultStreamCallback callback)
     : callback_(callback) {

cpp/src/arrow/python/flight.h

@@ -56,6 +56,10 @@ class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBas
  public:
   explicit PyFlightServer(PyObject* server, PyFlightServerVtable vtable);
 
+  // Like Serve(), but set up signals and invoke Python signal handlers
+  // if necessary.  This function may return with a Python exception set.
+  Status ServeWithSignals();
+
   Status ListFlights(const arrow::flight::Criteria* criteria,
                      std::unique_ptr<arrow::flight::FlightListing>* listings) override;
   Status GetFlightInfo(const arrow::flight::FlightDescriptor& request,

python/examples/flight/server.py

@@ -94,6 +94,7 @@ class FlightServer(pyarrow.flight.FlightServerBase):
 
     def _shutdown(self):
         """Shut down after a delay."""
+        print("Server is shutting down...")
         time.sleep(2)
         self.shutdown()
 

python/pyarrow/_flight.pyx

@@ -545,15 +545,20 @@ cdef class FlightServerBase:
         cdef:
             PyFlightServerVtable vtable = PyFlightServerVtable()
             int c_port = port
+            PyFlightServer* c_server
+
         vtable.list_flights = &_list_flights
         vtable.get_flight_info = &_get_flight_info
         vtable.do_put = &_do_put
         vtable.do_get = &_do_get
         vtable.list_actions = &_list_actions
         vtable.do_action = &_do_action
-        self.server.reset(new PyFlightServer(self, vtable))
+
+        c_server = new PyFlightServer(self, vtable)
+        self.server.reset(c_server)
         with nogil:
-            self.server.get().Run(c_port)
+            check_status(c_server.Init(c_port))
+            check_status(c_server.ServeWithSignals())
 
     def list_flights(self, criteria):
         raise NotImplementedError

python/pyarrow/includes/libarrow_flight.pxd

@@ -146,7 +146,9 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
 
     cdef cppclass PyFlightServer:
         PyFlightServer(object server, PyFlightServerVtable vtable)
-        void Run(int port)
+
+        CStatus Init(int port)
+        CStatus ServeWithSignals() except *
         void Shutdown()
 
     cdef cppclass CPyFlightResultStream\