Rework C++ client middleware to require infallability

cpp/src/arrow/flight/client.cc

@@ -107,7 +107,7 @@ class GrpcAddCallHeaders : public AddCallHeaders {
 class GrpcClientInterceptorAdapter : public grpc::experimental::Interceptor {
  public:
   GrpcClientInterceptorAdapter(std::vector<std::unique_ptr<ClientMiddleware>> middleware)
-      : middleware_(std::move(middleware)), failed_() {}
+      : middleware_(std::move(middleware)) {}
 
   void Intercept(grpc::experimental::InterceptorBatchMethods* methods) {
     using InterceptionHookPoints = grpc::experimental::InterceptionHookPoints;
@@ -115,47 +115,26 @@ class GrpcClientInterceptorAdapter : public grpc::experimental::Interceptor {
             InterceptionHookPoints::PRE_SEND_INITIAL_METADATA)) {
       GrpcAddCallHeaders add_headers(methods->GetSendInitialMetadata());
       for (const auto& middleware : middleware_) {
-        const Status& status = middleware->SendingHeaders(add_headers);
-        if (!status.ok()) {
-          failed_ = status;
-          break;
-        }
+        middleware->SendingHeaders(add_headers);
       }
     }
 
     // TODO: also check for trailing metadata if call failed?
     // (grpc-java and grpc-core differ on this point, it seems)
-
-    // Don't run this hook if an earlier middleware already failed the call
-    if (failed_.ok() && methods->QueryInterceptionHookPoint(
-                            InterceptionHookPoints::POST_RECV_INITIAL_METADATA)) {
+    if (methods->QueryInterceptionHookPoint(
+            InterceptionHookPoints::POST_RECV_INITIAL_METADATA)) {
       internal::GrpcCallHeaders headers(methods->GetRecvInitialMetadata());
       for (const auto& middleware : middleware_) {
-        const Status& status = middleware->ReceivedHeaders(headers);
-        if (!status.ok()) {
-          failed_ = status;
-          break;
-        }
+        middleware->ReceivedHeaders(headers);
       }
     }
 
     if (methods->QueryInterceptionHookPoint(InterceptionHookPoints::POST_RECV_STATUS)) {
       DCHECK_NE(nullptr, methods->GetRecvStatus());
-      const Status status =
-          failed_.ok() ? internal::FromGrpcStatus(*methods->GetRecvStatus()) : failed_;
+      const Status status = internal::FromGrpcStatus(*methods->GetRecvStatus());
 
       for (const auto& middleware : middleware_) {
-        const Status& substatus = middleware->CallCompleted(status);
-        if (!substatus.ok()) {
-          // Abandon the rest
-          failed_ = substatus;
-          break;
-        }
-      }
-
-      // Overwrite the gRPC status with the interceptor status
-      if (!failed_.ok()) {
-        *methods->GetRecvStatus() = internal::ToGrpcStatus(failed_);
+        middleware->CallCompleted(status);
       }
     }
 
@@ -164,8 +143,6 @@ class GrpcClientInterceptorAdapter : public grpc::experimental::Interceptor {
 
  private:
   std::vector<std::unique_ptr<ClientMiddleware>> middleware_;
-  // Stores an Arrow failure status returned by middleware (if any)
-  Status failed_;
 };
 
 class GrpcClientInterceptorAdapterFactory
@@ -202,22 +179,7 @@ class GrpcClientInterceptorAdapterFactory
     const CallInfo flight_info{flight_method};
     for (auto& factory : middleware_) {
       std::unique_ptr<ClientMiddleware> instance;
-      const Status& status = factory->StartCall(flight_info, &instance);
-      if (!status.ok()) {
-        // Call was rejected here, but we don't have a formal way to
-        // stop the call with a status.
-        ARROW_LOG(INFO) << "Call cancelled: " << status;
-        info->client_context()->TryCancel();
-        for (auto& existing_instance : middleware) {
-          const Status& sub_status = existing_instance->CallCompleted(status);
-          if (!sub_status.ok()) {
-            // Aborted twice, log and bail
-            ARROW_LOG(INFO) << "Call cancelled while processing cancellation: " << status;
-            break;
-          }
-        }
-        return nullptr;
-      }
+      factory->StartCall(flight_info, &instance);
       if (instance) {
         middleware.push_back(std::move(instance));
       }

cpp/src/arrow/flight/client_middleware.h

@@ -28,17 +28,32 @@ namespace arrow {
 
 namespace flight {
 
+/// \brief Client-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 ClientMiddleware {
  public:
   virtual ~ClientMiddleware() = default;
 
-  /// While a non-OK status will do its best to stop the call, this
-  /// isn't guaranteed.
-  virtual Status SendingHeaders(AddCallHeaders& outgoing_headers) = 0;
-  virtual Status ReceivedHeaders(const CallHeaders& incoming_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 when headers are received from the server.
+  virtual void ReceivedHeaders(const CallHeaders& incoming_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 client, this will be called for each RPC (including
+/// Handshake) to give the opportunity to intercept the call.
+///
+/// It is guaranteed that all client middleware methods are called
+/// from the same thread that calls the RPC method implementation.
 class ARROW_FLIGHT_EXPORT ClientMiddlewareFactory {
  public:
   virtual ~ClientMiddlewareFactory() = default;
@@ -49,11 +64,8 @@ class ARROW_FLIGHT_EXPORT ClientMiddlewareFactory {
   /// \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 status. A non-OK status will try to stop the
-  /// call, but the underlying implementation may not be able to
-  /// support this.
-  virtual Status StartCall(const CallInfo& info,
-                           std::unique_ptr<ClientMiddleware>* middleware) = 0;
+  virtual void StartCall(const CallInfo& info,
+                         std::unique_ptr<ClientMiddleware>* middleware) = 0;
 };
 
 }  // namespace flight

cpp/src/arrow/flight/flight_test.cc

@@ -646,20 +646,13 @@ class PropagatingClientMiddleware : public ClientMiddleware {
                                        std::vector<Status>* recorded_status)
       : received_headers_(received_headers), recorded_status_(recorded_status) {}
 
-  Status SendingHeaders(AddCallHeaders& outgoing_headers) {
+  void SendingHeaders(AddCallHeaders& outgoing_headers) {
     outgoing_headers.AddHeader("x-tracing-span-id", current_span_id);
-    return Status::OK();
   }
 
-  Status ReceivedHeaders(const CallHeaders& incoming_headers) {
-    (*received_headers_)++;
-    return Status::OK();
-  }
+  void ReceivedHeaders(const CallHeaders& incoming_headers) { (*received_headers_)++; }
 
-  Status CallCompleted(const Status& status) {
-    recorded_status_->push_back(status);
-    return Status::OK();
-  }
+  void CallCompleted(const Status& status) { recorded_status_->push_back(status); }
 
  private:
   std::atomic<int>* received_headers_;
@@ -668,11 +661,10 @@ class PropagatingClientMiddleware : public ClientMiddleware {
 
 class PropagatingClientMiddlewareFactory : public ClientMiddlewareFactory {
  public:
-  Status StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
+  void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
     recorded_calls_.push_back(info.method);
     *middleware = arrow::internal::make_unique<PropagatingClientMiddleware>(
         &received_headers_, &recorded_status_);
-    return Status::OK();
   }
 
   void Reset() {

cpp/src/arrow/python/flight.cc

@@ -273,7 +273,7 @@ void PyServerMiddleware::SendingHeaders(arrow::flight::AddCallHeaders& outgoing_
   });
 
   if (!status.ok()) {
-    ARROW_LOG(WARNING) << "Python middleware failed in SendingHeaders: " << status;
+    ARROW_LOG(WARNING) << "Python server middleware failed in SendingHeaders: " << status;
   }
 }
 
@@ -284,7 +284,7 @@ void PyServerMiddleware::CallCompleted(const Status& call_status) {
     return status;
   });
   if (!status.ok()) {
-    ARROW_LOG(WARNING) << "Python middleware failed in CallCompleted: " << status;
+    ARROW_LOG(WARNING) << "Python server middleware failed in CallCompleted: " << status;
   }
 }
 
@@ -297,14 +297,17 @@ PyClientMiddlewareFactory::PyClientMiddlewareFactory(PyObject* factory,
   factory_.reset(factory);
 }
 
-Status PyClientMiddlewareFactory::StartCall(
+void PyClientMiddlewareFactory::StartCall(
     const arrow::flight::CallInfo& info,
     std::unique_ptr<arrow::flight::ClientMiddleware>* middleware) {
-  return SafeCallIntoPython([&] {
+  const Status& status = SafeCallIntoPython([&] {
     const Status status = start_call_(factory_.obj(), info, middleware);
     RETURN_NOT_OK(CheckPyError());
     return status;
   });
+  if (!status.ok()) {
+    ARROW_LOG(WARNING) << "Python client middleware failed in StartCall: " << status;
+  }
 }
 
 PyClientMiddleware::PyClientMiddleware(PyObject* middleware, Vtable vtable)
@@ -313,30 +316,38 @@ PyClientMiddleware::PyClientMiddleware(PyObject* middleware, Vtable vtable)
   middleware_.reset(middleware);
 }
 
-Status PyClientMiddleware::SendingHeaders(
-    arrow::flight::AddCallHeaders& outgoing_headers) {
-  return SafeCallIntoPython([&] {
+void PyClientMiddleware::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 client middleware failed in StartCall: " << status;
+  }
 }
 
-Status PyClientMiddleware::ReceivedHeaders(
+void PyClientMiddleware::ReceivedHeaders(
     const arrow::flight::CallHeaders& incoming_headers) {
-  return SafeCallIntoPython([&] {
+  const Status& status = SafeCallIntoPython([&] {
     const Status status = vtable_.received_headers(middleware_.obj(), incoming_headers);
     RETURN_NOT_OK(CheckPyError());
     return status;
   });
+  if (!status.ok()) {
+    ARROW_LOG(WARNING) << "Python client middleware failed in StartCall: " << status;
+  }
 }
 
-Status PyClientMiddleware::CallCompleted(const Status& call_status) {
-  return SafeCallIntoPython([&] {
+void PyClientMiddleware::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 client middleware failed in StartCall: " << status;
+  }
 }
 
 Status CreateFlightInfo(const std::shared_ptr<arrow::Schema>& schema,

cpp/src/arrow/python/flight.h

@@ -237,8 +237,8 @@ class ARROW_PYTHON_EXPORT PyClientMiddlewareFactory
   /// \brief Must only be called while holding the GIL.
   explicit PyClientMiddlewareFactory(PyObject* factory, StartCallCallback start_call);
 
-  Status StartCall(const arrow::flight::CallInfo& info,
-                   std::unique_ptr<arrow::flight::ClientMiddleware>* middleware) override;
+  void StartCall(const arrow::flight::CallInfo& info,
+                 std::unique_ptr<arrow::flight::ClientMiddleware>* middleware) override;
 
  private:
   OwnedRefNoGIL factory_;
@@ -264,9 +264,9 @@ class ARROW_PYTHON_EXPORT PyClientMiddleware : public arrow::flight::ClientMiddl
   /// \brief Must only be called while holding the GIL.
   explicit PyClientMiddleware(PyObject* factory, Vtable vtable);
 
-  Status SendingHeaders(arrow::flight::AddCallHeaders& outgoing_headers) override;
-  Status ReceivedHeaders(const arrow::flight::CallHeaders& incoming_headers) override;
-  Status CallCompleted(const Status& status) override;
+  void SendingHeaders(arrow::flight::AddCallHeaders& outgoing_headers) override;
+  void ReceivedHeaders(const arrow::flight::CallHeaders& incoming_headers) override;
+  void CallCompleted(const Status& status) override;
 
  private:
   OwnedRefNoGIL middleware_;