Rework C++ middleware to require infalliability

cpp/src/arrow/flight/flight_test.cc

@@ -594,14 +594,13 @@ class CountingServerMiddleware : public ServerMiddleware {
  public:
   CountingServerMiddleware(std::atomic<int>* successful, std::atomic<int>* failed)
       : successful_(successful), failed_(failed) {}
-  Status SendingHeaders(AddCallHeaders& outgoing_headers) { return Status::OK(); }
-  Status CallCompleted(const Status& status) {
+  void SendingHeaders(AddCallHeaders& outgoing_headers) {}
+  void CallCompleted(const Status& status) {
     if (status.ok()) {
       ARROW_IGNORE_EXPR((*successful_)++);
     } else {
       ARROW_IGNORE_EXPR((*failed_)++);
     }
-    return Status::OK();
   }
 
  private:
@@ -623,15 +622,6 @@ class CountingServerMiddlewareFactory : public ServerMiddlewareFactory {
   std::atomic<int> failed_;
 };
 
-class TracingServerMiddleware : public ServerMiddleware {
- public:
-  TracingServerMiddleware(const std::string& span_id) : span_id_(span_id) {}
-  Status SendingHeaders(AddCallHeaders& outgoing_headers) { return Status::OK(); }
-  Status CallCompleted(const Status& status) { return Status::OK(); }
-
-  std::string span_id_;
-};
-
 static thread_local std::string current_span_id = "";
 
 class TracingServerMiddlewareFactory : public ServerMiddlewareFactory {
@@ -645,7 +635,6 @@ class TracingServerMiddlewareFactory : public ServerMiddlewareFactory {
       const auto& value = (*iter_pair.first).second;
       const std::string copy(value);
       current_span_id = copy;
-      *middleware = std::make_shared<TracingServerMiddleware>(copy);
     }
     return Status::OK();
   }
@@ -700,15 +689,8 @@ class PropagatingClientMiddlewareFactory : public ClientMiddlewareFactory {
 class ReportContextTestServer : public FlightServerBase {
   Status DoAction(const ServerCallContext& context, const Action& action,
                   std::unique_ptr<ResultStream>* result) override {
-    std::shared_ptr<TracingServerMiddleware> middleware =
-        std::static_pointer_cast<TracingServerMiddleware>(
-            context.GetMiddleware("tracing"));
     std::shared_ptr<Buffer> buf;
-    if (middleware) {
-      RETURN_NOT_OK(Buffer::FromString(middleware->span_id_, &buf));
-    } else {
-      RETURN_NOT_OK(Buffer::FromString("", &buf));
-    }
+    RETURN_NOT_OK(Buffer::FromString(current_span_id, &buf));
     *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
     return Status::OK();
   }
@@ -736,9 +718,7 @@ class TestServerMiddleware : public ::testing::Test {
 
     ASSERT_OK(Location::ForGrpcTcp("localhost", 0, &location));
     FlightServerOptions options(location);
-    options.middleware.push_back(
-        std::make_pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>(
-            "test", std::make_shared<RejectServerMiddlewareFactory>()));
+    options.middleware.push_back(std::make_shared<RejectServerMiddlewareFactory>());
     ASSERT_OK(server->Init(options));
 
     Location real_location;
@@ -767,9 +747,7 @@ class TestCountingServerMiddleware : public ::testing::Test {
 
     ASSERT_OK(Location::ForGrpcTcp("localhost", 0, &location));
     FlightServerOptions options(location);
-    options.middleware.push_back(
-        std::make_pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>(
-            "counting", request_counter_));
+    options.middleware.push_back(request_counter_);
     ASSERT_OK(server->Init(options));
 
     Location real_location;
@@ -808,9 +786,7 @@ class TestPropagatingMiddleware : public ::testing::Test {
 
     second_server_ = arrow::internal::make_unique<ReportContextTestServer>();
     FlightServerOptions second_server_options(location);
-    second_server_options.middleware.push_back(
-        std::make_pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>(
-            "tracing", server_middleware_));
+    second_server_options.middleware.push_back(server_middleware_);
     ASSERT_OK(second_server_->Init(second_server_options));
 
     Location real_location;
@@ -825,9 +801,7 @@ class TestPropagatingMiddleware : public ::testing::Test {
         arrow::internal::make_unique<PropagatingTestServer>(std::move(server_client));
 
     FlightServerOptions first_server_options(location);
-    first_server_options.middleware.push_back(
-        std::make_pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>(
-            "tracing", server_middleware_));
+    first_server_options.middleware.push_back(server_middleware_);
     ASSERT_OK(first_server_->Init(first_server_options));
 
     ASSERT_OK(Location::ForGrpcTcp("localhost", first_server_->port(), &real_location));

cpp/src/arrow/flight/server.cc

@@ -237,14 +237,6 @@ class FlightServiceImpl;
 class GrpcServerCallContext : public ServerCallContext {
   const std::string& peer_identity() const override { return peer_identity_; }
 
-  std::shared_ptr<ServerMiddleware> GetMiddleware(const std::string& key) const override {
-    auto result = middleware_map_.find(key);
-    if (result == middleware_map_.end()) {
-      return nullptr;
-    }
-    return result->second;
-  }
-
   // Helper method that runs interceptors given the result of an RPC,
   // then returns the final gRPC status to send to the client
   grpc::Status finish_request(const grpc::Status& status) {
@@ -255,8 +247,7 @@ class GrpcServerCallContext : public ServerCallContext {
 
   grpc::Status finish_request(const arrow::Status& status) {
     for (const auto& instance : middleware_) {
-      // If an error occurs here, remainder of middleware are aborted
-      GRPC_RETURN_NOT_OK(instance->CallCompleted(status));
+      instance->CallCompleted(status);
     }
     return internal::ToGrpcStatus(status);
   }
@@ -266,8 +257,6 @@ class GrpcServerCallContext : public ServerCallContext {
   ServerContext* context_;
   std::string peer_identity_;
   std::vector<std::shared_ptr<ServerMiddleware>> middleware_;
-  mutable std::unordered_map<std::string, std::shared_ptr<ServerMiddleware>>
-      middleware_map_;
 };
 
 class GrpcAddCallHeaders : public AddCallHeaders {
@@ -289,8 +278,7 @@ class FlightServiceImpl : public FlightService::Service {
  public:
   explicit FlightServiceImpl(
       std::shared_ptr<ServerAuthHandler> auth_handler,
-      std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
-          middleware,
+      std::vector<std::shared_ptr<ServerMiddlewareFactory>> middleware,
       FlightServerBase* server)
       : auth_handler_(auth_handler),
         middleware_(std::move(middleware)),
@@ -366,7 +354,7 @@ class FlightServiceImpl : public FlightService::Service {
     GrpcAddCallHeaders outgoing_headers(context);
     for (const auto& factory : middleware_) {
       std::shared_ptr<ServerMiddleware> instance;
-      Status result = factory.second->StartCall(info, incoming_headers, &instance);
+      Status result = factory->StartCall(info, incoming_headers, &instance);
       if (!result.ok()) {
         // Interceptor rejected call, end the request on all existing
         // interceptors
@@ -374,13 +362,7 @@ class FlightServiceImpl : public FlightService::Service {
       }
       if (instance) {
         flight_context.middleware_.push_back(instance);
-        flight_context.middleware_map_[factory.first] = instance;
-        result = instance->SendingHeaders(outgoing_headers);
-        if (!result.ok()) {
-          // Interceptor rejected call, end the request on all existing
-          // interceptors
-          return flight_context.finish_request(result);
-        }
+        instance->SendingHeaders(outgoing_headers);
       }
     }
 
@@ -580,8 +562,7 @@ class FlightServiceImpl : public FlightService::Service {
 
  private:
   std::shared_ptr<ServerAuthHandler> auth_handler_;
-  std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
-      middleware_;
+  std::vector<std::shared_ptr<ServerMiddlewareFactory>> middleware_;
   FlightServerBase* server_;
 };
 

cpp/src/arrow/flight/server.h

@@ -102,10 +102,6 @@ class ARROW_FLIGHT_EXPORT ServerCallContext {
   virtual ~ServerCallContext() = default;
   /// \brief The name of the authenticated peer (may be the empty string)
   virtual const std::string& peer_identity() const = 0;
-  /// \brief Get a middleware instance for this call.
-  /// \return nullptr if not found
-  virtual std::shared_ptr<ServerMiddleware> GetMiddleware(
-      const std::string& key) const = 0;
 };
 
 class ARROW_FLIGHT_EXPORT FlightServerOptions {
@@ -124,9 +120,8 @@ class ARROW_FLIGHT_EXPORT FlightServerOptions {
 
   /// \brief A list of server middleware to apply.
   ///
-  /// It is a vector rather than a map since the order of middleware matters.
-  std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
-      middleware;
+  /// Middleware are always applied in the order provided.
+  std::vector<std::shared_ptr<ServerMiddlewareFactory>> middleware;
 
   /// \brief A Flight implementation-specific callback to customize
   /// transport-specific options.

cpp/src/arrow/flight/server_middleware.h

@@ -28,14 +28,29 @@ namespace arrow {
 
 namespace flight {
 
+/// \brief Server-side middleware for a call, instantiated per RPC.
+///
+/// Middleware should be fast and must be infalliable: there is no way
+/// to reject the call or report errors from the middleware instance.
 class ARROW_FLIGHT_EXPORT ServerMiddleware {
  public:
   virtual ~ServerMiddleware() = default;
 
-  virtual Status SendingHeaders(AddCallHeaders& outgoing_headers) = 0;
-  virtual Status CallCompleted(const Status& status) = 0;
+  /// \brief A callback before headers are sent. Extra headers can be
+  /// added, but existing ones cannot be read.
+  virtual void SendingHeaders(AddCallHeaders& outgoing_headers) = 0;
+
+  /// \brief A callback after the call has completed.
+  virtual void CallCompleted(const Status& status) = 0;
 };
 
+/// \brief A factory for new middleware instances.
+///
+/// If added to a server, this will be called for each RPC (including
+/// Handshake) to give the opportunity to intercept the call.
+///
+/// It is guaranteed that all server middleware methods are called
+/// from the same thread that calls the RPC method implementation.
 class ARROW_FLIGHT_EXPORT ServerMiddlewareFactory {
  public:
   virtual ~ServerMiddlewareFactory() = default;
@@ -50,6 +65,10 @@ class ARROW_FLIGHT_EXPORT ServerMiddlewareFactory {
   /// \param[out] middleware The middleware instance for this call. If
   ///     unset, will not add middleware to this call instance from
   ///     this factory.
+  /// \return Status A non-OK status will reject the call with the
+  ///     given status. Middleware previously in the chain will have
+  ///     their CallCompleted callback called. Other middleware
+  ///     factories will not be called.
   virtual Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
                            std::shared_ptr<ServerMiddleware>* middleware) = 0;
 };

cpp/src/arrow/python/flight.cc

@@ -265,21 +265,27 @@ PyServerMiddleware::PyServerMiddleware(PyObject* middleware, Vtable vtable)
   middleware_.reset(middleware);
 }
 
-Status PyServerMiddleware::SendingHeaders(
-    arrow::flight::AddCallHeaders& outgoing_headers) {
-  return SafeCallIntoPython([&] {
+void PyServerMiddleware::SendingHeaders(arrow::flight::AddCallHeaders& outgoing_headers) {
+  const Status& status = SafeCallIntoPython([&] {
     const Status status = vtable_.sending_headers(middleware_.obj(), outgoing_headers);
     RETURN_NOT_OK(CheckPyError());
     return status;
   });
+
+  if (!status.ok()) {
+    ARROW_LOG(WARNING) << "Python middleware failed in SendingHeaders: " << status;
+  }
 }
 
-Status PyServerMiddleware::CallCompleted(const Status& call_status) {
-  return SafeCallIntoPython([&] {
+void PyServerMiddleware::CallCompleted(const Status& call_status) {
+  const Status& status = SafeCallIntoPython([&] {
     const Status status = vtable_.call_completed(middleware_.obj(), call_status);
     RETURN_NOT_OK(CheckPyError());
     return status;
   });
+  if (!status.ok()) {
+    ARROW_LOG(WARNING) << "Python middleware failed in CallCompleted: " << status;
+  }
 }
 
 // Flight Client Middleware

cpp/src/arrow/python/flight.h

@@ -217,8 +217,8 @@ class ARROW_PYTHON_EXPORT PyServerMiddleware : public arrow::flight::ServerMiddl
   /// \brief Must only be called while holding the GIL.
   explicit PyServerMiddleware(PyObject* middleware, Vtable vtable);
 
-  Status SendingHeaders(arrow::flight::AddCallHeaders& outgoing_headers) override;
-  Status CallCompleted(const Status& status) override;
+  void SendingHeaders(arrow::flight::AddCallHeaders& outgoing_headers) override;
+  void CallCompleted(const Status& status) override;
 
  private:
   OwnedRefNoGIL middleware_;

python/pyarrow/_flight.pyx

@@ -1870,6 +1870,70 @@ cdef class ServerMiddleware:
         c_instance[0].reset(new CPyServerMiddleware(py_middleware, vtable))
 
 
+cdef class _ServerMiddlewareFactoryWrapper(ServerMiddlewareFactory):
+    """
+
+    gRPC runs RPCs on C++ threads. We use PyGilState_Ensure/Release to
+    interact with the GIL. However, this has a major shortcoming: it
+    breaks thread locals, since PyGilState_Release will clean up
+    Python's thread state (and clean up any thread locals). This makes
+    development of middleware harder than it needs to be.
+
+    Instead, we'll always register a middleware that always runs
+    first/last that artifically makes sure the GIL state stays valid.
+
+    TODO: after this, remove Status from middleware (it just mucks things up)
+
+    """
+
+    cdef:
+        list factories
+
+    def __init__(self, factories):
+        self.factories = factories
+
+    def start_call(self, info, headers):
+        cdef PyGILState_STATE state = PyGILState_Ensure()
+
+        instances = []
+        try:
+            for factory in self.factories:
+                instance = factory.start_call(info, headers)
+                if instance:
+                    instances.append(instance)
+        except:
+            PyGILState_Release(state)
+            raise
+
+        if not instances:
+            PyGILState_Release(state)
+            return
+
+        wrapper = _ServerMiddlewareWrapper(instances)
+        wrapper.state = state
+        return wrapper
+
+
+cdef class _ServerMiddlewareWrapper(ServerMiddleware):
+    cdef:
+        PyGILState_STATE state
+        list middleware
+
+    def __init__(self, list middleware):
+        self.middleware = middleware
+
+    def sending_headers(self):
+        # TODO: need to merge all results
+        pass
+
+    def call_completed(self, exception):
+        try:
+            for instance in self.middleware:
+                instance.call_completed(exception)
+        finally:
+            PyGILState_Release(self.state)
+
+
 cdef class FlightServerBase:
     """A Flight service definition.
 
@@ -1886,7 +1950,7 @@ cdef class FlightServerBase:
         An authentication mechanism to use. May be None.
     tls_certificates : list optional, default None
         A list of (certificate, key) pairs.
-    middleware : dict optional, default None
+    middleware : list optional, default None
         A dictionary of (middleware key, middleware factory) items.
     """
 
@@ -1911,14 +1975,13 @@ cdef class FlightServerBase:
         self.init(location, auth_handler, tls_certificates, middleware)
 
     cdef init(self, Location location, ServerAuthHandler auth_handler,
-              list tls_certificates, dict middleware):
+              list tls_certificates, list middleware):
         cdef:
             PyFlightServerVtable vtable = PyFlightServerVtable()
             PyFlightServer* c_server
             unique_ptr[CFlightServerOptions] c_options
             CCertKeyPair c_cert
-            pair[c_string, shared_ptr[CServerMiddlewareFactory]] \
-                middleware_pair
+            shared_ptr[CServerMiddlewareFactory] c_middleware
             function[cb_server_middleware_start_call] start_call = \
                 &_server_middleware_start_call
 
@@ -1938,12 +2001,16 @@ cdef class FlightServerBase:
                 c_options.get().tls_certificates.push_back(c_cert)
 
         if middleware:
-            # Takes advantage of dict ordering in Python>=3.6
-            for key, factory in middleware.items():
-                middleware_pair.first = tobytes(key)
-                middleware_pair.second.reset(
-                    new CPyServerMiddlewareFactory(factory, start_call))
-                c_options.get().middleware.push_back(middleware_pair)
+            py_middleware = _ServerMiddlewareFactoryWrapper(middleware)
+            c_middleware.reset(new CPyServerMiddlewareFactory(
+                py_middleware,
+                start_call))
+            c_options.get().middleware.push_back(c_middleware)
+            # for key, factory in middleware.items():
+            #     middleware_pair.first = tobytes(key)
+            #     middleware_pair.second.reset(
+            #         new CPyServerMiddlewareFactory(factory, start_call))
+            #     c_options.get().middleware.push_back(middleware_pair)
 
         vtable.list_flights = &_list_flights
         vtable.get_flight_info = &_get_flight_info

python/pyarrow/includes/libarrow_flight.pxd

@@ -259,7 +259,7 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
         CLocation location
         unique_ptr[CServerAuthHandler] auth_handler
         vector[CCertKeyPair] tls_certificates
-        vector[pair[c_string, shared_ptr[CServerMiddlewareFactory]]] middleware
+        vector[shared_ptr[CServerMiddlewareFactory]] middleware
 
     cdef cppclass CFlightClientOptions" arrow::flight::FlightClientOptions":
         CFlightClientOptions()
@@ -482,3 +482,12 @@ cdef extern from "<utility>" namespace "std":
     unique_ptr[CFlightDataStream] move(unique_ptr[CFlightDataStream]) nogil
     unique_ptr[CServerAuthHandler] move(unique_ptr[CServerAuthHandler]) nogil
     unique_ptr[CClientAuthHandler] move(unique_ptr[CClientAuthHandler]) nogil
+
+# We have to manually write these definitions since they're broken in
+# Cython (until 82a0ed4 is released)
+cdef extern from "Python.h":
+    ctypedef struct PyGILState_STATE:
+        pass
+
+    PyGILState_STATE PyGILState_Ensure()
+    void PyGILState_Release(PyGILState_STATE)