Don't rely on thread locals for server middleware

cpp/src/arrow/flight/flight_test.cc

@@ -591,6 +591,8 @@ class CountingServerMiddleware : public ServerMiddleware {
     }
   }
 
+  std::string name() const override { return "CountingServerMiddleware"; }
+
  private:
   std::atomic<int>* successful_;
   std::atomic<int>* failed_;
@@ -610,8 +612,25 @@ class CountingServerMiddlewareFactory : public ServerMiddlewareFactory {
   std::atomic<int> failed_;
 };
 
+// The current span ID, used to emulate OpenTracing style distributed
+// tracing. Only used for communication between application code and
+// client middleware.
 static thread_local std::string current_span_id = "";
 
+// A server middleware that stores the current span ID, in an
+// emulation of OpenTracing style distributed tracing.
+class TracingServerMiddleware : public ServerMiddleware {
+ public:
+  explicit TracingServerMiddleware(const std::string& current_span_id)
+      : span_id(current_span_id) {}
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {}
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "TracingServerMiddleware"; }
+
+  std::string span_id;
+};
+
 class TracingServerMiddlewareFactory : public ServerMiddlewareFactory {
  public:
   TracingServerMiddlewareFactory() {}
@@ -622,7 +641,7 @@ class TracingServerMiddlewareFactory : public ServerMiddlewareFactory {
         incoming_headers.equal_range("x-tracing-span-id");
     if (iter_pair.first != iter_pair.second) {
       const util::string_view& value = (*iter_pair.first).second;
-      current_span_id = std::string(value);
+      *middleware = std::make_shared<TracingServerMiddleware>(std::string(value));
     }
     return Status::OK();
   }
@@ -638,6 +657,11 @@ class PropagatingClientMiddleware : public ClientMiddleware {
       : received_headers_(received_headers), recorded_status_(recorded_status) {}
 
   void SendingHeaders(AddCallHeaders* outgoing_headers) {
+    // Pick up the span ID from thread locals. We have to use a
+    // thread-local for communication, since we aren't even
+    // instantiated until after the application code has already
+    // started the call (and so there's no chance for application code
+    // to pass us parameters directly).
     outgoing_headers->AddHeader("x-tracing-span-id", current_span_id);
   }
 
@@ -673,7 +697,13 @@ class ReportContextTestServer : public FlightServerBase {
   Status DoAction(const ServerCallContext& context, const Action& action,
                   std::unique_ptr<ResultStream>* result) override {
     std::shared_ptr<Buffer> buf;
-    RETURN_NOT_OK(Buffer::FromString(current_span_id, &buf));
+    const ServerMiddleware* middleware = context.GetMiddleware("tracing");
+    if (middleware == nullptr || middleware->name() != "TracingServerMiddleware") {
+      RETURN_NOT_OK(Buffer::FromString("", &buf));
+    } else {
+      RETURN_NOT_OK(
+          Buffer::FromString(((TracingServerMiddleware*)middleware)->span_id, &buf));
+    }
     *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
     return Status::OK();
   }
@@ -686,6 +716,13 @@ class PropagatingTestServer : public FlightServerBase {
 
   Status DoAction(const ServerCallContext& context, const Action& action,
                   std::unique_ptr<ResultStream>* result) override {
+    const ServerMiddleware* middleware = context.GetMiddleware("tracing");
+    if (middleware == nullptr || middleware->name() != "TracingServerMiddleware") {
+      current_span_id = "";
+    } else {
+      current_span_id = ((TracingServerMiddleware*)middleware)->span_id;
+    }
+
     return client_->DoAction(action, result);
   }
 
@@ -700,7 +737,7 @@ class TestRejectServerMiddleware : public ::testing::Test {
         &server_, &client_,
         [](FlightServerOptions* options) {
           options->middleware.push_back(
-              std::make_shared<RejectServerMiddlewareFactory>());
+              {"reject", std::make_shared<RejectServerMiddlewareFactory>()});
           return Status::OK();
         },
         [](FlightClientOptions* options) { return Status::OK(); }));
@@ -720,7 +757,7 @@ class TestCountingServerMiddleware : public ::testing::Test {
     ASSERT_OK(MakeServer<MetadataTestServer>(
         &server_, &client_,
         [&](FlightServerOptions* options) {
-          options->middleware.push_back(request_counter_);
+          options->middleware.push_back({"request_counter", request_counter_});
           return Status::OK();
         },
         [](FlightClientOptions* options) { return Status::OK(); }));
@@ -752,7 +789,7 @@ class TestPropagatingMiddleware : public ::testing::Test {
     ASSERT_OK(MakeServer<ReportContextTestServer>(
         &second_server_, &server_client,
         [&](FlightServerOptions* options) {
-          options->middleware.push_back(server_middleware_);
+          options->middleware.push_back({"tracing", server_middleware_});
           return Status::OK();
         },
         [&](FlightClientOptions* options) {
@@ -763,7 +800,7 @@ class TestPropagatingMiddleware : public ::testing::Test {
     ASSERT_OK(MakeServer<PropagatingTestServer>(
         &first_server_, &client_,
         [&](FlightServerOptions* options) {
-          options->middleware.push_back(server_middleware_);
+          options->middleware.push_back({"tracing", server_middleware_});
           return Status::OK();
         },
         [&](FlightClientOptions* options) {

cpp/src/arrow/flight/server.cc

@@ -252,11 +252,20 @@ class GrpcServerCallContext : public ServerCallContext {
     return internal::ToGrpcStatus(status);
   }
 
+  ServerMiddleware* GetMiddleware(const std::string& key) const override {
+    const auto& instance = middleware_map_.find(key);
+    if (instance == middleware_map_.end()) {
+      return nullptr;
+    }
+    return instance->second.get();
+  }
+
  private:
   friend class FlightServiceImpl;
   ServerContext* context_;
   std::string peer_identity_;
   std::vector<std::shared_ptr<ServerMiddleware>> middleware_;
+  std::unordered_map<std::string, std::shared_ptr<ServerMiddleware>> middleware_map_;
 };
 
 class GrpcAddCallHeaders : public AddCallHeaders {
@@ -278,11 +287,10 @@ class FlightServiceImpl : public FlightService::Service {
  public:
   explicit FlightServiceImpl(
       std::shared_ptr<ServerAuthHandler> auth_handler,
-      std::vector<std::shared_ptr<ServerMiddlewareFactory>> middleware,
+      std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
+          middleware,
       FlightServerBase* server)
-      : auth_handler_(auth_handler),
-        middleware_(std::move(middleware)),
-        server_(server) {}
+      : auth_handler_(auth_handler), middleware_(middleware), server_(server) {}
 
   template <typename UserType, typename Iterator, typename ProtoType>
   grpc::Status WriteStream(Iterator* iterator, ServerWriter<ProtoType>* writer) {
@@ -360,7 +368,7 @@ class FlightServiceImpl : public FlightService::Service {
     GrpcAddCallHeaders outgoing_headers(context);
     for (const auto& factory : middleware_) {
       std::shared_ptr<ServerMiddleware> instance;
-      Status result = factory->StartCall(info, incoming_headers, &instance);
+      Status result = factory.second->StartCall(info, incoming_headers, &instance);
       if (!result.ok()) {
         // Interceptor rejected call, end the request on all existing
         // interceptors
@@ -368,6 +376,7 @@ class FlightServiceImpl : public FlightService::Service {
       }
       if (instance != nullptr) {
         flight_context.middleware_.push_back(instance);
+        flight_context.middleware_map_.insert({factory.first, instance});
         instance->SendingHeaders(&outgoing_headers);
       }
     }
@@ -568,7 +577,8 @@ class FlightServiceImpl : public FlightService::Service {
 
  private:
   std::shared_ptr<ServerAuthHandler> auth_handler_;
-  std::vector<std::shared_ptr<ServerMiddlewareFactory>> middleware_;
+  std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
+      middleware_;
   FlightServerBase* server_;
 };
 
@@ -635,10 +645,9 @@ FlightServerBase::FlightServerBase() { impl_.reset(new Impl); }
 
 FlightServerBase::~FlightServerBase() {}
 
-Status FlightServerBase::Init(FlightServerOptions& options) {
-  std::shared_ptr<ServerAuthHandler> handler = std::move(options.auth_handler);
+Status FlightServerBase::Init(const FlightServerOptions& options) {
   impl_->service_.reset(
-      new FlightServiceImpl(handler, std::move(options.middleware), this));
+      new FlightServiceImpl(options.auth_handler, options.middleware, this));
 
   grpc::ServerBuilder builder;
   // Allow uploading messages of any length

cpp/src/arrow/flight/server.h

@@ -102,6 +102,10 @@ 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 Look up a middleware by key. Do not maintain a reference
+  /// to the object beyond the request body.
+  /// \return The middleware, or nullptr if not found.
+  virtual ServerMiddleware* GetMiddleware(const std::string& key) const = 0;
 };
 
 class ARROW_FLIGHT_EXPORT FlightServerOptions {
@@ -114,14 +118,17 @@ class ARROW_FLIGHT_EXPORT FlightServerOptions {
   /// Use port 0 to bind to an available port.
   Location location;
   /// \brief The authentication handler to use.
-  std::unique_ptr<ServerAuthHandler> auth_handler;
+  std::shared_ptr<ServerAuthHandler> auth_handler;
   /// \brief A list of TLS certificate+key pairs to use.
   std::vector<CertKeyPair> tls_certificates;
 
-  /// \brief A list of server middleware to apply.
+  /// \brief A list of server middleware to apply, along with a key to
+  /// identify them by.
   ///
-  /// Middleware are always applied in the order provided.
-  std::vector<std::shared_ptr<ServerMiddlewareFactory>> middleware;
+  /// Middleware are always applied in the order provided. Duplicate
+  /// keys are an error.
+  std::vector<std::pair<std::string, std::shared_ptr<ServerMiddlewareFactory>>>
+      middleware;
 
   /// \brief A Flight implementation-specific callback to customize
   /// transport-specific options.
@@ -145,7 +152,7 @@ class ARROW_FLIGHT_EXPORT FlightServerBase {
   /// \brief Initialize a Flight server listening at the given location.
   /// This method must be called before any other method.
   /// \param[in] options The configuration for this server.
-  Status Init(FlightServerOptions& options);
+  Status Init(const FlightServerOptions& options);
 
   /// \brief Get the port that the Flight server is listening on.
   /// This method must only be called after Init().  Will return a

cpp/src/arrow/flight/server_middleware.h

@@ -21,6 +21,7 @@
 #pragma once
 
 #include <memory>
+#include <string>
 
 #include "arrow/flight/middleware.h"
 #include "arrow/flight/visibility.h"  // IWYU pragma: keep
@@ -37,6 +38,10 @@ class ARROW_FLIGHT_EXPORT ServerMiddleware {
  public:
   virtual ~ServerMiddleware() = default;
 
+  /// \brief Unique name of middleware, used as alternative to RTTI
+  /// \return the string name of the middleware
+  virtual std::string name() const = 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;

cpp/src/arrow/python/flight.cc

@@ -29,6 +29,8 @@ namespace arrow {
 namespace py {
 namespace flight {
 
+const char* kPyServerMiddlewareName = "arrow.py_server_middleware";
+
 PyServerAuthHandler::PyServerAuthHandler(PyObject* handler,
                                          const PyServerAuthHandlerVtable& vtable)
     : vtable_(vtable) {
@@ -288,6 +290,10 @@ void PyServerMiddleware::CallCompleted(const Status& call_status) {
   }
 }
 
+std::string PyServerMiddleware::name() const { return kPyServerMiddlewareName; }
+
+PyObject* PyServerMiddleware::py_object() const { return middleware_.obj(); }
+
 // Flight Client Middleware
 
 PyClientMiddlewareFactory::PyClientMiddlewareFactory(PyObject* factory,

cpp/src/arrow/python/flight.h

@@ -33,6 +33,8 @@ namespace py {
 
 namespace flight {
 
+extern const char* kPyServerMiddlewareName;
+
 /// \brief A table of function pointers for calling from C++ into
 /// Python.
 class ARROW_PYTHON_EXPORT PyFlightServerVtable {
@@ -219,6 +221,9 @@ class ARROW_PYTHON_EXPORT PyServerMiddleware : public arrow::flight::ServerMiddl
 
   void SendingHeaders(arrow::flight::AddCallHeaders* outgoing_headers) override;
   void CallCompleted(const Status& status) override;
+  std::string name() const override;
+  /// \brief Get the underlying Python object.
+  PyObject* py_object() const;
 
  private:
   OwnedRefNoGIL middleware_;

docs/source/cpp/api/flight.rst

@@ -47,10 +47,6 @@ Common Types
    :project: arrow_cpp
    :members:
 
-.. doxygenclass:: arrow::flight::CallHeaders
-   :project: arrow_cpp
-   :members:
-
 .. doxygenstruct:: arrow::flight::Criteria
    :project: arrow_cpp
    :members:

python/pyarrow/_flight.pyx

@@ -1218,6 +1218,24 @@ cdef class ServerCallContext:
         """
         return tobytes(self.context.peer_identity())
 
+    def get_middleware(self, key):
+        """
+        Get a middleware instance by key.
+
+        Returns None if the middleware was not found.
+        """
+        cdef:
+            CServerMiddleware* c_middleware = \
+                self.context.GetMiddleware(CPyServerMiddlewareName)
+            CPyServerMiddleware* middleware
+        if c_middleware == NULL:
+            return None
+        if c_middleware.name() != CPyServerMiddlewareName:
+            return None
+        middleware = <CPyServerMiddleware*> c_middleware
+        py_middleware = <_ServerMiddlewareWrapper> middleware.py_object()
+        return py_middleware.middleware.get(key)
+
     @staticmethod
     cdef ServerCallContext wrap(const CServerCallContext& context):
         cdef ServerCallContext result = \
@@ -1987,64 +2005,37 @@ cdef class ServerMiddleware:
 
 
 cdef class _ServerMiddlewareFactoryWrapper(ServerMiddlewareFactory):
-    """A wrapper that bundles server middleware into a single C++ one.
-
-    gRPC runs RPCs on C++ threads, so we use PyGilState_Ensure/Release
-    to interact with the GIL before running Python code. However, this
-    has a major shortcoming: it breaks thread locals. When we enter
-    Python, it creates new thread state. When we exit, that thread
-    state is freed (and Python thread locals cleaned up), since the
-    thread isn't persistent from Python's point of view. This makes
-    development of middleware harder than it needs to be.
-
-    This middleware solves this by first incrementing the reference
-    count on the thread state, then running all Python middleware. It
-    then doesn't decrement the reference count until the call is
-    finished.
-
-    """
+    """Wrapper to bundle server middleware into a single C++ one."""
 
     cdef:
-        list factories
+        dict factories
 
-    def __init__(self, factories):
+    def __init__(self, dict 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:  # noqa: E722
-            # Don't artifically keep thread state around if an error
-            # happens.
-            # We want a bare except since we always want to release
-            # thread state.
-            PyGILState_Release(state)
-            raise
-
-        # Even if there's no middleware, we need to keep thread state
-        # alive
-        wrapper = _ServerMiddlewareWrapper(instances)
-        wrapper.state = state
-        return wrapper
+        instances = {}
+        for key, factory in self.factories.items():
+            instance = factory.start_call(info, headers)
+            if instance:
+                # TODO: prevent duplicate keys
+                instances[key] = instance
+        if instances:
+            wrapper = _ServerMiddlewareWrapper(instances)
+            return wrapper
+        return None
 
 
 cdef class _ServerMiddlewareWrapper(ServerMiddleware):
     cdef:
-        PyGILState_STATE state
-        list middleware
+        dict middleware
 
-    def __init__(self, list middleware):
+    def __init__(self, dict middleware):
         self.middleware = middleware
 
     def sending_headers(self):
         headers = collections.defaultdict(list)
-        for instance in self.middleware:
+        for instance in self.middleware.values():
             more_headers = instance.sending_headers()
             if not more_headers:
                 continue
@@ -2057,11 +2048,8 @@ cdef class _ServerMiddlewareWrapper(ServerMiddleware):
         return headers
 
     def call_completed(self, exception):
-        try:
-            for instance in self.middleware:
-                instance.call_completed(exception)
-        finally:
-            PyGILState_Release(self.state)
+        for instance in self.middleware.values():
+            instance.call_completed(exception)
 
 
 cdef class FlightServerBase:
@@ -2081,7 +2069,10 @@ cdef class FlightServerBase:
     tls_certificates : list optional, default None
         A list of (certificate, key) pairs.
     middleware : list optional, default None
-        A dictionary of :class:`ServerMiddlewareFactory` items.
+        A dictionary of :class:`ServerMiddlewareFactory` items. The
+        keys are used to retrieve the middleware instance during calls
+        (see :meth:`ServerCallContext.get_middleware`).
+
     """
 
     cdef:
@@ -2105,15 +2096,15 @@ cdef class FlightServerBase:
         self.init(location, auth_handler, tls_certificates, middleware)
 
     cdef init(self, Location location, ServerAuthHandler auth_handler,
-              list tls_certificates, list middleware):
+              list tls_certificates, dict middleware):
         cdef:
             PyFlightServerVtable vtable = PyFlightServerVtable()
             PyFlightServer* c_server
             unique_ptr[CFlightServerOptions] c_options
             CCertKeyPair c_cert
-            shared_ptr[CServerMiddlewareFactory] c_middleware
             function[cb_server_middleware_start_call] start_call = \
                 &_server_middleware_start_call
+            pair[c_string, shared_ptr[CServerMiddlewareFactory]] c_middleware
 
         c_options.reset(new CFlightServerOptions(Location.unwrap(location)))
 
@@ -2132,7 +2123,8 @@ cdef class FlightServerBase:
 
         if middleware:
             py_middleware = _ServerMiddlewareFactoryWrapper(middleware)
-            c_middleware.reset(new CPyServerMiddlewareFactory(
+            c_middleware.first = CPyServerMiddlewareName
+            c_middleware.second.reset(new CPyServerMiddlewareFactory(
                 py_middleware,
                 start_call))
             c_options.get().middleware.push_back(c_middleware)

python/pyarrow/includes/libarrow_flight.pxd

@@ -24,6 +24,9 @@ from pyarrow.includes.libarrow cimport *
 
 
 cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
+    cdef char* CPyServerMiddlewareName\
+        " arrow::py::flight::kPyServerMiddlewareName"
+
     cdef cppclass CActionType" arrow::flight::ActionType":
         c_string type
         c_string description
@@ -188,6 +191,7 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
 
     cdef cppclass CServerCallContext" arrow::flight::ServerCallContext":
         c_string& peer_identity()
+        CServerMiddleware* GetMiddleware(const c_string& key)
 
     cdef cppclass CTimeoutDuration" arrow::flight::TimeoutDuration":
         CTimeoutDuration(double)
@@ -243,7 +247,7 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
         void AddHeader(const c_string& key, const c_string& value)
 
     cdef cppclass CServerMiddleware" arrow::flight::ServerMiddleware":
-        pass
+        c_string name()
 
     cdef cppclass CServerMiddlewareFactory\
             " arrow::flight::ServerMiddlewareFactory":
@@ -261,7 +265,7 @@ cdef extern from "arrow/flight/api.h" namespace "arrow" nogil:
         CLocation location
         unique_ptr[CServerAuthHandler] auth_handler
         vector[CCertKeyPair] tls_certificates
-        vector[shared_ptr[CServerMiddlewareFactory]] middleware
+        vector[pair[c_string, shared_ptr[CServerMiddlewareFactory]]] middleware
 
     cdef cppclass CFlightClientOptions" arrow::flight::FlightClientOptions":
         CFlightClientOptions()
@@ -440,6 +444,7 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil:
     cdef cppclass CPyServerMiddleware\
             " arrow::py::flight::PyServerMiddleware"(CServerMiddleware):
         CPyServerMiddleware(object middleware, PyServerMiddlewareVtable vtable)
+        void* py_object()
 
     cdef cppclass CPyServerMiddlewareFactory\
             " arrow::py::flight::PyServerMiddlewareFactory"\
@@ -484,12 +489,3 @@ 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)

python/pyarrow/tests/test_flight.py

@@ -416,34 +416,33 @@ class TokenClientAuthHandler(ClientAuthHandler):
         return self.token
 
 
-_thread_locals = threading.local()
-_thread_locals.sentinel = "wrong value"
+class HeaderServerMiddleware(ServerMiddleware):
+    """Expose a per-call value to the RPC method body."""
+
+    def __init__(self, special_value):
+        self.special_value = special_value
 
 
-class ThreadLocalServerMiddleware(ServerMiddleware):
-    def sending_headers(self):
-        assert _thread_locals.sentinel == "right value"
-
-    def call_completed(self, exception):
-        ThreadLocalServerMiddlewareFactory.last_sentinel = \
-            _thread_locals.sentinel
-
-
-class ThreadLocalServerMiddlewareFactory(ServerMiddlewareFactory):
-    last_sentinel = None
+class HeaderServerMiddlewareFactory(ServerMiddlewareFactory):
+    """Expose a per-call hard-coded value to the RPC method body."""
 
     def start_call(self, info, headers):
-        _thread_locals.sentinel = "right value"
-        return ThreadLocalServerMiddleware()
+        return HeaderServerMiddleware("right value")
 
 
-class ThreadLocalFlightServer(FlightServerBase):
+class HeaderFlightServer(FlightServerBase):
+    """Echo back the per-call hard-coded value."""
+
     def do_action(self, context, action):
-        sentinel = getattr(_thread_locals, "sentinel", "wrong value")
-        return iter([flight.Result(sentinel.encode())])
+        middleware = context.get_middleware("test")
+        if middleware:
+            return iter([flight.Result(middleware.special_value.encode())])
+        return iter([flight.Result("".encode())])
 
 
 class SelectiveAuthServerMiddlewareFactory(ServerMiddlewareFactory):
+    """Deny access to certain methods based on a header."""
+
     def start_call(self, info, headers):
         if info.method == flight.FlightMethod.LIST_ACTIONS:
             # No auth needed
@@ -457,7 +456,7 @@ class SelectiveAuthServerMiddlewareFactory(ServerMiddlewareFactory):
         if token != b"password":
             raise flight.FlightUnauthenticatedError("Invalid token")
 
-        _thread_locals.sentinel = "password"
+        return HeaderServerMiddleware(token.decode())
 
 
 class SelectiveAuthClientMiddlewareFactory(ClientMiddlewareFactory):
@@ -987,27 +986,27 @@ def test_do_put_independent_read_write():
 
 def test_server_middleware_same_thread():
     """Ensure that server middleware run on the same thread as the RPC."""
-    with ThreadLocalFlightServer(middleware=[
-            ThreadLocalServerMiddlewareFactory(),
-    ]) as server:
+    with HeaderFlightServer(middleware={
+            "test": HeaderServerMiddlewareFactory(),
+    }) as server:
         client = FlightClient(('localhost', server.port))
         results = list(client.do_action(flight.Action(b"test", b"")))
         assert len(results) == 1
         value = results[0].body.to_pybytes()
         assert b"right value" == value
-        assert "right value" == \
-            ThreadLocalServerMiddlewareFactory.last_sentinel
 
 
 def test_middleware_reject():
     """Test rejecting an RPC with server middleware."""
-    with ThreadLocalFlightServer(middleware=[
-            SelectiveAuthServerMiddlewareFactory(),
-    ]) as server:
+    with HeaderFlightServer(middleware={
+            "test": SelectiveAuthServerMiddlewareFactory(),
+    }) as server:
         client = FlightClient(('localhost', server.port))
+        # The middleware allows this through without auth.
         with pytest.raises(pa.ArrowNotImplementedError):
             list(client.list_actions())
 
+        # But not anything else.
         with pytest.raises(flight.FlightUnauthenticatedError):
             list(client.do_action(flight.Action(b"", b"")))