Skip to content

Commit 05ce0a2

Browse files
pitroupitrou
committed
ARROW-4687: [Python] Stop Flight server on incoming signals
Author: Antoine Pitrou <antoine@python.org> Closes apache#3770 from pitrou/ARROW-4687-python-flight-sigint and squashes the following commits: 5e54883 <Antoine Pitrou> Use std::atomic<> 8a28061 <Antoine Pitrou> ARROW-4687: Stop Flight server on incoming signals
1 parent 07e6f5a commit 05ce0a2

10 files changed

Lines changed: 175 additions & 50 deletions

File tree

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -192,12 +192,11 @@ void Shutdown(int signal) {
192192
int main(int argc, char** argv) {
193193
gflags::ParseCommandLineFlags(&argc, &argv, true);
194194

195-
// SIGTERM shuts down the server
196-
signal(SIGTERM, Shutdown);
197-
198195
g_server.reset(new arrow::flight::FlightPerfServer);
199196

200-
// TODO(wesm): How can we tell if the server failed to start for some reason?
201-
g_server->Run(FLAGS_port);
197+
ARROW_CHECK_OK(g_server->Init(FLAGS_port));
198+
// Exit with a clean error code (0) on SIGTERM
199+
ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
200+
ARROW_CHECK_OK(g_server->Serve());
202201
return 0;
203202
}

‎cpp/src/arrow/flight/server.cc‎

Lines changed: 86 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
#include "arrow/flight/server.h"
1919
#include "arrow/flight/protocol-internal.h"
2020

21+
#include <signal.h>
22+
#include <atomic>
2123
#include <cstdint>
2224
#include <memory>
2325
#include <string>
@@ -294,30 +296,101 @@ class FlightServiceImpl : public FlightService::Service {
294296
FlightServerBase* server_;
295297
};
296298

297-
struct FlightServerBase::FlightServerBaseImpl {
298-
std::unique_ptr<grpc::Server> server;
299+
#if (ATOMIC_INT_LOCK_FREE != 2 || ATOMIC_POINTER_LOCK_FREE != 2)
300+
#error "atomic ints and atomic pointers not always lock-free!"
301+
#endif
302+
303+
struct FlightServerBase::Impl {
304+
std::string address_;
305+
std::unique_ptr<FlightServiceImpl> service_;
306+
std::unique_ptr<grpc::Server> server_;
307+
308+
// Signal handling
309+
std::vector<int> signals_;
310+
std::vector<struct sigaction> old_signal_handlers_;
311+
std::atomic<int> got_signal_;
312+
static thread_local std::atomic<Impl*> running_instance_;
313+
314+
static void HandleSignal(int signum);
315+
316+
void DoHandleSignal(int signum) {
317+
got_signal_ = signum;
318+
server_->Shutdown();
319+
}
299320
};
300321

301-
FlightServerBase::FlightServerBase() { impl_.reset(new FlightServerBaseImpl); }
322+
thread_local std::atomic<FlightServerBase::Impl*>
323+
FlightServerBase::Impl::running_instance_;
324+
325+
void FlightServerBase::Impl::HandleSignal(int signum) {
326+
auto instance = running_instance_.load();
327+
if (instance != nullptr) {
328+
instance->DoHandleSignal(signum);
329+
}
330+
}
331+
332+
FlightServerBase::FlightServerBase() { impl_.reset(new Impl); }
302333

303334
FlightServerBase::~FlightServerBase() {}
304335

305-
void FlightServerBase::Run(int port) {
306-
std::string address = "localhost:" + std::to_string(port);
336+
Status FlightServerBase::Init(int port) {
337+
impl_->address_ = "localhost:" + std::to_string(port);
338+
impl_->service_.reset(new FlightServiceImpl(this));
307339

308-
FlightServiceImpl service(this);
309340
grpc::ServerBuilder builder;
310-
builder.AddListeningPort(address, grpc::InsecureServerCredentials());
311-
builder.RegisterService(&service);
341+
builder.AddListeningPort(impl_->address_, grpc::InsecureServerCredentials());
342+
builder.RegisterService(impl_->service_.get());
343+
344+
impl_->server_ = builder.BuildAndStart();
345+
return Status::OK();
346+
}
312347

313-
impl_->server = builder.BuildAndStart();
314-
std::cout << "Server listening on " << address << std::endl;
315-
impl_->server->Wait();
348+
Status FlightServerBase::SetShutdownOnSignals(const std::vector<int> sigs) {
349+
impl_->signals_ = sigs;
350+
impl_->old_signal_handlers_.clear();
351+
return Status::OK();
316352
}
317353

354+
Status FlightServerBase::Serve() {
355+
impl_->got_signal_ = 0;
356+
impl_->running_instance_ = impl_.get();
357+
358+
// Setup signal handlers
359+
impl_->old_signal_handlers_.clear();
360+
for (size_t i = 0; i < impl_->signals_.size(); ++i) {
361+
int signum = impl_->signals_[i];
362+
// Override with our own handler so as to stop the server.
363+
struct sigaction sa, old_handler;
364+
sa.sa_handler = &Impl::HandleSignal;
365+
sa.sa_flags = 0;
366+
sigemptyset(&sa.sa_mask);
367+
int ret = sigaction(signum, &sa, &old_handler);
368+
if (ret != 0) {
369+
return Status::IOError("sigaction call failed");
370+
}
371+
impl_->old_signal_handlers_.push_back(old_handler);
372+
}
373+
374+
// TODO(wesm): How can we tell if the server failed to start for some reason?
375+
impl_->server_->Wait();
376+
impl_->running_instance_ = nullptr;
377+
378+
// Restore signal handlers
379+
for (size_t i = 0; i < impl_->signals_.size(); ++i) {
380+
int ret = sigaction(impl_->signals_[i], &impl_->old_signal_handlers_[i], nullptr);
381+
if (ret != 0) {
382+
return Status::IOError("sigaction call failed");
383+
}
384+
}
385+
386+
return Status::OK();
387+
}
388+
389+
int FlightServerBase::GotSignal() const { return impl_->got_signal_; }
390+
318391
void FlightServerBase::Shutdown() {
319-
DCHECK(impl_->server);
320-
impl_->server->Shutdown();
392+
DCHECK(impl_->server_);
393+
impl_->server_->Shutdown();
321394
}
322395

323396
Status FlightServerBase::ListFlights(const Criteria* criteria,
@@ -327,7 +400,6 @@ Status FlightServerBase::ListFlights(const Criteria* criteria,
327400

328401
Status FlightServerBase::GetFlightInfo(const FlightDescriptor& request,
329402
std::unique_ptr<FlightInfo>* info) {
330-
std::cout << "GetFlightInfo" << std::endl;
331403
return Status::NotImplemented("NYI");
332404
}
333405

‎cpp/src/arrow/flight/server.h‎

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -79,13 +79,31 @@ class ARROW_EXPORT FlightServerBase {
7979
FlightServerBase();
8080
virtual ~FlightServerBase();
8181

82-
/// \brief Run an insecure server on localhost at the indicated port. Block
83-
/// until server is shut down or otherwise terminates
84-
/// \param[in] port the port to bind to
85-
void Run(int port);
82+
// Lifecycle methods.
83+
84+
/// \brief Initialize an insecure TCP server listening on localhost
85+
/// at the given port.
86+
/// This method must be called before any other method.
87+
Status Init(int port);
88+
89+
/// \brief Set the server to stop when receiving any of the given signal
90+
/// numbers.
91+
/// This method must be called before Serve().
92+
Status SetShutdownOnSignals(const std::vector<int> sigs);
93+
94+
/// \brief Start serving.
95+
/// This method blocks until either Shutdown() is called or one of the signals
96+
/// registered in SetShutdownOnSignals() is received.
97+
Status Serve();
98+
99+
/// \brief Query whether Serve() was interrupted by a signal.
100+
/// This method must be called after Serve() has returned.
101+
///
102+
/// \return int the signal number that interrupted Serve(), if any, otherwise 0
103+
int GotSignal() const;
86104

87105
/// \brief Shut down the server. Can be called from signal handler or another
88-
/// thread while Run blocks
106+
/// thread while Serve() blocks.
89107
///
90108
/// TODO(wesm): Shutdown with deadline
91109
void Shutdown();
@@ -132,8 +150,8 @@ class ARROW_EXPORT FlightServerBase {
132150
virtual Status ListActions(std::vector<ActionType>* actions);
133151

134152
private:
135-
struct FlightServerBaseImpl;
136-
std::unique_ptr<FlightServerBaseImpl> impl_;
153+
struct Impl;
154+
std::unique_ptr<Impl> impl_;
137155
};
138156

139157
} // namespace flight

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

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -117,20 +117,16 @@ class FlightIntegrationTestServer : public FlightServerBase {
117117

118118
std::unique_ptr<arrow::flight::FlightIntegrationTestServer> g_server;
119119

120-
void Shutdown(int signal) {
121-
if (g_server != nullptr) {
122-
g_server->Shutdown();
123-
}
124-
}
125-
126120
int main(int argc, char** argv) {
127121
gflags::SetUsageMessage("Integration testing server for Flight.");
128122
gflags::ParseCommandLineFlags(&argc, &argv, true);
129123

130-
// SIGTERM shuts down the server
131-
signal(SIGTERM, Shutdown);
132-
133124
g_server.reset(new arrow::flight::FlightIntegrationTestServer);
134-
g_server->Run(FLAGS_port);
125+
ARROW_CHECK_OK(g_server->Init(FLAGS_port));
126+
// Exit with a clean error code (0) on SIGTERM
127+
ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
128+
129+
std::cout << "Server listening on localhost:" << FLAGS_port << std::endl;
130+
ARROW_CHECK_OK(g_server->Serve());
135131
return 0;
136132
}

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

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727

2828
#include "arrow/io/test-common.h"
2929
#include "arrow/record_batch.h"
30+
#include "arrow/util/logging.h"
3031

3132
#include "arrow/flight/server.h"
3233
#include "arrow/flight/test-util.h"
@@ -121,21 +122,15 @@ class FlightTestServer : public FlightServerBase {
121122

122123
std::unique_ptr<arrow::flight::FlightTestServer> g_server;
123124

124-
void Shutdown(int signal) {
125-
if (g_server != nullptr) {
126-
g_server->Shutdown();
127-
}
128-
}
129-
130125
int main(int argc, char** argv) {
131126
gflags::ParseCommandLineFlags(&argc, &argv, true);
132127

133-
// SIGTERM shuts down the server
134-
signal(SIGTERM, Shutdown);
135-
136128
g_server.reset(new arrow::flight::FlightTestServer);
129+
ARROW_CHECK_OK(g_server->Init(FLAGS_port));
130+
// Exit with a clean error code (0) on SIGTERM
131+
ARROW_CHECK_OK(g_server->SetShutdownOnSignals({SIGTERM}));
137132

138-
// TODO(wesm): How can we tell if the server failed to start for some reason?
139-
g_server->Run(FLAGS_port);
133+
std::cout << "Server listening on localhost:" << FLAGS_port << std::endl;
134+
ARROW_CHECK_OK(g_server->Serve());
140135
return 0;
141136
}

‎cpp/src/arrow/python/flight.cc‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,12 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
#include <signal.h>
1819
#include <utility>
1920

2021
#include "arrow/flight/internal.h"
2122
#include "arrow/python/flight.h"
23+
#include "arrow/util/logging.h"
2224

2325
namespace arrow {
2426
namespace py {
@@ -71,6 +73,37 @@ Status PyFlightServer::ListActions(std::vector<arrow::flight::ActionType>* actio
7173
return CheckPyError();
7274
}
7375

76+
Status PyFlightServer::ServeWithSignals() {
77+
// Respect the current Python settings, i.e. only interrupt the server if there is
78+
// an active signal handler for SIGINT and SIGTERM.
79+
std::vector<int> signals;
80+
for (const int signum : {SIGINT, SIGTERM}) {
81+
struct sigaction handler;
82+
int ret = sigaction(signum, nullptr, &handler);
83+
if (ret != 0) {
84+
return Status::IOError("sigaction call failed");
85+
}
86+
if (handler.sa_handler != SIG_DFL && handler.sa_handler != SIG_IGN) {
87+
signals.push_back(signum);
88+
}
89+
}
90+
RETURN_NOT_OK(SetShutdownOnSignals(signals));
91+
92+
// Serve until we got told to shutdown or a signal interrupted us
93+
RETURN_NOT_OK(Serve());
94+
int signum = GotSignal();
95+
if (signum != 0) {
96+
// Issue the signal again with Python's signal handlers restored
97+
PyAcquireGIL lock;
98+
raise(signum);
99+
// XXX Ideally we would loop and serve again if no exception was raised.
100+
// Unfortunately, gRPC will return immediately if Serve() is called again.
101+
ARROW_UNUSED(PyErr_CheckSignals());
102+
}
103+
104+
return Status::OK();
105+
}
106+
74107
PyFlightResultStream::PyFlightResultStream(PyObject* generator,
75108
PyFlightResultStreamCallback callback)
76109
: callback_(callback) {

‎cpp/src/arrow/python/flight.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,10 @@ class ARROW_PYTHON_EXPORT PyFlightServer : public arrow::flight::FlightServerBas
5656
public:
5757
explicit PyFlightServer(PyObject* server, PyFlightServerVtable vtable);
5858

59+
// Like Serve(), but set up signals and invoke Python signal handlers
60+
// if necessary. This function may return with a Python exception set.
61+
Status ServeWithSignals();
62+
5963
Status ListFlights(const arrow::flight::Criteria* criteria,
6064
std::unique_ptr<arrow::flight::FlightListing>* listings) override;
6165
Status GetFlightInfo(const arrow::flight::FlightDescriptor& request,

‎python/examples/flight/server.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ def do_action(self, action):
9494

9595
def _shutdown(self):
9696
"""Shut down after a delay."""
97+
print("Server is shutting down...")
9798
time.sleep(2)
9899
self.shutdown()
99100

‎python/pyarrow/_flight.pyx‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -545,15 +545,20 @@ cdef class FlightServerBase:
545545
cdef:
546546
PyFlightServerVtable vtable = PyFlightServerVtable()
547547
int c_port = port
548+
PyFlightServer* c_server
549+
548550
vtable.list_flights = &_list_flights
549551
vtable.get_flight_info = &_get_flight_info
550552
vtable.do_put = &_do_put
551553
vtable.do_get = &_do_get
552554
vtable.list_actions = &_list_actions
553555
vtable.do_action = &_do_action
554-
self.server.reset(new PyFlightServer(self, vtable))
556+
557+
c_server = new PyFlightServer(self, vtable)
558+
self.server.reset(c_server)
555559
with nogil:
556-
self.server.get().Run(c_port)
560+
check_status(c_server.Init(c_port))
561+
check_status(c_server.ServeWithSignals())
557562

558563
def list_flights(self, criteria):
559564
raise NotImplementedError

‎python/pyarrow/includes/libarrow_flight.pxd‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,9 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
146146

147147
cdef cppclass PyFlightServer:
148148
PyFlightServer(object server, PyFlightServerVtable vtable)
149-
void Run(int port)
149+
150+
CStatus Init(int port)
151+
CStatus ServeWithSignals() except *
150152
void Shutdown()
151153

152154
cdef cppclass CPyFlightResultStream\

0 commit comments

Comments
 (0)