From 62621131bd5771279420bc8edde6f6012bac55b9 Mon Sep 17 00:00:00 2001 From: Alexis Placet <2400067+Alex-PLACET@users.noreply.github.com> Date: Thu, 5 Nov 2026 17:13:18 +0100 Subject: [PATCH] [C++][FlightRPC] Add an async gRPC Flight server --- cpp/examples/arrow/CMakeLists.txt | 14 + .../flight_async_generic_doget_example.cc | 379 ++++++++ .../flight_async_generic_doput_example.cc | 287 ++++++ cpp/src/arrow/flight/CMakeLists.txt | 42 + cpp/src/arrow/flight/README.md | 2 +- cpp/src/arrow/flight/api.h | 1 + cpp/src/arrow/flight/flight_data_decoder.cc | 324 +++++++ cpp/src/arrow/flight/flight_data_decoder.h | 243 +++++ cpp/src/arrow/flight/flight_internals_test.cc | 385 ++++++++ cpp/src/arrow/flight/meson.build | 59 +- cpp/src/arrow/flight/serialization_internal.h | 1 + cpp/src/arrow/flight/server.cc | 55 +- cpp/src/arrow/flight/server_async.cc | 156 ++++ cpp/src/arrow/flight/server_async.h | 381 ++++++++ cpp/src/arrow/flight/transport.cc | 27 + cpp/src/arrow/flight/transport.h | 9 + .../flight/transport/grpc/async_grpc_compat.h | 45 + .../transport/grpc/async_grpc_do_get.cc | 186 ++++ .../transport/grpc/async_grpc_do_get_test.cc | 329 +++++++ .../transport/grpc/async_grpc_do_put.cc | 255 ++++++ .../transport/grpc/async_grpc_do_put_test.cc | 762 +++++++++++++++ .../transport/grpc/async_grpc_exchange.cc | 123 +++ .../grpc/async_grpc_exchange_internal.h | 218 +++++ .../grpc/async_grpc_exchange_reader.cc | 229 +++++ .../grpc/async_grpc_exchange_test.cc | 206 +++++ .../grpc/async_grpc_exchange_writer.cc | 234 +++++ .../transport/grpc/async_grpc_handshake.cc | 82 ++ .../grpc/async_grpc_handshake_test.cc | 509 +++++++++++ .../transport/grpc/async_grpc_service.cc | 150 +++ .../transport/grpc/async_grpc_service.h | 81 ++ .../grpc/async_grpc_service_internal.h | 232 +++++ .../transport/grpc/async_grpc_service_test.cc | 268 ++++++ .../transport/grpc/async_grpc_streaming.cc | 254 +++++ .../grpc/async_grpc_streaming_test.cc | 304 ++++++ .../transport/grpc/async_grpc_test_fixtures.h | 865 ++++++++++++++++++ .../flight/transport/grpc/async_grpc_unary.cc | 124 +++ .../transport/grpc/async_grpc_unary_test.cc | 218 +++++ .../flight/transport/grpc/customize_grpc.h | 4 +- .../flight/transport/grpc/grpc_server.cc | 108 ++- .../transport/grpc/grpc_server_internal.h | 34 +- .../transport/grpc/serialization_internal.cc | 30 +- .../transport/grpc/serialization_internal.h | 30 + .../flight/transport/grpc/util_internal.cc | 23 +- .../flight/transport/grpc/util_internal.h | 7 + .../arrow/flight/transport_server_internal.cc | 56 ++ .../arrow/flight/transport_server_internal.h | 36 + cpp/src/arrow/flight/types.cc | 3 + cpp/src/arrow/flight/types.h | 2 + cpp/src/arrow/util/windows_fixup.h | 8 - 49 files changed, 8297 insertions(+), 83 deletions(-) create mode 100644 cpp/examples/arrow/flight_async_generic_doget_example.cc create mode 100644 cpp/examples/arrow/flight_async_generic_doput_example.cc create mode 100644 cpp/src/arrow/flight/flight_data_decoder.cc create mode 100644 cpp/src/arrow/flight/flight_data_decoder.h create mode 100644 cpp/src/arrow/flight/server_async.cc create mode 100644 cpp/src/arrow/flight/server_async.h create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_compat.h create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_do_get.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_do_get_test.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_do_put.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_do_put_test.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_exchange.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_internal.h create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_reader.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_test.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_writer.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_handshake.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_handshake_test.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_service.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_service.h create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_service_internal.h create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_service_test.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_streaming.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_streaming_test.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_test_fixtures.h create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_unary.cc create mode 100644 cpp/src/arrow/flight/transport/grpc/async_grpc_unary_test.cc diff --git a/cpp/examples/arrow/CMakeLists.txt b/cpp/examples/arrow/CMakeLists.txt index 82c075c51df1..6d9332464a51 100644 --- a/cpp/examples/arrow/CMakeLists.txt +++ b/cpp/examples/arrow/CMakeLists.txt @@ -127,6 +127,20 @@ if(ARROW_FLIGHT) "${CMAKE_CURRENT_BINARY_DIR}/helloworld.pb.cc" "${CMAKE_CURRENT_BINARY_DIR}/helloworld.grpc.pb.cc") + add_arrow_example(flight_async_generic_doget_example + EXTRA_LINK_LIBS + ${FLIGHT_EXAMPLES_LINK_LIBS} + gRPC::grpc++ + ${ARROW_PROTOBUF_LIBPROTOBUF} + ${GFLAGS_LIBRARIES}) + + add_arrow_example(flight_async_generic_doput_example + EXTRA_LINK_LIBS + ${FLIGHT_EXAMPLES_LINK_LIBS} + gRPC::grpc++ + ${ARROW_PROTOBUF_LIBPROTOBUF} + ${GFLAGS_LIBRARIES}) + if(ARROW_FLIGHT_SQL) if(ARROW_BUILD_SHARED AND ARROW_GRPC_USE_SHARED) set(FLIGHT_SQL_EXAMPLES_LINK_LIBS arrow_flight_sql_shared) diff --git a/cpp/examples/arrow/flight_async_generic_doget_example.cc b/cpp/examples/arrow/flight_async_generic_doget_example.cc new file mode 100644 index 000000000000..07dac1c296dc --- /dev/null +++ b/cpp/examples/arrow/flight_async_generic_doget_example.cc @@ -0,0 +1,379 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#ifdef _WIN32 +# include +#else +# include +#endif +#include +#if defined(__linux__) +# include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +// Serve DoGet over the async generic Flight transport. +// +// Deriving from AsyncGenericFlightServerBase (arrow/flight/server_async.h) +// selects the async generic gRPC transport. DoGet is answered by DoGetAsync: +// it returns a future for a stream of payloads (AsyncFlightDataStream) instead +// of the payloads themselves, so the gRPC callback thread serving the RPC +// never waits for data. +// +// DelayedAsyncStream shows what that is for: each batch is produced on the +// stream's own worker thread, --batch_delay_ms after the transport asks for +// it. The transport adds a callback to the future and returns, so the delay +// never blocks a gRPC callback thread (which would stall every RPC multiplexed +// onto it). With --batch_delay_ms=0 the delay is gone and this is a plain +// fast DoGet. +// +// Any ticket streams the same batches. Every other RPC - GetFlightInfo +// included - answers UNIMPLEMENTED, so call DoGet directly instead of +// resolving the flight first. +// +// Usage: +// flight-async-generic-doget-example --port=31337 --batch_delay_ms=200 +// flight-async-generic-doget-example --port=0 --demo --batches=5 +// +// Run with no arguments to print this message and exit. + +DEFINE_int32(port, 0, "Port to listen on (0 picks a free port)"); +DEFINE_int32(batches, 5, "Number of batches in the DoGet stream"); +DEFINE_int32(batch_delay_ms, 50, + "Delay before each batch, spent on the stream's own worker thread"); +DEFINE_bool(demo, false, "Run the DoGet self-check against the server, then exit"); + +namespace flight = arrow::flight; + +/// \brief The current process' id, portably. +int64_t CurrentPid() { +#ifdef _WIN32 + return static_cast(::GetCurrentProcessId()); +#else + return static_cast(::getpid()); +#endif +} + +/// \brief A reader with `num_batches` batches of three int64 values. +class BatchReader : public arrow::RecordBatchReader { + public: + BatchReader(std::shared_ptr schema, int32_t num_batches) + : schema_(std::move(schema)), num_batches_(num_batches) {} + + std::shared_ptr schema() const override { return schema_; } + + arrow::Status ReadNext(std::shared_ptr* batch) override { + if (next_batch_ >= num_batches_) { + *batch = nullptr; // end of stream + return arrow::Status::OK(); + } + arrow::Int64Builder builder; + for (int64_t i = 0; i < 3; i++) { + ARROW_RETURN_NOT_OK(builder.Append(3 * next_batch_ + i)); + } + ARROW_ASSIGN_OR_RAISE(auto array, builder.Finish()); + *batch = arrow::RecordBatch::Make(schema_, array->length(), {array}); + ++next_batch_; + return arrow::Status::OK(); + } + + private: + std::shared_ptr schema_; + int32_t num_batches_; + int32_t next_batch_ = 0; +}; + +/// \brief The stream's worker state. Heap-allocated and shared with the worker +/// thread: the transport deletes the stream as soon as a payload completes +/// (the end-of-stream payload is completed *on the worker*), so the worker must +/// not touch the stream object after that. Everything it needs lives here. +struct WorkerState { + std::unique_ptr stream; + std::shared_ptr schema; + int32_t delay_ms = 0; + std::mutex mutex; + std::condition_variable ready; + bool closed = false; + /// The payload the transport is waiting for, if one was requested. + arrow::Future> pending; + bool has_pending = false; + /// Whether the wrapped stream has been closed, so Close() closes it once. + bool stream_closed = false; + std::thread thread; +}; + +/// \brief The worker: wait for a payload request, sleep delay_ms, produce it, +/// complete the future the transport holds. +void RunWorker(std::shared_ptr state) { +#if defined(__linux__) + // Name the thread, so /proc//task/*/comm distinguishes this server's own + // per-stream workers (doget_worker) from gRPC's pool threads (event_engine / + // lifeguard / grpc_global_timer). Thread names are inherited, and gRPC names + // its threads, so an anonymous thread created from a gRPC callback would + // otherwise show up as another "event_engine". comm is capped at 15 + // characters. + prctl(PR_SET_NAME, "doget_worker", 0, 0, 0); +#endif + std::unique_lock lock(state->mutex); + // Read the schema on this thread: it is the only one allowed to touch the + // wrapped stream, and the transport asks for the schema first. + state->schema = state->stream->schema(); + while (!state->closed) { + state->ready.wait(lock, [&] { return state->has_pending || state->closed; }); + if (state->closed) { + break; + } + // Take the request: after this, only this thread resolves the future. + state->has_pending = false; + arrow::Future> future = state->pending; + lock.unlock(); + std::this_thread::sleep_for(std::chrono::milliseconds(state->delay_ms)); + lock.lock(); + const bool cancelled = state->closed; + lock.unlock(); + // Runs the transport's callback inline, on this thread: it serializes the + // payload, and on end of stream it calls Close() (hence the lock dance). + if (cancelled) { + future.MarkFinished(arrow::Status::Cancelled()); + } else { + auto payload = state->stream->Next(); + if (!payload.ok()) { + future.MarkFinished(payload.status()); + } else if (payload->ipc_message.metadata == nullptr) { + // End of stream: the async interface reports it as a nullopt. + future.MarkFinished(std::nullopt); + } else { + future.MarkFinished(std::move(*payload)); + } + } + lock.lock(); + } + // The worker is the only thread that touched the wrapped stream, so it is the + // one that closes it (see Close(), which only asks it to stop). + lock.unlock(); + ARROW_WARN_NOT_OK(state->stream->Close(), "closing the DoGet stream failed"); + lock.lock(); + state->ready.notify_all(); +} + +/// \brief The stream the async transport pulls from, with a sleep per batch. +/// +/// NextAsync() hands the transport a future and returns at once, the delay and +/// the serialization of the wrapped synchronous stream happen on this stream's +/// own worker thread. The transport keeps one payload in flight at a time -- +/// it asks for the next only once the write completed -- so a single worker is +/// enough. +/// +/// One thread per stream. A server with many slow streams may prefer a shared +/// timer or thread pool, swapping WorkerState's thread for a task on that pool. +class DelayedAsyncStream : public flight::AsyncFlightDataStream { + public: + DelayedAsyncStream(std::unique_ptr stream, int32_t delay_ms) + : state_(std::make_shared()) { + state_->stream = std::move(stream); + state_->delay_ms = delay_ms; + std::shared_ptr state = state_; + state_->thread = std::thread([state] { RunWorker(state); }); + } + + ~DelayedAsyncStream() override { + Stop(); + if (state_->thread.joinable()) { + if (state_->thread.get_id() == std::this_thread::get_id()) { + // Deleted from inside its own worker (the transport releases the + // stream as a payload completes): joining would deadlock. + state_->thread.detach(); + } else { + state_->thread.join(); + } + } + } + + std::shared_ptr schema() override { return state_->schema; } + + arrow::Future GetSchemaPayloadAsync() override { + // The schema is ready now: no delay, no worker. This runs before the first + // NextAsync(), so it cannot overlap the worker's use of the stream. + return arrow::Future::MakeFinished( + state_->stream->GetSchemaPayload()); + } + + arrow::Future> NextAsync() override { + arrow::Future> future = + arrow::Future>::Make(); + { + std::lock_guard lock(state_->mutex); + if (state_->closed) { + future.MarkFinished(arrow::Status::Cancelled()); + return future; + } + state_->pending = future; + state_->has_pending = true; + } + state_->ready.notify_one(); + return future; + } + + arrow::Status Close() override { + // Only ask the worker to stop: the wrapped stream belongs to the worker + // thread, which may be inside stream->Next() right now, and Next()/Close() + // both mutate the IPC reader. Closing from this (gRPC) thread would race + // with that call and can corrupt the stream. Stop() is also what resolves + // a payload the worker has not taken yet, so this stays non-blocking; the + // worker closes the stream after it stops producing. + Stop(); + return arrow::Status::OK(); + } + + private: + /// \brief Stop the worker, resolve a payload it has not taken yet. Idempotent. + /// Does not join: called by the transport, whose threads must not block on + /// the worker's sleep. + void Stop() { + arrow::Future> pending; + { + std::lock_guard lock(state_->mutex); + state_->closed = true; + if (state_->has_pending) { + state_->has_pending = false; + pending = state_->pending; + } + } + state_->ready.notify_all(); + if (pending.is_valid()) { + // Runs the transport's callback inline, it reacts by calling Close(). + pending.MarkFinished(arrow::Status::Cancelled()); + } + } + + std::shared_ptr state_; +}; + +/// \brief The server: DoGetAsync is the only handler overridden, so every other +/// RPC answers UNIMPLEMENTED. +class ExampleServer : public flight::AsyncGenericFlightServerBase { + public: + arrow::Future> DoGetAsync( + const flight::ServerCallContext& context, const flight::Ticket& request) override { + std::cout << "DoGet: ticket=" << request.ticket << std::endl; + auto schema = arrow::schema({arrow::field("value", arrow::int64())}); + auto reader = std::make_shared(schema, FLAGS_batches); + std::shared_ptr stream = + std::make_shared( + std::make_unique(reader), FLAGS_batch_delay_ms); + // The stream is ready here, but its payloads are not: the transport will + // wait on the futures NextAsync() returns. + return arrow::Future>::MakeFinished( + std::move(stream)); + } +}; + +/// \brief The --demo self-check: drain one DoGet stream and count its batches. +bool RunDemo(int port) { + auto location = flight::Location::ForGrpcTcp("127.0.0.1", port); + if (!location.ok()) { + std::cerr << "Failed to build client location: " << location.status() << std::endl; + return false; + } + auto client = flight::FlightClient::Connect(*location); + if (!client.ok()) { + std::cerr << "Failed to connect: " << client.status() << std::endl; + return false; + } + auto reader = (*client)->DoGet(flight::Ticket("demo")); + if (!reader.ok()) { + std::cerr << "DoGet failed: " << reader.status() << std::endl; + return false; + } + int64_t num_batches = 0; + while (true) { + auto chunk = (*reader)->Next(); + if (!chunk.ok()) { + std::cerr << "DoGet drain failed: " << chunk.status() << std::endl; + return false; + } + if (chunk->data == nullptr) break; // end of stream + ++num_batches; + } + if (num_batches == FLAGS_batches) { + std::cout << "PASS: DoGet streamed " << num_batches << " batches" << std::endl; + return true; + } + std::cout << "FAIL: DoGet streamed " << num_batches << " batches, expected " + << FLAGS_batches << std::endl; + return false; +} + +int main(int argc, char** argv) { + if (argc == 1) { + // As in the other examples: a bare run (e.g. from ctest) does not start a + // server, it just prints the usage. + std::cout << "Usage: " << argv[0] + << " [--port=PORT] [--batches=N] [--batch_delay_ms=MS] [--demo]" + << std::endl; + return EXIT_SUCCESS; + } + gflags::ParseCommandLineFlags(&argc, &argv, true); + + ExampleServer server; + auto location = flight::Location::ForGrpcTcp("0.0.0.0", FLAGS_port); + if (!location.ok()) { + std::cerr << location.status() << std::endl; + return EXIT_FAILURE; + } + auto status = server.Init(flight::FlightServerOptions(*location)); + if (!status.ok()) { + std::cerr << status.ToString() << std::endl; + return EXIT_FAILURE; + } + std::cout << "Serving Flight on " << server.location().ToString() << std::endl; + std::cout << "Server pid " << CurrentPid() << ", port " << server.port() << std::endl; + + if (FLAGS_demo) { + bool ok = RunDemo(server.port()); + ARROW_WARN_NOT_OK(server.Shutdown(), "Error shutting down server"); + ARROW_WARN_NOT_OK(server.Wait(), "Error waiting for server shutdown"); + return ok ? EXIT_SUCCESS : EXIT_FAILURE; + } + + status = server.SetShutdownOnSignals({SIGINT, SIGTERM}); + if (!status.ok()) { + std::cerr << status.ToString() << std::endl; + return EXIT_FAILURE; + } + status = server.Serve(); + if (!status.ok()) { + std::cerr << status.ToString() << std::endl; + return EXIT_FAILURE; + } + return EXIT_SUCCESS; +} diff --git a/cpp/examples/arrow/flight_async_generic_doput_example.cc b/cpp/examples/arrow/flight_async_generic_doput_example.cc new file mode 100644 index 000000000000..2adee004332b --- /dev/null +++ b/cpp/examples/arrow/flight_async_generic_doput_example.cc @@ -0,0 +1,287 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include +#ifdef _WIN32 +# include +#else +# include +#endif + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +// Serve DoPut over the async generic Flight transport. +// +// Deriving from AsyncGenericFlightServerBase (arrow/flight/server_async.h) +// selects the async generic gRPC transport. On this path DoPut is not answered +// by a handler method: the transport hands every decoded upload to the +// AsyncFlightDataListener the server class hands out (CreateDoPutListener), one +// listener per RPC. The callbacks are: +// +// * OnDescriptor(const FlightDescriptor&) - the descriptor of the upload; +// returning non-OK rejects the upload (the client sees that status when it +// closes the writer). +// * OnSchemaDecoded(std::shared_ptr) - from ipc::Listener, the decoded +// schema. +// * OnNext(FlightStreamChunk) - once per decoded message, chunk.data is the +// RecordBatch, or nullptr for metadata-only messages. +// +// All of them run on gRPC callback threads, so they must return promptly: a +// slow callback stalls the RPC and every other call that thread would serve. +// This listener only counts and prints, a server with real work would hand the +// chunk to its own queue or thread pool and return at once. +// +// Usage: +// flight-async-generic-doput-example --port=31337 +// flight-async-generic-doput-example --port=0 --demo +// +// Run with no arguments to print this message and exit. + +DEFINE_int32(port, 0, "Port to listen on (0 picks a free port)"); +DEFINE_bool(demo, false, "Run the DoPut self-check against the server, then exit"); + +namespace flight = arrow::flight; + +/// \brief The current process' id, portably. +int64_t CurrentPid() { +#ifdef _WIN32 + return static_cast(::GetCurrentProcessId()); +#else + return static_cast(::getpid()); +#endif +} + +constexpr int kNumBatches = 3; +constexpr int64_t kRowsPerBatch = 4; + +/// \brief Make a one-column batch of `num_rows` int64 values starting at `value`. +arrow::Result> MakeInt64Batch( + const std::shared_ptr& schema, int64_t value, int64_t num_rows) { + arrow::Int64Builder builder; + for (int64_t i = 0; i < num_rows; i++) { + ARROW_RETURN_NOT_OK(builder.Append(value + i)); + } + ARROW_ASSIGN_OR_RAISE(auto array, builder.Finish()); + return arrow::RecordBatch::Make(schema, array->length(), {array}); +} + +/// \brief The AsyncFlightDataListener serving one DoPut RPC: it records what it is +/// handed as the upload streams in. The counters are atomics because the +/// callbacks run on gRPC threads. +class CountingUploadListener : public flight::AsyncFlightDataListener { + public: + arrow::Future<> OnDescriptor(const flight::FlightDescriptor& descriptor) override { + std::cout << "DoPut descriptor: " << descriptor.ToString() << std::endl; + // Returning non-OK here would reject the upload, the client sees that + // status when it closes the writer. + return arrow::Future<>::MakeFinished(arrow::Status::OK()); + } + + arrow::Status OnSchemaDecoded(std::shared_ptr schema) override { + std::cout << "DoPut schema:"; + for (const auto& field : schema->fields()) { + std::cout << " " << field->name() << ":" << field->type()->ToString(); + } + std::cout << std::endl; + return arrow::Status::OK(); + } + + arrow::Future<> OnNext(flight::FlightStreamChunk chunk) override { + if (chunk.data != nullptr) { // null for metadata-only messages + num_batches_.fetch_add(1); + num_rows_.fetch_add(chunk.data->num_rows()); + } + return arrow::Future<>::MakeFinished(arrow::Status::OK()); + } + + /// \brief The upload ended, whichever way: here is where a consumer commits. + /// status == OK means the client ended it normally, anything else is why it + /// did not complete. + arrow::Future<> OnFinish(arrow::Status status) override { + finished_ = true; + if (status.ok()) { + std::cout << "DoPut finished: the client ended the upload" << std::endl; + } else { + std::cout << "DoPut finished: " << status.ToString() << std::endl; + } + return arrow::Future<>::MakeFinished(arrow::Status::OK()); + } + + int64_t num_batches() const { return num_batches_.load(); } + int64_t num_rows() const { return num_rows_.load(); } + bool finished() const { return finished_; } + + private: + std::atomic num_batches_{0}; + std::atomic num_rows_{0}; + std::atomic finished_{false}; +}; + +/// \brief The server. Deriving from AsyncGenericFlightServerBase selects the +/// async generic transport. Uploads are served by the listener the server +/// class itself hands out: the transport calls CreateDoPutListener() once per +/// DoPut RPC, the server-class counterpart of DoGetAsync. Returning nullptr +/// refuses uploads (they answer UNIMPLEMENTED). +/// +/// One listener per RPC, as overlapping uploads require. This demo has a +/// single --demo client, so it keeps the listener it handed out in order to +/// check it in main(), once the server has stopped. +class UploadServer : public flight::AsyncGenericFlightServerBase { + public: + std::shared_ptr CreateDoPutListener( + const flight::ServerCallContext& context) override { + auto listener = std::make_shared(); + // Called concurrently, once per DoPut RPC, from gRPC callback threads: an + // unsynchronised shared_ptr assignment here is a data race that corrupts + // the heap. The lock only serializes the bookkeeping for the demo's + // assertion; a real server keeps no such member. + std::lock_guard lock(listener_mutex_); + last_listener_ = listener; + return listener; + } + + /// The listener the last DoPut RPC got, for the demo's assertions. + std::shared_ptr last_listener() const { + std::lock_guard lock(listener_mutex_); + return last_listener_; + } + + private: + mutable std::mutex listener_mutex_; + std::shared_ptr last_listener_; +}; + +/// \brief The --demo client half: upload kNumBatches batches and close the +/// writer. What the listener saw is read back in main(), once the server has +/// stopped. +bool RunDemo(int port) { + auto location = flight::Location::ForGrpcTcp("127.0.0.1", port); + if (!location.ok()) { + std::cerr << "Failed to build client location: " << location.status() << std::endl; + return false; + } + auto client = flight::FlightClient::Connect(*location); + if (!client.ok()) { + std::cerr << "Failed to connect: " << client.status() << std::endl; + return false; + } + + auto schema = arrow::schema({arrow::field("value", arrow::int64())}); + auto put = (*client)->DoPut(flight::FlightDescriptor::Command("upload"), schema); + if (!put.ok()) { + std::cerr << "DoPut failed: " << put.status() << std::endl; + return false; + } + arrow::Status status; + for (int i = 0; i < kNumBatches && status.ok(); i++) { + auto batch = MakeInt64Batch(schema, i * kRowsPerBatch, kRowsPerBatch); + if (!batch.ok()) { + std::cerr << "Failed to make a batch: " << batch.status() << std::endl; + return false; + } + status = put->writer->WriteRecordBatch(**batch); + } + if (status.ok()) { + status = put->writer->DoneWriting(); + } + if (status.ok()) { + status = put->writer->Close(); + } + if (!status.ok()) { + std::cerr << "DoPut upload failed: " << status << std::endl; + return false; + } + return true; +} + +int main(int argc, char** argv) { + if (argc == 1) { + // As in the other examples: a bare run (e.g. from ctest) does not start a + // server, it just prints the usage. + std::cout << "Usage: " << argv[0] << " [--port=PORT] [--demo]" << std::endl; + return EXIT_SUCCESS; + } + gflags::ParseCommandLineFlags(&argc, &argv, true); + + UploadServer server; + + auto location = flight::Location::ForGrpcTcp("0.0.0.0", FLAGS_port); + if (!location.ok()) { + std::cerr << location.status() << std::endl; + return EXIT_FAILURE; + } + flight::FlightServerOptions options(*location); + // UploadServer derives from AsyncGenericFlightServerBase: the server class + // alone selects the async generic transport, so there is no option to set. + // Uploads come from the server's CreateDoPutListener() override. + auto status = server.Init(options); + if (!status.ok()) { + std::cerr << status.ToString() << std::endl; + return EXIT_FAILURE; + } + int port = server.port(); + std::cout << "Serving Flight on " << server.location().ToString() << std::endl; + std::cout << "Server pid " << CurrentPid() << ", port " << port << std::endl; + + if (FLAGS_demo) { + bool ok = RunDemo(port); + ARROW_WARN_NOT_OK(server.Shutdown(), "Error shutting down server"); + ARROW_WARN_NOT_OK(server.Wait(), "Error waiting for server shutdown"); + // The listener callbacks run on gRPC threads: read its counters only once + // the server has stopped (the same pattern the flight tests use). + const auto listener = server.last_listener(); + if (listener == nullptr) { + std::cout << "FAIL: the server handed out no DoPut listener" << std::endl; + ok = false; + } else if (listener->num_batches() == kNumBatches && + listener->num_rows() == kNumBatches * kRowsPerBatch && + listener->finished()) { + std::cout << "PASS: listener saw " << listener->num_batches() << " batches, " + << listener->num_rows() << " rows, and its OnFinish said the client " + << "ended the upload" << std::endl; + } else { + std::cout << "FAIL: listener saw " << listener->num_batches() << " batches, " + << listener->num_rows() << " rows, finished=" << listener->finished() + << ", expected " << kNumBatches << " batches, " + << kNumBatches * kRowsPerBatch << " rows, finished" << std::endl; + ok = false; + } + return ok ? EXIT_SUCCESS : EXIT_FAILURE; + } + + status = server.SetShutdownOnSignals({SIGINT, SIGTERM}); + if (!status.ok()) { + std::cerr << status.ToString() << std::endl; + return EXIT_FAILURE; + } + status = server.Serve(); + if (!status.ok()) { + std::cerr << status.ToString() << std::endl; + return EXIT_FAILURE; + } + return EXIT_SUCCESS; +} diff --git a/cpp/src/arrow/flight/CMakeLists.txt b/cpp/src/arrow/flight/CMakeLists.txt index 09eda0e24445..0e41da9bcf8a 100644 --- a/cpp/src/arrow/flight/CMakeLists.txt +++ b/cpp/src/arrow/flight/CMakeLists.txt @@ -162,6 +162,7 @@ else() endif() # Was in a different namespace, or simply not supported, prior to this +# (the async *client* paths in grpc_client.cc). if(ARROW_GRPC_VERSION VERSION_GREATER_EQUAL "1.40") add_definitions(-DGRPC_ENABLE_ASYNC) endif() @@ -179,9 +180,11 @@ set(ARROW_FLIGHT_SRCS client_cookie_middleware.cc client_tracing_middleware.cc cookie_internal.cc + flight_data_decoder.cc middleware.cc serialization_internal.cc server.cc + server_async.cc server_auth.cc server_tracing_middleware.cc transport.cc @@ -196,6 +199,23 @@ set(ARROW_FLIGHT_SRCS transport/grpc/util_internal.cc types.cc) +# The async generic *server* needs gRPC's generic callback API, which exists from +# 1.65 on (see transport/grpc/async_grpc_compat.h). This floor is higher than +# GRPC_ENABLE_ASYNC above, which only gates the async client. +if(ARROW_GRPC_VERSION VERSION_GREATER_EQUAL "1.65") + list(APPEND + ARROW_FLIGHT_SRCS + transport/grpc/async_grpc_service.cc + transport/grpc/async_grpc_do_get.cc + transport/grpc/async_grpc_do_put.cc + transport/grpc/async_grpc_exchange.cc + transport/grpc/async_grpc_exchange_reader.cc + transport/grpc/async_grpc_exchange_writer.cc + transport/grpc/async_grpc_handshake.cc + transport/grpc/async_grpc_streaming.cc + transport/grpc/async_grpc_unary.cc) +endif() + # Handle Unity build header conflicts on Windows. if(CMAKE_UNITY_BUILD AND WIN32) set_source_files_properties(client.cc @@ -206,6 +226,7 @@ if(CMAKE_UNITY_BUILD AND WIN32) transport/grpc/protocol_grpc_internal.cc transport/grpc/util_internal.cc types.cc + transport_server_internal.cc PROPERTIES SKIP_UNITY_BUILD_INCLUSION TRUE) endif() @@ -336,6 +357,27 @@ add_arrow_test(flight_test LABELS "arrow_flight") +# Async Flight server over the gRPC generic callback API (>= 1.65, like the +# sources it tests; higher than the async client floor above). +if(ARROW_GRPC_VERSION VERSION_GREATER_EQUAL "1.65") + add_arrow_test(flight_async_grpc_test + SOURCES + transport/grpc/async_grpc_service_test.cc + transport/grpc/async_grpc_do_get_test.cc + transport/grpc/async_grpc_do_put_test.cc + transport/grpc/async_grpc_unary_test.cc + transport/grpc/async_grpc_streaming_test.cc + transport/grpc/async_grpc_exchange_test.cc + transport/grpc/async_grpc_handshake_test.cc + STATIC_LINK_LIBS + ${ARROW_FLIGHT_TEST_LINK_LIBS} + LABELS + "arrow_flight") + if(TEST arrow-flight-async-grpc-test) + set_tests_properties(arrow-flight-async-grpc-test PROPERTIES TIMEOUT 300) + endif() +endif() + # Build test server for unit tests or benchmarks if(ARROW_BUILD_TESTS OR ARROW_BUILD_BENCHMARKS) add_executable(flight-test-server test_server.cc) diff --git a/cpp/src/arrow/flight/README.md b/cpp/src/arrow/flight/README.md index 5156973acf33..f23f88c78dfe 100644 --- a/cpp/src/arrow/flight/README.md +++ b/cpp/src/arrow/flight/README.md @@ -33,4 +33,4 @@ your current working directory or you need to add it to your path, e.g. ``` PATH=debug:$PATH debug/flight-test -``` \ No newline at end of file +``` diff --git a/cpp/src/arrow/flight/api.h b/cpp/src/arrow/flight/api.h index ed31b5c8fa41..c3b0e307d7a7 100644 --- a/cpp/src/arrow/flight/api.h +++ b/cpp/src/arrow/flight/api.h @@ -23,6 +23,7 @@ #include "arrow/flight/client_tracing_middleware.h" #include "arrow/flight/middleware.h" #include "arrow/flight/server.h" +#include "arrow/flight/server_async.h" #include "arrow/flight/server_auth.h" #include "arrow/flight/server_middleware.h" #include "arrow/flight/server_tracing_middleware.h" diff --git a/cpp/src/arrow/flight/flight_data_decoder.cc b/cpp/src/arrow/flight/flight_data_decoder.cc new file mode 100644 index 000000000000..6d60ceb9692a --- /dev/null +++ b/cpp/src/arrow/flight/flight_data_decoder.cc @@ -0,0 +1,324 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/flight_data_decoder.h" + +#include +#include + +#include "arrow/flight/serialization_internal.h" +#include "arrow/flight/transport.h" +#include "arrow/ipc/message.h" +#include "arrow/ipc/reader.h" +#include "arrow/record_batch.h" +#include "arrow/result.h" +#include "arrow/status.h" + +namespace arrow::flight { + +namespace { + +// FlightDataMessageReader is an ipc::MessageReader that accepts one at a time +// FlightData messages. Analogous to MessageReader::Open(InputStream*) but for +// individual FlightData messages directly read from the received buffers. +class FlightDataMessageReader : public ipc::MessageReader { + public: + /// \brief Queue one message, with the app metadata it arrived with. + /// \param message The message to queue. + /// \param app_metadata The application metadata associated with the message. + void Push(std::unique_ptr message, std::shared_ptr app_metadata) { + queue_.emplace_back(std::move(message), std::move(app_metadata)); + } + + /// \brief Check if there is a record batch message queued. + /// \return True if a record batch message is queued, false otherwise. + bool HasRecordBatch() const { + return std::ranges::any_of(queue_, [](const QueuedMessage& queued) { + return queued.message->type() == ipc::MessageType::RECORD_BATCH; + }); + } + + /// \brief Take the app metadata of the last queued message, leaving that message + /// with none: its app metadata has been delivered already. + /// \return The app metadata of the last queued message, or nullptr if no messages are + /// queued. + std::shared_ptr TakeLastAppMetadata() { + if (queue_.empty()) { + return nullptr; + } + return std::move(queue_.back().app_metadata); + } + + /// Take the next queued message, if any, and return it. The app metadata of + /// the message taken is stored internally and can be retrieved with + /// ReadAppMetadata(). + /// \return The next queued message, or nullptr if no messages are queued. + ::arrow::Result> ReadNextMessage() override { + if (queue_.empty()) { + return nullptr; + } + QueuedMessage queued = std::move(queue_.front()); + queue_.pop_front(); + app_metadata_ = std::move(queued.app_metadata); + return std::move(queued.message); + } + + /// \brief Read the application metadata of the last message taken by ReadNextMessage(). + /// \return The application metadata of the last message taken, or nullptr if no + /// messages have been taken. + std::shared_ptr ReadAppMetadata() { return app_metadata_; } + + private: + /// \brief Queue of messages that have been pushed but not yet read. + struct QueuedMessage { + QueuedMessage() = default; + QueuedMessage(std::unique_ptr message, + std::shared_ptr app_metadata) + : message(std::move(message)), app_metadata(std::move(app_metadata)) {} + std::unique_ptr message; + std::shared_ptr app_metadata; + }; + + std::deque queue_; + std::shared_ptr app_metadata_; +}; + +// FlightDataMessageReader is owned by the decoder, not by the +// RecordBatchStreamReader that reads through it: Open() takes a unique_ptr and +// adopts it, so handing over the decoder's reader would free it on any failure +// path (a malformed first message, a first message that is not a schema), +// leaving the decoder with a dangling pointer. This forwards to the decoder's +// reader and stays owned by it. +class FlightDataMessageReaderRef final : public ipc::MessageReader { + public: + explicit FlightDataMessageReaderRef(FlightDataMessageReader* reader) + : reader_(reader) {} + + arrow::Result> ReadNextMessage() override { + return reader_->ReadNextMessage(); + } + + private: + FlightDataMessageReader* reader_; +}; + +} // namespace + +class AsyncFlightMessageDecoder::AsyncFlightMessageDecoderImpl { + public: + AsyncFlightMessageDecoderImpl(std::shared_ptr listener, + ipc::IpcReadOptions options) + : listener_(std::move(listener)), + options_(std::move(options)), + message_reader_(std::make_unique()) {} + + /// \brief Consume a chunk of Flight data. + /// \param data The Flight data to consume. + /// \return A future carrying the listener's status for this message: a + /// non-OK result rejects the upload. + /// The descriptor of an upload is eported before the rest of its first message is + /// decoded. + Future<> ConsumeData(internal::FlightData data) { + if (!data.descriptor) { + return ConsumeDataAfterDescriptor(std::move(data)); + } else { + return listener_->OnDescriptor(*data.descriptor) + .Then([this, data = std::move(data)]() mutable { + return ConsumeDataAfterDescriptor(std::move(data)); + }); + } + } + + std::shared_ptr schema() const { + return batch_reader_ ? batch_reader_->schema() : nullptr; + } + + private: + /// \brief Decode a FlightData message whose descriptor (if any) was already + /// reported to the listener. + /// \param data The Flight data to consume after the descriptor has been reported. + /// \return A future carrying the listener's status for this message: a non-OK result + /// rejects the upload. + /// \note The descriptor of an upload is assumed to have been reported + /// before calling this method. + /// + Future<> ConsumeDataAfterDescriptor(internal::FlightData data) { + if (!data.metadata) { + // Metadata-only message: no IPC content, just Flight app_metadata. + if (data.app_metadata && data.app_metadata->size() > 0) { + FlightStreamChunk chunk; + chunk.app_metadata = std::move(data.app_metadata); + return listener_->OnNext(std::move(chunk)); + } + return Status::OK(); + } + + ARROW_ASSIGN_OR_RAISE(auto message, data.OpenMessage()); + + if (!batch_reader_) { + // Initialize RecordBatchStreamReader and read the first IPC message. + // It must be a schema. + message_reader_->Push(std::move(message), std::move(data.app_metadata)); + ARROW_ASSIGN_OR_RAISE( + batch_reader_, + ipc::RecordBatchStreamReader::Open( + std::make_unique(message_reader_.get()), + options_)); + return listener_->OnSchemaDecoded(batch_reader_->schema()); + } + + message_reader_->Push(std::move(message), std::move(data.app_metadata)); + + if (!message_reader_->HasRecordBatch()) { + auto app_metadata = message_reader_->TakeLastAppMetadata(); + if (app_metadata && app_metadata->size() > 0) { + FlightStreamChunk chunk; + chunk.app_metadata = std::move(app_metadata); + return listener_->OnNext(std::move(chunk)); + } + return arrow::Future<>::MakeFinished(Status::OK()); + } + + std::shared_ptr batch; + RETURN_NOT_OK(batch_reader_->ReadNext(&batch)); + auto app_metadata = message_reader_->ReadAppMetadata(); + + if (batch) { + FlightStreamChunk chunk{std::move(batch), std::move(app_metadata)}; + return listener_->OnNext(std::move(chunk)); + } + + // This has to be a Dictionary batch. + if (app_metadata && app_metadata->size() > 0) { + FlightStreamChunk chunk{nullptr, std::move(app_metadata)}; + return listener_->OnNext(std::move(chunk)); + } + return arrow::Future<>::MakeFinished(Status::OK()); + } + + std::shared_ptr listener_; + ipc::IpcReadOptions options_; + std::unique_ptr message_reader_; + std::shared_ptr batch_reader_; +}; + +AsyncFlightMessageDecoder::AsyncFlightMessageDecoder( + std::shared_ptr listener, ipc::IpcReadOptions options) + : impl_(std::make_unique(std::move(listener), + std::move(options))) {} + +AsyncFlightMessageDecoder::~AsyncFlightMessageDecoder() = default; + +Future<> AsyncFlightMessageDecoder::Consume(std::shared_ptr buffer) { + ARROW_ASSIGN_OR_RAISE(auto data, internal::DeserializeFlightData(buffer)); + return impl_->ConsumeData(std::move(data)); +} + +Future<> AsyncFlightMessageDecoder::Consume(internal::FlightData data) { + return impl_->ConsumeData(std::move(data)); +} + +std::shared_ptr AsyncFlightMessageDecoder::schema() const { + return impl_->schema(); +} + +// --- AsyncFlightDataListener: transport state and the terminal callbacks --- + +AsyncFlightDataListener::AsyncFlightDataListener() = default; +AsyncFlightDataListener::~AsyncFlightDataListener() = default; +AsyncFlightDataWriter::~AsyncFlightDataWriter() = default; + +Status AsyncFlightDataWriter::WriteMetadata(const Buffer& app_metadata) { + // Install() and Clear() swap transport_ under writer_mutex_, so both the read + // and the call must happen under it: otherwise this can read the old pointer + // just before DoPutReactor::OnDone() runs Clear(), and then call through a + // destructed reactor. The transport's own WriteMetadata only takes its + // private write_mutex_, so there is no lock-order cycle. + std::lock_guard lock(writer_mutex_); + if (transport_ == nullptr) { + return Status::Invalid( + "no upload in flight to write metadata to: the writer is only usable " + "while the upload's RPC is running"); + } + return transport_->WriteMetadata(app_metadata); +} + +namespace internal { + +FlightDataListenerTransport::~FlightDataListenerTransport() = default; + +void FlightDataListenerTransport::Install( + const std::shared_ptr& listener, + FlightDataListenerTransport* transport) { + if (listener == nullptr) { + return; + } + std::lock_guard lock(listener->transport_mutex_); + listener->transport_ = transport; + if (listener->writer_ != nullptr) { + std::lock_guard writer_lock(listener->writer_->writer_mutex_); + listener->writer_->transport_ = transport; + } +} + +void FlightDataListenerTransport::Clear( + const std::shared_ptr& listener) { + if (listener == nullptr) { + return; + } + std::lock_guard lock(listener->transport_mutex_); + listener->transport_ = nullptr; + if (listener->writer_ != nullptr) { + std::lock_guard writer_lock(listener->writer_->writer_mutex_); + listener->writer_->transport_ = nullptr; + } +} + +Future<> FlightDataListenerTransport::ReportFinish( + const std::shared_ptr& listener, Status status) { + if (listener == nullptr) { + return Future<>::MakeFinished(); + } + // Report exactly once: the first ending wins. + bool expected = false; + if (!listener->finished_.compare_exchange_strong(expected, true)) { + return Future<>::MakeFinished(); + } + return listener->OnFinish(std::move(status)); +} + +} // namespace internal + +Future<> AsyncFlightDataListener::Cancel(Status status) { + internal::FlightDataListenerTransport* transport; + { + std::lock_guard lock(transport_mutex_); + transport = transport_; + if (transport == nullptr) { + return arrow::Future<>::MakeFinished(Status::Invalid( + "no upload in flight to cancel: Cancel() must be called while the " + "upload's RPC is running")); + } + transport->Hold(); + } + // Unlocked, so OnFinish() may re-enter Cancel() without deadlocking. + transport->CancelUpload(std::move(status)); + transport->ReleaseHold(); + return arrow::Future<>::MakeFinished(Status::OK()); +} + +} // namespace arrow::flight diff --git a/cpp/src/arrow/flight/flight_data_decoder.h b/cpp/src/arrow/flight/flight_data_decoder.h new file mode 100644 index 000000000000..315547371274 --- /dev/null +++ b/cpp/src/arrow/flight/flight_data_decoder.h @@ -0,0 +1,243 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include + +#include "arrow/flight/types.h" +#include "arrow/flight/visibility.h" +#include "arrow/ipc/options.h" +#include "arrow/ipc/reader.h" +#include "arrow/status.h" + +namespace arrow::flight { + +namespace internal { +struct FlightData; +} // namespace internal + +class AsyncFlightDataListener; + +namespace internal { + +/// \brief The per-RPC handle a transport gives an AsyncFlightDataListener so the +/// application can act on the RPC while it is in flight. +/// +/// The transport owns this, the listener only holds a reference for as long as +/// the RPC is live, and asks for it under a lock the transport takes to clear +/// it (see FlightDataListenerTransport::Clear) +class ARROW_FLIGHT_EXPORT FlightDataListenerTransport { + public: + virtual ~FlightDataListenerTransport(); + + /// \brief Finish the upload's RPC with `status`, from any thread. + /// \param status The status to finish the upload with. + virtual void CancelUpload(Status status) = 0; + + /// \brief Send one application metadata message to the client, from any thread. + /// \param app_metadata The metadata to attach to the next PutResult message. + /// \return OK when the message was handed to the transport, or a non-OK status + /// when the upload is already over. + virtual Status WriteMetadata(const Buffer& app_metadata) = 0; + + /// \brief Pin this transport's lifetime, from any thread. + /// + /// Install()/Clear() publish and retire the transport through the listener's + /// mutex, but a caller that reads the pointer and then releases the mutex + /// before calling in (as Cancel() must, so OnFinish() can re-enter) still + /// races with Clear() and the reactor's deletion. Hold under the mutex and + /// pair with ReleaseHold() around the call so the object stays alive. + virtual void Hold() = 0; + + /// \brief Drop a reference taken with Hold(). Deletes the transport at zero. + virtual void ReleaseHold() = 0; + + /// \brief Install `transport` on `listener` for the life of the RPC. + /// \param listener The listener to install the transport on. + /// \param transport The transport to install. + static void Install(const std::shared_ptr& listener, + FlightDataListenerTransport* transport); + + /// \brief Clear the installed transport. + /// \param listener The listener to clear the transport from. + static void Clear(const std::shared_ptr& listener); + + /// \brief Report the upload's terminal status to the listener, once. + /// \param listener The listener to report the status to. + /// \param status The terminal status of the upload. + /// \return the listener's OnFinish() future. The caller must not block on it + /// (this runs on a gRPC callback thread) but should chain the RPC's + /// completion on it, so an asynchronous failure to commit is reported to the + /// client instead of a success. Returns a finished future when the status + /// was already reported. + static Future<> ReportFinish(const std::shared_ptr& listener, + Status status); +}; + +} // namespace internal + +/// \brief Sends application-specific metadata back to the client during an upload. +/// +/// The async counterpart of FlightMetadataWriter: the transport hands one to the +/// AsyncFlightDataListener at the start of an upload (AsyncFlightDataListener::writer) +/// and it stays valid for as long as the upload's RPC runs. Every call sends one +/// PutResult message immediately, the client reads them with ReadMetadata() on +/// its DoPutResult. +class ARROW_FLIGHT_EXPORT AsyncFlightDataWriter { + public: + virtual ~AsyncFlightDataWriter(); + + /// \brief Send a metadata message to the client. + /// + /// Safe to call from any thread while the upload is in flight. Once the + /// upload has ended the call reports the ending (Cancelled or FailedPrecondition) + /// instead of writing to a finished RPC. + /// \param app_metadata The application metadata to send. + /// \return Status indicating success or failure. + Status WriteMetadata(const Buffer& app_metadata); + + private: + friend class internal::FlightDataListenerTransport; + + /// The transport's RPC state, set and cleared through + /// FlightDataListenerTransport::Install/Clear, like + /// AsyncFlightDataListener::transport_. + internal::FlightDataListenerTransport* transport_ = nullptr; + /// Guards transport_, since WriteMetadata may be called from any thread while + /// the transport installs or clears the pointer. + mutable std::mutex writer_mutex_; +}; + +/// \brief A general listener class to receive events from FlightMessageDecoder +/// +/// User must implement callback methods for interested events. +class ARROW_FLIGHT_EXPORT AsyncFlightDataListener : public ipc::Listener { + public: + AsyncFlightDataListener(); + ~AsyncFlightDataListener() override; + + /// \brief Called for each decoded FlightStreamChunk. + /// \param chunk The decoded FlightStreamChunk. + /// \return A future that completes when the chunk has been processed. + virtual Future<> OnNext(FlightStreamChunk chunk) = 0; + + /// \brief Called when the descriptor of an upload is decoded. + /// + /// Fired before any schema or data of that upload, so the listener knows + /// which upload it is being handed. A non-OK status rejects the upload. + /// \param descriptor The decoded FlightDescriptor. + /// \return A future that completes when the descriptor has been processed. + virtual Future<> OnDescriptor(const FlightDescriptor& descriptor) { + return Future<>::MakeFinished(); + } + + /// \brief Called once, when the upload ends, whichever way it ends. + /// + /// Runs on a transport thread, so it must not block. + /// `status` is OK for an upload if the client ended normally, and the failure otherwise + /// (the client went away, the transport failed, or the upload was rejected). + /// + /// \param status The terminal status of the upload. + /// \return A future that completes when the finish has been processed. + virtual Future<> OnFinish(Status status) { return Future<>::MakeFinished(); } + + /// \brief The writer for application metadata sent back to the client. + /// + /// Valid for as long as the upload's RPC is live (until OnFinish is reported). + /// Never null, calls after the upload ended report the ending instead of writing. + const std::shared_ptr& writer() const { return writer_; } + + /// \brief Cancel the upload with `status`, from any thread. + /// + /// Finishes the RPC with that status without waiting for the client to end + /// the upload: the client's pending write or Close() reports it. + /// Use for a server-side rejection discovered mid-upload (quota, bad batch, upstream + /// error). + /// Safe to call from any thread at any time: once the upload is no + /// longer in flight this reports Invalid instead of reaching a finished RPC. + /// \param status The status to cancel the upload with. + /// + /// \return OK when the cancel was handed to the transport, or Invalid when + /// there is no upload in flight. + Future<> Cancel(Status status); + + private: + /// The transport installs and clears the state through + /// FlightDataListenerTransport::Install/Clear, both under this lock. + friend class internal::FlightDataListenerTransport; + + /// Whether OnFinish was already reported. + std::atomic finished_{false}; + + /// Never null: the listener may use it without waiting for OnDescriptor. + std::shared_ptr writer_ = + std::make_shared(); + + mutable std::mutex transport_mutex_; + /// Not owned: the transport's own RPC state, valid until Clear(). + internal::FlightDataListenerTransport* transport_ = nullptr; +}; + +/// \brief Push style stream decoder that turns raw arrow Buffers into +/// FlightStreamChunks. +/// +/// This class decodes Apache Arrow Flight data format from arrow::Buffer +/// and fires events on the provided AsyncFlightDataListener. +class ARROW_FLIGHT_EXPORT AsyncFlightMessageDecoder { + public: + /// \brief Construct an AsyncFlightMessageDecoder with the given listener and IPC read + /// options. + /// + /// \param listener The listener that will receive decoded Flight messages. + /// \param options The IPC read options to use for decoding. + explicit AsyncFlightMessageDecoder( + std::shared_ptr listener, + ipc::IpcReadOptions options = ipc::IpcReadOptions::Defaults()); + ~AsyncFlightMessageDecoder(); + + /// \brief Decode one FlightData message directly from a buffer. + /// + /// Fires listener->OnSchemaDecoded() on the first message containing + /// a schema, listener->OnNext() for each subsequent record batch, + /// metadata-only message or dictionary batch. + /// + /// \param[in] buffer a raw buffer directly from the transport. Example + /// the arrow::Buffer extracted from the grpc::ByteBuffer from the gRPC transport. + /// \return Status + Future<> Consume(std::shared_ptr buffer); + + /// \brief Decode one FlightData message already deserialized by the transport. + /// + /// Same as Consume(Buffer) but avoids a serialization round-trip for transports + /// that hand out an internal::FlightData directly (the gRPC client reads FlightData + /// straight into internal::FlightData). + /// \internal This method is intended for internal use within the Arrow library. + Future<> Consume(internal::FlightData data); + + /// \brief The decoded schema. + /// \return The decoded schema, or nullptr if no schema has been decoded yet. + std::shared_ptr schema() const; + + private: + class AsyncFlightMessageDecoderImpl; + std::unique_ptr impl_; +}; + +} // namespace arrow::flight diff --git a/cpp/src/arrow/flight/flight_internals_test.cc b/cpp/src/arrow/flight/flight_internals_test.cc index 5ec9b75dcfda..6061e120724f 100644 --- a/cpp/src/arrow/flight/flight_internals_test.cc +++ b/cpp/src/arrow/flight/flight_internals_test.cc @@ -31,10 +31,14 @@ #include "arrow/flight/client_cookie_middleware.h" #include "arrow/flight/client_middleware.h" #include "arrow/flight/cookie_internal.h" +#include "arrow/flight/flight_data_decoder.h" #include "arrow/flight/serialization_internal.h" #include "arrow/flight/server.h" +#include "arrow/flight/server_async.h" #include "arrow/flight/test_util.h" #include "arrow/flight/transport.h" +#include "arrow/flight/transport/grpc/async_grpc_compat.h" +#include "arrow/flight/transport/grpc/serialization_internal.h" #include "arrow/flight/transport/grpc/util_internal.h" #include "arrow/flight/types.h" #include "arrow/ipc/reader.h" @@ -740,6 +744,43 @@ TEST(GrpcTransport, FlightDataDeserialize) { #endif } +TEST(GrpcTransport, BytesFromBufferConcatenatesSlices) { + // gRPC hands the caller whatever slice layout the message arrived in: the + // loop has to stitch the slices back together, embedded NULs included. + const std::string first("ab\0c", 4); + const std::string empty; + const std::string third("\xff\x01", 2); + std::vector<::grpc::Slice> slices = {::grpc::Slice(first), ::grpc::Slice(empty), + ::grpc::Slice(third)}; + ::grpc::ByteBuffer buffer(slices.data(), slices.size()); + + ASSERT_OK_AND_ASSIGN(auto bytes, + flight::transport::grpc::BytesFromBuffer(buffer, "unused")); + ASSERT_EQ(6, static_cast(bytes.size())); + ASSERT_EQ(first + third, bytes); +} + +TEST(GrpcTransport, BytesFromBufferEmptyBuffer) { + // A valid buffer holding no slices is an empty message, not an error. + ::grpc::ByteBuffer buffer(nullptr, 0); + ASSERT_OK_AND_ASSIGN(auto bytes, + flight::transport::grpc::BytesFromBuffer(buffer, "unused")); + ASSERT_EQ("", bytes); +} + +TEST(GrpcTransport, BytesFromBufferInvalidBufferIsAnInternalError) { + // A default-constructed buffer has no underlying grpc_byte_buffer: the failure + // must keep the caller's message and the transport's Internal detail. + ::grpc::ByteBuffer buffer; + auto result = flight::transport::grpc::BytesFromBuffer(buffer, "no bytes"); + ASSERT_FALSE(result.ok()) << "a buffer with no underlying message cannot be dumped"; + ASSERT_TRUE(result.status().IsIOError()) << result.status(); + ASSERT_NE(std::string::npos, result.status().ToString().find("no bytes")); + auto detail = FlightStatusDetail::UnwrapStatus(result.status()); + ASSERT_NE(nullptr, detail); + ASSERT_EQ(FlightStatusCode::Internal, detail->code()); +} + // ---------------------------------------------------------------------- // Transport-agnostic serialization roundtrip tests @@ -808,6 +849,329 @@ TEST(FlightSerialization, RoundtripMetadataOnly) { ASSERT_EQ(data.app_metadata->ToString(), "metadata-only-message"); } +// ---------------------------------------------------------------------- +// Push-style decoder (AsyncFlightMessageDecoder) tests + +namespace { + +// Serialize one FlightPayload into the single contiguous buffer a transport +// hands to AsyncFlightMessageDecoder::Consume (exactly one FlightData message). +::arrow::Result> DecoderTestWireBuffer( + const FlightPayload& payload) { + ARROW_ASSIGN_OR_RAISE(auto buffers, internal::SerializePayloadToBuffers(payload)); + return ConcatenateBuffers(buffers); +} + +// Captures the events AsyncFlightMessageDecoder fires, in order. +class DecoderTestListener : public AsyncFlightDataListener { + public: + explicit DecoderTestListener(Status next_status = Status::OK(), + Status descriptor_status = Status::OK()) + : next_status_(std::move(next_status)), + descriptor_status_(std::move(descriptor_status)) {} + + Future<> OnDescriptor(const FlightDescriptor& descriptor) override { + events.emplace_back("OnDescriptor"); + descriptors.push_back(descriptor); + return Future<>::MakeFinished(descriptor_status_); + } + + Future<> OnNext(FlightStreamChunk chunk) override { + events.emplace_back("OnNext"); + chunks.push_back(std::move(chunk)); + return Future<>::MakeFinished(next_status_); + } + + // The decoder calls the single-argument overload. + Status OnSchemaDecoded(std::shared_ptr schema) override { + events.emplace_back("OnSchemaDecoded"); + schemas.push_back(std::move(schema)); + return Status::OK(); + } + + std::vector events; + std::vector> schemas; + std::vector chunks; + std::vector descriptors; + + private: + Status next_status_; + Status descriptor_status_; +}; + +} // namespace + +TEST(AsyncFlightMessageDecoder, SchemaAndBatches) { + auto schema = arrow::schema( + {arrow::field("a", arrow::int32()), arrow::field("b", arrow::int32())}); + auto batch1 = RecordBatch::Make( + schema, 2, + {ArrayFromJSON(arrow::int32(), "[1, 2]"), ArrayFromJSON(arrow::int32(), "[3, 4]")}); + auto batch2 = RecordBatch::Make( + schema, 2, + {ArrayFromJSON(arrow::int32(), "[5, 6]"), ArrayFromJSON(arrow::int32(), "[7, 8]")}); + auto reader = RecordBatchReader::Make({batch1, batch2}).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + + auto listener = std::make_shared(); + AsyncFlightMessageDecoder decoder(listener); + EXPECT_EQ(decoder.schema(), nullptr); + + // One wire message per payload: schema first, then one per record batch. + ASSERT_OK_AND_ASSIGN(auto schema_payload, stream.GetSchemaPayload()); + ASSERT_OK_AND_ASSIGN(auto schema_wire, DecoderTestWireBuffer(schema_payload)); + ASSERT_OK(decoder.Consume(schema_wire).status()); + + EXPECT_THAT(listener->events, ::testing::ElementsAre("OnSchemaDecoded")); + ASSERT_EQ(listener->schemas.size(), 1); + EXPECT_TRUE(listener->schemas[0]->Equals(*schema)); + ASSERT_NE(decoder.schema(), nullptr); + EXPECT_TRUE(decoder.schema()->Equals(*schema)); + + ASSERT_OK_AND_ASSIGN(auto payload1, stream.Next()); + ASSERT_OK_AND_ASSIGN(auto payload2, stream.Next()); + ASSERT_OK_AND_ASSIGN(auto wire1, DecoderTestWireBuffer(payload1)); + ASSERT_OK_AND_ASSIGN(auto wire2, DecoderTestWireBuffer(payload2)); + ASSERT_OK(decoder.Consume(wire1).status()); + ASSERT_OK(decoder.Consume(wire2).status()); + + EXPECT_THAT(listener->events, + ::testing::ElementsAre("OnSchemaDecoded", "OnNext", "OnNext")); + ASSERT_EQ(listener->chunks.size(), 2); + ASSERT_NE(listener->chunks[0].data, nullptr); + EXPECT_TRUE(listener->chunks[0].data->Equals(*batch1)); + EXPECT_EQ(listener->chunks[0].app_metadata, nullptr); + ASSERT_NE(listener->chunks[1].data, nullptr); + EXPECT_TRUE(listener->chunks[1].data->Equals(*batch2)); + EXPECT_EQ(listener->chunks[1].app_metadata, nullptr); +} + +TEST(AsyncFlightMessageDecoder, MetadataOnlyChunkIsAValue) { + // A message with app_metadata but no IPC body is a value, not an issue . + FlightPayload payload; + payload.app_metadata = Buffer::FromString("only-app-metadata"); + ASSERT_OK_AND_ASSIGN(auto wire, DecoderTestWireBuffer(payload)); + + auto listener = std::make_shared(); + AsyncFlightMessageDecoder decoder(listener); + ASSERT_OK(decoder.Consume(wire).status()); + + EXPECT_THAT(listener->events, ::testing::ElementsAre("OnNext")); + ASSERT_EQ(listener->chunks.size(), 1); + EXPECT_EQ(listener->chunks[0].data, nullptr); + ASSERT_NE(listener->chunks[0].app_metadata, nullptr); + EXPECT_EQ(listener->chunks[0].app_metadata->ToString(), "only-app-metadata"); + EXPECT_EQ(decoder.schema(), nullptr); +} + +TEST(AsyncFlightMessageDecoder, AppMetadataOnBatch) { + auto schema = arrow::schema({arrow::field("a", arrow::int32())}); + auto batch = RecordBatch::Make(schema, 3, {ArrayFromJSON(arrow::int32(), "[1, 2, 3]")}); + auto reader = RecordBatchReader::Make({batch}).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + + auto listener = std::make_shared(); + AsyncFlightMessageDecoder decoder(listener); + + ASSERT_OK_AND_ASSIGN(auto schema_payload, stream.GetSchemaPayload()); + ASSERT_OK_AND_ASSIGN(auto schema_wire, DecoderTestWireBuffer(schema_payload)); + ASSERT_OK(decoder.Consume(schema_wire).status()); + + ASSERT_OK_AND_ASSIGN(auto payload, stream.Next()); + payload.app_metadata = Buffer::FromString("batch-app-metadata"); + ASSERT_OK_AND_ASSIGN(auto wire, DecoderTestWireBuffer(payload)); + ASSERT_OK(decoder.Consume(wire).status()); + + EXPECT_THAT(listener->events, ::testing::ElementsAre("OnSchemaDecoded", "OnNext")); + ASSERT_EQ(listener->chunks.size(), 1); + ASSERT_NE(listener->chunks[0].data, nullptr); + EXPECT_TRUE(listener->chunks[0].data->Equals(*batch)); + ASSERT_NE(listener->chunks[0].app_metadata, nullptr); + EXPECT_EQ(listener->chunks[0].app_metadata->ToString(), "batch-app-metadata"); +} + +TEST(AsyncFlightMessageDecoder, SchemaOnlyStream) { + auto schema = arrow::schema({arrow::field("a", arrow::int32())}); + auto batch = RecordBatch::Make(schema, 1, {ArrayFromJSON(arrow::int32(), "[42]")}); + auto reader = RecordBatchReader::Make({batch}).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + + auto listener = std::make_shared(); + AsyncFlightMessageDecoder decoder(listener); + + ASSERT_OK_AND_ASSIGN(auto schema_payload, stream.GetSchemaPayload()); + ASSERT_OK_AND_ASSIGN(auto wire, DecoderTestWireBuffer(schema_payload)); + ASSERT_OK(decoder.Consume(wire).status()); + + EXPECT_THAT(listener->events, ::testing::ElementsAre("OnSchemaDecoded")); + EXPECT_TRUE(listener->chunks.empty()); + ASSERT_NE(decoder.schema(), nullptr); + EXPECT_TRUE(decoder.schema()->Equals(*schema)); +} + +TEST(AsyncFlightMessageDecoder, ListenerErrorIsPropagated) { + auto schema = arrow::schema({arrow::field("a", arrow::int32())}); + auto batch = RecordBatch::Make(schema, 1, {ArrayFromJSON(arrow::int32(), "[1]")}); + auto reader = RecordBatchReader::Make({batch}).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + + auto listener = + std::make_shared(Status::Invalid("listener failed")); + AsyncFlightMessageDecoder decoder(listener); + + ASSERT_OK_AND_ASSIGN(auto schema_payload, stream.GetSchemaPayload()); + ASSERT_OK_AND_ASSIGN(auto schema_wire, DecoderTestWireBuffer(schema_payload)); + ASSERT_OK(decoder.Consume(schema_wire).status()); + + // A listener error on a record batch must reach the caller of Consume. + ASSERT_OK_AND_ASSIGN(auto payload, stream.Next()); + ASSERT_OK_AND_ASSIGN(auto wire, DecoderTestWireBuffer(payload)); + auto status = decoder.Consume(wire).status(); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), StatusCode::Invalid); + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("listener failed")); + + // and so must one on a metadata-only message. + FlightPayload metadata_only; + metadata_only.app_metadata = Buffer::FromString("meta"); + ASSERT_OK_AND_ASSIGN(auto meta_wire, DecoderTestWireBuffer(metadata_only)); + auto meta_status = decoder.Consume(meta_wire).status(); + EXPECT_FALSE(meta_status.ok()); + EXPECT_EQ(meta_status.code(), StatusCode::Invalid); +} + +TEST(AsyncFlightMessageDecoder, FailedOpenLeavesTheDecoderUsable) { + // A first IPC message that is not a schema makes RecordBatchStreamReader::Open + // fail. The decoder must survive that: it owns its message reader, rather + // than handing it over to Open, which destroys it on failure and would leave + // the decoder's cached pointer dangling for the next Consume(). + auto schema = arrow::schema({arrow::field("a", arrow::int32())}); + auto batch = RecordBatch::Make(schema, 1, {ArrayFromJSON(arrow::int32(), "[1]")}); + auto reader = RecordBatchReader::Make({batch}).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + + auto listener = std::make_shared(); + AsyncFlightMessageDecoder decoder(listener); + + // A record batch as the first message: not a schema, so Open() fails. + ASSERT_OK_AND_ASSIGN(auto batch_payload, stream.Next()); + ASSERT_OK_AND_ASSIGN(auto batch_wire, DecoderTestWireBuffer(batch_payload)); + auto failed = decoder.Consume(batch_wire).status(); + EXPECT_FALSE(failed.ok()) << "a first message that is not a schema must be rejected"; + EXPECT_TRUE(listener->events.empty()) + << "nothing may be surfaced for a rejected stream"; + + // The decoder must still work: consume a well-formed stream and watch the + // schema and the batch come through. + auto good_reader = RecordBatchReader::Make({batch}).ValueOrDie(); + RecordBatchStream good_stream(std::move(good_reader)); + ASSERT_OK_AND_ASSIGN(auto schema_payload, good_stream.GetSchemaPayload()); + ASSERT_OK_AND_ASSIGN(auto schema_wire, DecoderTestWireBuffer(schema_payload)); + ASSERT_OK(decoder.Consume(schema_wire).status()); + ASSERT_OK_AND_ASSIGN(auto good_batch, good_stream.Next()); + ASSERT_OK_AND_ASSIGN(auto good_wire, DecoderTestWireBuffer(good_batch)); + ASSERT_OK(decoder.Consume(good_wire).status()); + + EXPECT_THAT(listener->events, ::testing::ElementsAre("OnSchemaDecoded", "OnNext")); + ASSERT_EQ(listener->chunks.size(), 1); + ASSERT_NE(listener->chunks[0].data, nullptr); + EXPECT_TRUE(listener->chunks[0].data->Equals(*batch)); +} + +TEST(AsyncFlightMessageDecoder, DictionaryStream) { + // A dictionary-encoded stream arrives as: schema -> dictionary message -> batch. + // Tthe IPC reader reads through the dictionary message while + // looking for a record batch, so a pushed message must only be handed out once + // (returning it twice re-reads it forever). + const auto dict_type = dictionary(arrow::int32(), utf8()); + const auto schema = arrow::schema({arrow::field("d", dict_type)}); + const auto batch = + RecordBatch::Make(schema, 3, {ArrayFromJSON(dict_type, R"(["a", "b", "a"])")}); + auto reader = RecordBatchReader::Make({batch}, schema).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + + auto listener = std::make_shared(); + AsyncFlightMessageDecoder decoder(listener); + + ASSERT_OK_AND_ASSIGN(auto schema_payload, stream.GetSchemaPayload()); + ASSERT_OK_AND_ASSIGN(auto schema_wire, DecoderTestWireBuffer(schema_payload)); + ASSERT_OK(decoder.Consume(schema_wire).status()); + + EXPECT_THAT(listener->events, ::testing::ElementsAre("OnSchemaDecoded")); + ASSERT_NE(decoder.schema(), nullptr); + EXPECT_TRUE(decoder.schema()->Equals(*schema)); + + // The dictionary message is consumed without surfacing as a chunk of its own. + int payloads = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(auto payload, stream.Next()); + if (payload.ipc_message.metadata == nullptr) { + break; // end of stream + } + ASSERT_OK_AND_ASSIGN(auto wire, DecoderTestWireBuffer(payload)); + ASSERT_OK(decoder.Consume(wire).status()); + ++payloads; + } + EXPECT_EQ(payloads, 2); // dictionary message + record batch + + // and the batch arrives exactly once, decoded against the dictionary. + EXPECT_THAT(listener->events, ::testing::ElementsAre("OnSchemaDecoded", "OnNext")); + ASSERT_EQ(listener->chunks.size(), 1); + ASSERT_NE(listener->chunks[0].data, nullptr); + EXPECT_TRUE(listener->chunks[0].data->Equals(*batch)); + EXPECT_EQ(listener->chunks[0].app_metadata, nullptr); +} + +TEST(AsyncFlightMessageDecoder, DescriptorReachesListener) { + // An upload message carries the descriptor of the DoPut command, the + // listener sees it before anything else from that message is decoded. + const auto schema = arrow::schema({arrow::field("a", arrow::int32())}); + const auto batch = + RecordBatch::Make(schema, 1, {ArrayFromJSON(arrow::int32(), "[42]")}); + auto reader = RecordBatchReader::Make({batch}).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + + auto listener = std::make_shared(); + AsyncFlightMessageDecoder decoder(listener); + + auto descriptor = FlightDescriptor::Path({"wave", "h1"}); + ASSERT_OK_AND_ASSIGN(auto payload, stream.GetSchemaPayload()); + ASSERT_OK(internal::ToPayload(descriptor, &payload.descriptor)); + ASSERT_OK_AND_ASSIGN(auto wire, DecoderTestWireBuffer(payload)); + ASSERT_OK(decoder.Consume(wire).status()); + + EXPECT_THAT(listener->events, + ::testing::ElementsAre("OnDescriptor", "OnSchemaDecoded")); + ASSERT_EQ(listener->descriptors.size(), 1); + EXPECT_EQ(listener->descriptors[0], descriptor); +} + +TEST(AsyncFlightMessageDecoder, DescriptorErrorRejectsUpload) { + const auto schema = arrow::schema({arrow::field("a", arrow::int32())}); + const auto batch = + RecordBatch::Make(schema, 1, {ArrayFromJSON(arrow::int32(), "[42]")}); + auto reader = RecordBatchReader::Make({batch}).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + + auto listener = std::make_shared( + Status::OK(), Status::Invalid("descriptor rejected")); + AsyncFlightMessageDecoder decoder(listener); + + auto descriptor = FlightDescriptor::Path({"wave", "h1"}); + ASSERT_OK_AND_ASSIGN(auto payload, stream.GetSchemaPayload()); + ASSERT_OK(internal::ToPayload(descriptor, &payload.descriptor)); + ASSERT_OK_AND_ASSIGN(auto wire, DecoderTestWireBuffer(payload)); + auto status = decoder.Consume(wire).status(); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), StatusCode::Invalid); + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("descriptor rejected")); + + // The upload is rejected before the schema is decoded. + EXPECT_THAT(listener->events, ::testing::ElementsAre("OnDescriptor")); + EXPECT_EQ(decoder.schema(), nullptr); +} + // ---------------------------------------------------------------------- // Transport abstraction tests @@ -922,5 +1286,26 @@ TEST(TestServerSignalState, ShutdownUninitializedKeepsOtherServerSignalable) { #endif +// An async server must report why it cannot start, in every build: either it +// starts, or Init() answers NotImplemented naming the gRPC requirement. Before +// this was pinned, a build against gRPC < 1.65 never registered an async server +// factory and Init() failed with the generic "no async server transport +// implementation" KeyError, hiding the real cause. +TEST(TestAsyncServerAvailability, InitReportsWhetherAsyncServerIsSupported) { + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + AsyncGenericFlightServerBase server; + const auto status = server.Init(options); +#if ARROW_FLIGHT_HAS_ASYNC_SERVER + ASSERT_OK(status); + ASSERT_OK(server.Shutdown()); +#else + ASSERT_FALSE(status.ok()); + ASSERT_EQ(StatusCode::NotImplemented, status.code()); + ASSERT_NE(std::string::npos, status.message().find("1.65")) + << "the error must name the gRPC version requirement, got: " << status.message(); +#endif +} + } // namespace flight } // namespace arrow diff --git a/cpp/src/arrow/flight/meson.build b/cpp/src/arrow/flight/meson.build index dcdacc42ffdc..465aabad2fa2 100644 --- a/cpp/src/arrow/flight/meson.build +++ b/cpp/src/arrow/flight/meson.build @@ -23,11 +23,13 @@ install_headers( 'client.h', 'client_middleware.h', 'client_tracing_middleware.h', + 'flight_data_decoder.h', 'middleware.h', 'otel_logging.h', 'platform.h', 'server_auth.h', 'server.h', + 'server_async.h', 'server_middleware.h', 'server_tracing_middleware.h', 'test_auth_handlers.h', @@ -45,6 +47,16 @@ install_headers( ) grpc_dep = dependency('grpc++') + +# GRPC_ENABLE_ASYNC gates the async *client* paths in grpc_client.cc; 1.40 is +# where those gRPC APIs landed. +has_async_client = grpc_dep.version().version_compare('>=1.40.0') +grpc_enable_async_args = has_async_client ? ['-DGRPC_ENABLE_ASYNC'] : [] + +# The async generic *server* needs gRPC's generic callback API, which exists +# from 1.65 on; see transport/grpc/async_grpc_compat.h, which does the C++ +# version check. Higher floor than the async client above. +has_async_server = grpc_dep.version().version_compare('>=1.65.0') abseil_sync_dep = dependency('absl_synchronization') fs = import('fs') @@ -81,9 +93,11 @@ arrow_flight_srcs = [ 'client_cookie_middleware.cc', 'client_tracing_middleware.cc', 'cookie_internal.cc', + 'flight_data_decoder.cc', 'middleware.cc', 'serialization_internal.cc', 'server.cc', + 'server_async.cc', 'server_auth.cc', 'server_tracing_middleware.cc', 'transport.cc', @@ -98,6 +112,24 @@ arrow_flight_srcs = [ 'types.cc', ] +# The async generic server needs gRPC's generic callback API; see +# async_grpc_compat.h and the CMake equivalent. +if has_async_server + arrow_flight_srcs += [ + 'transport/grpc/async_grpc_service.cc', + 'transport/grpc/async_grpc_do_get.cc', + 'transport/grpc/async_grpc_do_put.cc', + 'transport/grpc/async_grpc_exchange.cc', + 'transport/grpc/async_grpc_exchange_reader.cc', + 'transport/grpc/async_grpc_exchange_writer.cc', + 'transport/grpc/async_grpc_handshake.cc', + 'transport/grpc/async_grpc_streaming.cc', + 'transport/grpc/async_grpc_unary.cc', + ] +endif + + + thread_dep = dependency('threads') arrow_flight = library( @@ -117,8 +149,8 @@ arrow_flight = library( abseil_sync_dep, thread_dep, ], - cpp_shared_args: ['-DARROW_FLIGHT_EXPORTING'], - cpp_static_args: ['-DARROW_FLIGHT_STATIC'], + cpp_shared_args: ['-DARROW_FLIGHT_EXPORTING'] + grpc_enable_async_args, + cpp_static_args: ['-DARROW_FLIGHT_STATIC'] + grpc_enable_async_args, gnu_symbol_visibility: 'inlineshidden', ) @@ -166,6 +198,29 @@ foreach flight_test : flight_tests test(test_name, exc) endforeach +if has_async_server + async_grpc_test = executable( + 'flight-async-grpc-test', + sources: [ + 'transport/grpc/async_grpc_service_test.cc', + 'transport/grpc/async_grpc_do_get_test.cc', + 'transport/grpc/async_grpc_do_put_test.cc', + 'transport/grpc/async_grpc_unary_test.cc', + 'transport/grpc/async_grpc_streaming_test.cc', + 'transport/grpc/async_grpc_exchange_test.cc', + 'transport/grpc/async_grpc_handshake_test.cc', + # these transitively include Flight.grpc.pb.h via the fixtures + # header, so we must declare it here to avoid a race condition + flight_proto_grpc_files[1], + ], + cpp_args: grpc_enable_async_args, + dependencies: [arrow_test_dep, arrow_flight_test_dep], + ) + # The timeout bounds the tests that would otherwise hang the suite instead + # of failing when a cancel/finish path regresses. + test('flight-async-grpc-test', async_grpc_test, timeout: 300) +endif + flight_test_dep_no_main = [ arrow_dep, arrow_flight_test_dep, diff --git a/cpp/src/arrow/flight/serialization_internal.h b/cpp/src/arrow/flight/serialization_internal.h index 896ca5b67567..b4f6a3f0e69b 100644 --- a/cpp/src/arrow/flight/serialization_internal.h +++ b/cpp/src/arrow/flight/serialization_internal.h @@ -180,6 +180,7 @@ ARROW_FLIGHT_EXPORT Status ToProto(const CloseSessionRequest& request, ARROW_FLIGHT_EXPORT Status ToProto(const CloseSessionResult& result, pb::CloseSessionResult* pb_result); +ARROW_FLIGHT_EXPORT Status ToPayload(const FlightDescriptor& descr, std::shared_ptr* out); /// \brief Serialize a FlightPayload to a vector of buffers. diff --git a/cpp/src/arrow/flight/server.cc b/cpp/src/arrow/flight/server.cc index da66d704c56c..1bb50ec191aa 100644 --- a/cpp/src/arrow/flight/server.cc +++ b/cpp/src/arrow/flight/server.cc @@ -43,8 +43,7 @@ namespace flight { /// Server implementation. Manages the lifecycle of the "real" server /// (ServerTransport) and contains struct FlightServerBase::Impl { - std::unique_ptr transport_; - internal::ServerSignalState signal_state_; + internal::ServerLifecycle lifecycle{"FlightServerBase"}; }; FlightServerOptions::FlightServerOptions(const Location& location_) @@ -67,63 +66,29 @@ Status FlightServerBase::Init(const FlightServerOptions& options) { flight::transport::grpc::InitializeFlightGrpcServer(); const auto scheme = options.location.scheme(); - ARROW_ASSIGN_OR_RAISE(impl_->transport_, + ARROW_ASSIGN_OR_RAISE(auto transport, internal::GetDefaultTransportRegistry()->MakeServer( scheme, this, options.memory_manager)); - ARROW_ASSIGN_OR_RAISE(auto uri, options.location.uri()); - return impl_->transport_->Init(options, uri); + return impl_->lifecycle.Init(options, std::move(transport)); } -int FlightServerBase::port() const { return location().port(); } +int FlightServerBase::port() const { return impl_->lifecycle.port(); } -Location FlightServerBase::location() const { return impl_->transport_->location(); } +Location FlightServerBase::location() const { return impl_->lifecycle.location(); } Status FlightServerBase::SetShutdownOnSignals(const std::vector sigs) { - return impl_->signal_state_.SetShutdownOnSignals(sigs); + return impl_->lifecycle.SetShutdownOnSignals(sigs); } -Status FlightServerBase::Serve() { - // Check before entering the signal helper: a failing wait() callback would - // otherwise leave the process signal handlers overridden, and clearing the - // signal target on the error path would disable signals for another server. - if (!impl_->transport_) { - return Status::UnknownError("Server did not start properly"); - } - return impl_->signal_state_.Serve( - [this]() -> Status { return impl_->transport_->Wait(); }, - [this](const std::chrono::system_clock::time_point* deadline) -> Status { - if (deadline) { - return impl_->transport_->Shutdown(*deadline); - } - return impl_->transport_->Shutdown(); - }, - "Error shutting down server"); -} +Status FlightServerBase::Serve() { return impl_->lifecycle.Serve(); } -int FlightServerBase::GotSignal() const { return impl_->signal_state_.GotSignal(); } +int FlightServerBase::GotSignal() const { return impl_->lifecycle.GotSignal(); } Status FlightServerBase::Shutdown(const std::chrono::system_clock::time_point* deadline) { - if (!impl_->transport_) { - return Status::Invalid("Shutdown() on uninitialized FlightServerBase"); - } - return impl_->signal_state_.Shutdown( - [this](const std::chrono::system_clock::time_point* maybe_deadline) -> Status { - if (maybe_deadline) { - return impl_->transport_->Shutdown(*maybe_deadline); - } - return impl_->transport_->Shutdown(); - }, - deadline); + return impl_->lifecycle.Shutdown(deadline); } -Status FlightServerBase::Wait() { - return impl_->signal_state_.Wait([this] { - if (!impl_->transport_) { - return Status::Invalid("Wait() on uninitialized FlightServerBase"); - } - return impl_->transport_->Wait(); - }); -} +Status FlightServerBase::Wait() { return impl_->lifecycle.Wait(); } Status FlightServerBase::ListFlights(const ServerCallContext& context, const Criteria* criteria, diff --git a/cpp/src/arrow/flight/server_async.cc b/cpp/src/arrow/flight/server_async.cc new file mode 100644 index 000000000000..19529c79479a --- /dev/null +++ b/cpp/src/arrow/flight/server_async.cc @@ -0,0 +1,156 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Standalone async Flight server implementation on top of the transport +// interface. + +#include "arrow/flight/server_async.h" + +#include +#include +#include +#include +#include + +#include "arrow/flight/transport.h" +#include "arrow/flight/transport/grpc/grpc_server.h" +#include "arrow/flight/transport_server.h" +#include "arrow/flight/transport_server_internal.h" +#include "arrow/flight/types.h" + +namespace arrow::flight { + +AsyncFlightDataStream::~AsyncFlightDataStream() = default; +AsyncFlightListing::~AsyncFlightListing() = default; +AsyncResultStream::~AsyncResultStream() = default; +AsyncFlightMessageReader::~AsyncFlightMessageReader() = default; +AsyncFlightMessageWriter::~AsyncFlightMessageWriter() = default; + +struct AsyncGenericFlightServerBase::Impl { + internal::ServerLifecycle lifecycle{"AsyncGenericFlightServerBase"}; +}; + +AsyncGenericFlightServerBase::AsyncGenericFlightServerBase() : impl_(new Impl) {} + +AsyncGenericFlightServerBase::~AsyncGenericFlightServerBase() = default; + +Status AsyncGenericFlightServerBase::Init(const FlightServerOptions& options) { + flight::transport::grpc::InitializeFlightGrpcServer(); + + const auto scheme = options.location.scheme(); + ARROW_ASSIGN_OR_RAISE(auto transport, + internal::GetDefaultTransportRegistry()->MakeServerAsync( + scheme, this, options.memory_manager)); + return impl_->lifecycle.Init(options, std::move(transport)); +} + +int AsyncGenericFlightServerBase::port() const { return impl_->lifecycle.port(); } + +Location AsyncGenericFlightServerBase::location() const { + return impl_->lifecycle.location(); +} + +Status AsyncGenericFlightServerBase::SetShutdownOnSignals( + const std::vector signals) { + return impl_->lifecycle.SetShutdownOnSignals(signals); +} + +Status AsyncGenericFlightServerBase::Serve() { return impl_->lifecycle.Serve(); } + +int AsyncGenericFlightServerBase::GotSignal() const { + return impl_->lifecycle.GotSignal(); +} + +Status AsyncGenericFlightServerBase::Shutdown( + const std::chrono::system_clock::time_point* deadline) { + return impl_->lifecycle.Shutdown(deadline); +} + +Status AsyncGenericFlightServerBase::Wait() { return impl_->lifecycle.Wait(); } + +// --- handler defaults: the same answers FlightServerBase's defaults give --- + +arrow::Future> +AsyncGenericFlightServerBase::ListFlightsAsync(const ServerCallContext&, + const Criteria*) { + return arrow::Future>::MakeFinished( + Status::NotImplemented("NYI")); +} + +arrow::Future> +AsyncGenericFlightServerBase::GetFlightInfoAsync(const ServerCallContext&, + const FlightDescriptor&) { + return arrow::Future>::MakeFinished( + Status::NotImplemented("NYI")); +} + +arrow::Future> +AsyncGenericFlightServerBase::PollFlightInfoAsync(const ServerCallContext&, + const FlightDescriptor&) { + return arrow::Future>::MakeFinished( + Status::NotImplemented("NYI")); +} + +arrow::Future> AsyncGenericFlightServerBase::GetSchemaAsync( + const ServerCallContext&, const FlightDescriptor&) { + return arrow::Future>::MakeFinished( + Status::NotImplemented("NYI")); +} + +arrow::Future> +AsyncGenericFlightServerBase::DoActionAsync(const ServerCallContext&, const Action&) { + return arrow::Future>::MakeFinished( + Status::NotImplemented("NYI")); +} + +arrow::Future> AsyncGenericFlightServerBase::ListActionsAsync( + const ServerCallContext&) { + return arrow::Future>::MakeFinished( + Status::NotImplemented("NYI")); +} + +arrow::Future<> AsyncGenericFlightServerBase::DoExchangeAsync( + const ServerCallContext&, std::shared_ptr, + std::shared_ptr) { + return arrow::Future<>::MakeFinished(Status::NotImplemented("NYI")); +} + +arrow::Future> +AsyncGenericFlightServerBase::DoGetAsync(const ServerCallContext&, const Ticket&) { + return arrow::Future>::MakeFinished( + Status::NotImplemented("DoGetAsync is not implemented")); +} + +std::shared_ptr +AsyncGenericFlightServerBase::CreateDoPutListener(const ServerCallContext&) { + // The default refuses uploads, a server that accepts them overrides this. + return nullptr; +} + +Status AsyncGenericFlightServerBase::Handshake(const ServerCallContext&, + const std::string&, std::string*) { + return Status::NotImplemented( + "This service does not have an authentication mechanism enabled."); +} + +Status AsyncGenericFlightServerBase::ValidateToken(const ServerCallContext&, + const std::string&, std::string*) { + // Keep the transport-level (TLS) identity, ignore the token. + return Status::OK(); +} + +} // namespace arrow::flight diff --git a/cpp/src/arrow/flight/server_async.h b/cpp/src/arrow/flight/server_async.h new file mode 100644 index 000000000000..84532d6103d6 --- /dev/null +++ b/cpp/src/arrow/flight/server_async.h @@ -0,0 +1,381 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include + +#include "arrow/flight/flight_data_decoder.h" +#include "arrow/flight/server.h" +#include "arrow/flight/types.h" +#include "arrow/util/future.h" +#include "arrow/util/macros.h" + +namespace arrow::flight { + +/// \brief Interface that produces a sequence of IPC payloads to be sent in +/// FlightData protobuf messages, where a payload may arrive later. +class ARROW_FLIGHT_EXPORT AsyncFlightDataStream { + public: + virtual ~AsyncFlightDataStream(); + + /// \brief The schema of the payloads this stream produces. + /// + /// The transport never calls this: the schema the client sees is the one + /// GetSchemaPayloadAsync() returns. + /// \return The schema of the payloads this stream produces. + virtual std::shared_ptr schema() = 0; + + /// \brief Like GetSchemaPayload(), but the payload arrives later. + /// + /// Called once, before any NextAsync() payload is requested. + /// \return A future completed with the schema payload. A non-OK status fails the RPC. + virtual arrow::Future GetSchemaPayloadAsync() = 0; + + /// \brief Like Next(), but the payload arrives later. + /// + /// The transport keeps one call in flight at a time and only asks again once + /// the previous payload has been written. + /// + /// \return A future completed with the next payload, or with a nullopt once the stream + /// is exhausted. A non-OK status fails the RPC. + virtual arrow::Future> NextAsync() = 0; + + /// \brief Stop the stream: the RPC it was producing for is over. + /// + /// Called at most once, possibly from any thread. A pending + /// NextAsync()/GetSchemaPayloadAsync() future must be resolved here (or + /// abandoned knowingly): the transport does not wait for it. + /// Its return status becomes the RPC's final status only on the normal end-of-stream + /// path (after NextAsync() reports end of stream). + /// On every other path (write failure, cancellation) the RPC already has a status and + /// the returned status is discarded. + /// \return The status of closing the stream. + virtual Status Close() { return Status::OK(); } +}; + +/// \brief An asynchronous iterator of FlightInfo instances, returned by +/// ListFlightsAsync. +/// +/// The transport pulls it one NextAsync() at a time and writes each +/// FlightInfo to the client as it becomes available. +class ARROW_FLIGHT_EXPORT AsyncFlightListing { + public: + virtual ~AsyncFlightListing(); + + /// \brief The next FlightInfo in the listing. + /// \return A future completed with the next FlightInfo, or with a nullptr + /// once the listing is exhausted. A non-OK status fails the RPC. + virtual arrow::Future> NextAsync() = 0; +}; + +/// \brief An asynchronous iterator of Result instances, returned by +/// DoActionAsync. +class ARROW_FLIGHT_EXPORT AsyncResultStream { + public: + virtual ~AsyncResultStream(); + + /// \brief The next Result in the stream. + /// \return A future completed with the next Result, or with a nullptr once + /// the stream is exhausted. A non-OK status fails the RPC. + virtual arrow::Future> NextAsync() = 0; +}; + +/// \brief The reader of one DoExchange: the messages the client sends, pulled +/// one NextAsync() at a time. +/// +/// Runs on the thread that calls it, at most one NextAsync() may be +/// outstanding. +/// Schema and dictionary messages are read through (they are not chunks), a metadata-only +/// message is a chunk with a null data member, and the chunk whose members are both null +/// ends the exchange. +/// The reader must not be used after the future returned by DoExchangeAsync completes. +class ARROW_FLIGHT_EXPORT AsyncFlightMessageReader { + public: + virtual ~AsyncFlightMessageReader(); + + /// \brief The descriptor of this exchange. + virtual const FlightDescriptor& descriptor() const = 0; + + /// \brief The next message of the exchange, decoded. + /// \return A future completed with the next chunk, both of its members are + /// null at the end of the exchange. A non-OK status fails the RPC. + virtual arrow::Future NextAsync() = 0; +}; + +/// \brief The writer of one DoExchange: the messages the server sends back. +/// +/// Runs on the thread that calls it, await each returned future before +/// writing again. The writer must not be used after the future returned by +/// DoExchangeAsync completes. +class ARROW_FLIGHT_EXPORT AsyncFlightMessageWriter { + public: + virtual ~AsyncFlightMessageWriter(); + + /// \brief Start the writer with the schema of the messages it will write. + /// Must be called before any record batch is written. + /// \param schema The schema of the messages to be written. + /// \return A future that completes once the writer is ready to accept messages. + virtual arrow::Future<> BeginAsync(std::shared_ptr schema) = 0; + + /// \brief Write one record batch. + /// \param batch The record batch to write. + /// \return A future that completes once the batch has been written. + virtual arrow::Future<> WriteRecordBatchAsync(const RecordBatch& batch) = 0; + + /// \brief Write one application-metadata-only message. + /// \param app_metadata The application metadata to write. + /// \return A future that completes once the metadata has been written. + virtual arrow::Future<> WriteMetadataAsync(std::shared_ptr app_metadata) = 0; + + /// \brief Write one record batch with application metadata attached. + /// \param batch The record batch to write. + /// \param app_metadata The application metadata to attach. + /// \return A future that completes once the batch and metadata have been written. + virtual arrow::Future<> WriteWithMetadataAsync( + const RecordBatch& batch, std::shared_ptr app_metadata) = 0; + + /// \brief Finish the write side: end the IPC stream. + /// + /// Call once, after the last message, then await it before completing + /// DoExchangeAsync's future. + /// \return A future that completes once the write side has been closed. + virtual arrow::Future<> CloseAsync() = 0; +}; + +/// \brief a Flight server served by the async generic gRPC +/// callback service. +/// +/// Derive from this class alone and implement the handlers you serve: +/// +/// \code +/// class MyServer : public arrow::flight::AsyncGenericFlightServerBase { +/// arrow::Future> GetFlightInfoAsync( +/// const arrow::flight::ServerCallContext&, +/// const arrow::flight::FlightDescriptor&) override; +/// +/// arrow::Future> DoGetAsync( +/// const arrow::flight::ServerCallContext&, +/// const arrow::flight::Ticket&) override; +/// }; +/// \endcode +/// +/// Every handler returns a future: complete it from your own thread or pool +/// when the answer is ready, and no gRPC callback thread is ever held waiting +/// for one. Handlers that do run inline should still return promptly: blocking +/// a callback thread starves new connections. +/// +/// All handlers must be non-blocking and return promptly. +/// Non overriding handlers will return `Status::NotImplemented` by default. +/// +/// Every request argument a handler receives (`Ticket`, `FlightDescriptor`, +/// `Criteria`, `Action`) is only valid for the duration of that handler call: +/// copy what you need before completing the returned future from another +/// thread. The `ServerCallContext` handed to a handler is likewise the call's +/// own context and must not outlive it. +/// +/// Authentication goes through Handshake()/ValidateToken(): a blocking +/// ServerAuthHandler in FlightServerOptions is not supported, and Init() +/// rejects it. +class ARROW_FLIGHT_EXPORT AsyncGenericFlightServerBase { + public: + AsyncGenericFlightServerBase(); + virtual ~AsyncGenericFlightServerBase(); + + /// \brief Initialize the server listening at the given location. + /// Must be called before any other method. + /// \param options The server options to use for initialization. + /// \return Status indicating success or failure of the initialization. + 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 non-positive + /// value if no port exists (e.g. when listening on a domain socket). + /// \return The port number the server is listening on, or a non-positive value. + int port() const; + + /// \return The address the server is listening on. + /// \warning Only valid after Init() has been called. + Location location() const; + + /// \brief Set the server to stop on any of the given signal numbers. + /// \warning Must be called before Serve(). + /// \param signals The signal numbers on which the server should shut down. + /// \return Status indicating success or failure of setting the shutdown signals. + Status SetShutdownOnSignals(const std::vector signals); + + /// \brief Start serving. Blocks until the server shuts down, which happens + /// when Shutdown() is called or a registered signal arrives. + /// \return Status indicating success or failure of starting the server. + Status Serve(); + + /// \return The signal number that interrupted Serve(), or 0. + /// \note Only valid after Serve() has returned. + int GotSignal() const; + + /// \brief Shut down the server, blocking until in-flight RPCs finish. + /// May be called from another thread or a signal handler while Serve() + /// blocks, with an optional deadline. + /// + /// An RPC whose handler never completes its future is not finished, so this + /// blocks until that future resolves: a handler that abandons its future + /// (see the per-handler notes on the transport not waiting for it) needs a + /// deadline here, or the future must be given up on before shutdown. + /// \param[in] deadline is an optional time point by which the server should complete + /// shutdown. + /// \return Status indicating success or failure of the shutdown operation. + /// \warning Call at most once. + Status Shutdown(const std::chrono::system_clock::time_point* deadline = NULLPTR); + + /// \brief Block until the server shuts down. Does not respond to signals. + /// + /// Like Shutdown(), this does not return while a handler future on a live RPC + /// is still unresolved. + Status Wait(); + + /// \brief Asynchronously retrieve a list of available fields given an + /// optional opaque criteria. + /// \param[in] context is the server call context. + /// \param[in] criteria is the optional opaque criteria for listing flights. + /// \see FlightServerBase::ListFlights + /// + /// \return A future completed with the listing, or with a non-OK status to + /// fail the RPC. + /// A future completed with a nullptr listing serves an empty listing. + /// The transport pulls the listing one NextAsync() at a time, the criteria pointer is + /// never null and is only valid for the duration of the call. + virtual arrow::Future> ListFlightsAsync( + const ServerCallContext& context, const Criteria* criteria); + + /// \brief Asynchronously retrieve the schema and an access plan for the + /// indicated descriptor. + /// \see FlightServerBase::GetFlightInfo + /// \param[in] context is the server call context. + /// \param[in] request is the flight descriptor for which to retrieve the info. + /// + /// \return A future completed with the info, or with a non-OK status to + /// fail the RPC. A `nullptr` info answers `KeyError("Flight not found")`, as + /// a synchronous server that returns OK without setting the response does. + virtual arrow::Future> GetFlightInfoAsync( + const ServerCallContext& context, const FlightDescriptor& request); + + /// \brief Asynchronously retrieve the current status of the target query. + /// \param[in] context is the server call context. + /// \param[in] request is the flight descriptor for which to retrieve the status. + /// \see FlightServerBase::PollFlightInfo. + /// \return A future completed with the info, or with a non-OK status to + /// fail the RPC. A nullptr info answers `KeyError("Flight not found")`. + virtual arrow::Future> PollFlightInfoAsync( + const ServerCallContext& context, const FlightDescriptor& request); + + /// \brief Asynchronously retrieve the schema for the indicated descriptor. + /// \param[in] context is the server call context. + /// \param[in] request is the flight descriptor for which to retrieve the schema. + /// \see FlightServerBase::GetSchema. + /// \return A future completed with the schema, or with a non-OK status to + /// fail the RPC. + /// A `nullptr` schema answers `KeyError("Flight not found")`. + virtual arrow::Future> GetSchemaAsync( + const ServerCallContext& context, const FlightDescriptor& request); + + /// \brief Asynchronously execute an action, returning a stream of zero or + /// more results. + /// \param[in] context is the server call context. + /// \param[in] action is the action to execute. + /// \see FlightServerBase::DoAction. + /// \return A future completed with the result stream, or with a non-OK + /// status to fail the RPC. A future completed with a `nullptr` stream + /// answers CANCELLED, as a synchronous server that returns OK without a + /// stream does. + /// The transport pulls the stream one NextAsync() at a time. + virtual arrow::Future> DoActionAsync( + const ServerCallContext& context, const Action& action); + + /// \brief Asynchronously retrieve the list of available actions. + /// \param[in] context is the server call context. + /// \see FlightServerBase::ListActions. + virtual arrow::Future> ListActionsAsync( + const ServerCallContext& context); + + /// \brief Asynchronously process a bidirectional stream of IPC payloads. + /// \see FlightServerBase::DoExchange. + /// \param[in] context is the server call context. + /// \param[in] reader is the asynchronous reader for the client's messages. + /// \param[in] writer is the asynchronous writer for the server's messages. + /// \return A future completed when the exchange is over, its status is the + /// RPC's status. + /// While it runs, consume the client's messages through the reader and send the + /// server's through the writer. Both must not be used after the future completes. + virtual arrow::Future<> DoExchangeAsync( + const ServerCallContext& context, std::shared_ptr reader, + std::shared_ptr writer); + + /// \brief The stream serving one DoGet RPC. + /// + /// Called after the ticket is parsed, the request is only valid for the + /// duration of the call. + /// \param[in] context is the server call context. + /// \param[in] request is an opaque ticket. + /// \return A future completed with the stream to serve, or with a non-OK + /// status to fail the RPC. + /// A future completed with a `nullptr` stream answers `KeyError("No data in this + /// flight")`. Complete the future from your ownthread or pool: the transport writes + /// whatever stream it resolves to, and must not be blocked waiting for it. + /// + /// If the RPC is cancelled while the returned future is pending, the + /// transport finishes the RPC without waiting for it, and the stream it + /// eventually resolves to is dropped without Close() ever being delivered + /// (the transport no longer holds the stream). + virtual arrow::Future> DoGetAsync( + const ServerCallContext& context, const Ticket& request); + + /// \brief Create the AsyncFlightDataListener serving one DoPut RPC. + /// \see AsyncFlightDataListener + /// \param[in] context is the server call context. + /// \return A shared pointer to the AsyncFlightDataListener handling the upload, or + /// `nullptr` to refuse the upload. + virtual std::shared_ptr CreateDoPutListener( + const ServerCallContext& context); + + /// \brief Handle the handshake protocol with the client. + /// \param[in] context is the server call context. + /// \param[in] request is the client's handshake message. + /// \param[out] response is where the server writes its handshake response. + /// \return Status indicating success or failure of the handshake. + virtual Status Handshake(const ServerCallContext& context, const std::string& request, + std::string* response); + + /// \brief Validate the token sent in the `auth-token-bin` header of RPCs + /// issued after a successful Handshake(). + /// \param[in] context is the server call context. + /// \param[in] token is the authentication token sent by the client. + /// \param[out] peer_identity is set to the authenticated identity on success. + /// \return Status indicating success or failure of the token validation. + virtual Status ValidateToken(const ServerCallContext& context, const std::string& token, + std::string* peer_identity); + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace arrow::flight diff --git a/cpp/src/arrow/flight/transport.cc b/cpp/src/arrow/flight/transport.cc index fd74b6d95a59..b14b2735bd4c 100644 --- a/cpp/src/arrow/flight/transport.cc +++ b/cpp/src/arrow/flight/transport.cc @@ -136,6 +136,15 @@ class TransportRegistry::Impl final { } return it->second(base, std::move(memory_manager)); } + arrow::Result> MakeServerAsync( + const std::string& scheme, AsyncGenericFlightServerBase* async_base, + std::shared_ptr memory_manager) const { + auto it = async_server_factories_.find(scheme); + if (it == async_server_factories_.end()) { + return Status::KeyError("No async server transport implementation for ", scheme); + } + return it->second(async_base, std::move(memory_manager)); + } Status RegisterClient(const std::string& scheme, ClientFactory factory) { auto it = client_factories_.insert({scheme, std::move(factory)}); if (!it.second) { @@ -150,10 +159,19 @@ class TransportRegistry::Impl final { } return Status::OK(); } + Status RegisterAsyncServer(const std::string& scheme, AsyncServerFactory factory) { + auto it = async_server_factories_.insert({scheme, std::move(factory)}); + if (!it.second) { + return Status::Invalid("Async server transport already registered for ", scheme); + } + return Status::OK(); + } private: std::unordered_map client_factories_; std::unordered_map server_factories_; + std::unordered_map + async_server_factories_; }; TransportRegistry::TransportRegistry() { impl_ = std::make_unique(); } @@ -167,6 +185,11 @@ arrow::Result> TransportRegistry::MakeServer( std::shared_ptr memory_manager) const { return impl_->MakeServer(scheme, base, std::move(memory_manager)); } +arrow::Result> TransportRegistry::MakeServerAsync( + const std::string& scheme, AsyncGenericFlightServerBase* async_base, + std::shared_ptr memory_manager) const { + return impl_->MakeServerAsync(scheme, async_base, std::move(memory_manager)); +} Status TransportRegistry::RegisterClient(const std::string& scheme, ClientFactory factory) { return impl_->RegisterClient(scheme, std::move(factory)); @@ -175,6 +198,10 @@ Status TransportRegistry::RegisterServer(const std::string& scheme, ServerFactory factory) { return impl_->RegisterServer(scheme, std::move(factory)); } +Status TransportRegistry::RegisterAsyncServer(const std::string& scheme, + AsyncServerFactory factory) { + return impl_->RegisterAsyncServer(scheme, std::move(factory)); +} TransportRegistry* GetDefaultTransportRegistry() { static TransportRegistry kRegistry; diff --git a/cpp/src/arrow/flight/transport.h b/cpp/src/arrow/flight/transport.h index 6b9694318915..4f30afc7927d 100644 --- a/cpp/src/arrow/flight/transport.h +++ b/cpp/src/arrow/flight/transport.h @@ -72,6 +72,7 @@ namespace ipc { class Message; } namespace flight { +class AsyncGenericFlightServerBase; class FlightStatusDetail; namespace internal { @@ -216,6 +217,9 @@ class ARROW_FLIGHT_EXPORT TransportRegistry { using ClientFactory = std::function>()>; using ServerFactory = std::function>( FlightServerBase*, std::shared_ptr memory_manager)>; + using AsyncServerFactory = + std::function>( + AsyncGenericFlightServerBase*, std::shared_ptr memory_manager)>; TransportRegistry(); ~TransportRegistry(); @@ -226,8 +230,13 @@ class ARROW_FLIGHT_EXPORT TransportRegistry { const std::string& scheme, FlightServerBase* base, std::shared_ptr memory_manager) const; + arrow::Result> MakeServerAsync( + const std::string& scheme, AsyncGenericFlightServerBase* async_base, + std::shared_ptr memory_manager) const; + Status RegisterClient(const std::string& scheme, ClientFactory factory); Status RegisterServer(const std::string& scheme, ServerFactory factory); + Status RegisterAsyncServer(const std::string& scheme, AsyncServerFactory factory); private: class Impl; diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_compat.h b/cpp/src/arrow/flight/transport/grpc/async_grpc_compat.h new file mode 100644 index 000000000000..1647c3895dc5 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_compat.h @@ -0,0 +1,45 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +// Version gating for the async Flight server, which is built on gRPC's generic +// callback API (grpc::CallbackGenericService, grpc::GenericCallbackServerContext +// and grpc::ServerGenericBidiReactor). +// +// That API appeared in gRPC 1.65. One header covers every supported version: +// grpcpp/generic/async_generic_service.h declares the callback API itself in +// 1.65, and from 1.66 on it just includes +// grpcpp/generic/callback_generic_service.h (which does not exist before 1.66, +// so including it directly would break a 1.65 build). +// +// ARROW_FLIGHT_HAS_ASYNC_SERVER hides that detail from the rest of the +// transport: when it is 0 the async server is not compiled and the gRPC +// transport answers async server requests with NotImplemented. +// +// GRPC_CPP_VERSION_CHECK comes from util_internal.h and evaluates to 0 for +// gRPC releases predating the version macros, which covers the older system +// gRPC packages some builds link against. + +#include "arrow/flight/transport/grpc/util_internal.h" + +#if GRPC_CPP_VERSION_CHECK(1, 65, 0) +# define ARROW_FLIGHT_HAS_ASYNC_SERVER 1 +# include +#else +# define ARROW_FLIGHT_HAS_ASYNC_SERVER 0 +#endif diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_do_get.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_do_get.cc new file mode 100644 index 000000000000..1f5fd6a18bd0 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_do_get.cc @@ -0,0 +1,186 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_service_internal.h" + +#include +#include +#include +#include +#include + +namespace arrow::flight::transport::grpc::detail { + +namespace { + +/// Serve one DoGet RPC over the generic callback API. +class DoGetReactor : public AsyncReactorBase { + public: + DoGetReactor(AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) + : AsyncReactorBase(std::move(flight_context)), base_(base) { + StartRead(&request_buf_); + } + + void OnReadDone(bool ok) override { + if (!ok) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, "Failed to read request")); + return; + } + + const auto ticket = ParseTicket(); + if (!ticket.ok()) { + FinishOnce(ticket.status()); + return; + } + + arrow::Future> future = + base_->DoGetAsync(flight_context(), *ticket); + future.AddCallback( + [this, token = hold()]( + const arrow::Result>& result) { + if (!finished()) { + if (!result.ok()) { + FinishOnce(result.status()); + } else { + std::shared_ptr stream = result.ValueOrDie(); + if (stream == nullptr) { + FinishOnce(arrow::Status::KeyError("No data in this flight")); + } else { + bool already_closed; + { + std::lock_guard lock(stream_mutex_); + async_data_stream_ = stream; + already_closed = closed_; + } + if (already_closed) { + ARROW_WARN_NOT_OK(stream->Close(), + "DoGet: closing a cancelled call's stream failed"); + } else { + WriteNextPayload(); + } + } + } + } + }); + } + + void OnWriteDone(bool ok) override { + if (finished()) { + return; + } + if (!ok) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, "Write failed")); + ARROW_WARN_NOT_OK(CloseStreamOnce(), + "DoGet: closing the data stream after a failed write failed"); + return; + } + WriteNextPayload(); + } + + void OnCancel() override { + FinishOnce(arrow::Status::Cancelled()); + ARROW_WARN_NOT_OK(CloseStreamOnce(), + "DoGet: closing the stream of a cancelled call failed"); + } + + private: + arrow::Result ParseTicket() { + return ParseProtoRequest(request_buf_, "Ticket"); + } + + Status CloseStreamOnce() { + std::shared_ptr stream; + { + std::lock_guard lock(stream_mutex_); + if (closed_) { + return Status::OK(); + } + closed_ = true; + stream = async_data_stream_; + } + if (stream == nullptr) { + return Status::OK(); + } + return stream->Close(); + } + + void WriteNextPayload() { + if (finished()) { + return; + } + HoldToken token = hold(); + + arrow::Future> next; + if (!wrote_schema_) { + next = async_data_stream_->GetSchemaPayloadAsync().Then( + [](FlightPayload payload) -> std::optional { + return {std::move(payload)}; + }); + } else { + next = async_data_stream_->NextAsync(); + } + next.AddCallback([this, token = token.release()]( + arrow::Result> result) { + if (!finished()) { + if (!result.ok()) { + FinishOnce(result.status()); + ARROW_WARN_NOT_OK( + CloseStreamOnce(), + "DoGet: closing the data stream after a failed payload failed"); + } else { + wrote_schema_ = true; + std::optional payload = std::move(*result); + if (!payload.has_value()) { + FinishOnce(CloseStreamOnce()); + } else { + bool own_buffer = false; + const ::grpc::Status grpc_status = + FlightDataSerialize(*payload, &write_buf_, &own_buffer); + if (!grpc_status.ok()) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, + grpc_status.error_message())); + ARROW_WARN_NOT_OK( + CloseStreamOnce(), + "DoGet: closing the data stream after a failed serialization failed"); + } else { + WriteIfLive([this] { StartWrite(&write_buf_); }); + } + } + } + } + }); + } + + AsyncGenericFlightServerBase* base_; + ::grpc::ByteBuffer request_buf_; + ::grpc::ByteBuffer write_buf_; + std::shared_ptr async_data_stream_; + bool wrote_schema_ = false; + /// Guards async_data_stream_ and closed_, so adopting the stream and the cancellation + /// teardown cannot interleave. + std::mutex stream_mutex_; + bool closed_ = false; +}; + +} // namespace + +::grpc::ServerGenericBidiReactor* MakeDoGetReactor(AsyncCallContext flight_context, + AsyncGenericFlightServerBase* base) { + return new DoGetReactor(std::move(flight_context), base); +} + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_do_get_test.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_do_get_test.cc new file mode 100644 index 000000000000..1bed2683febe --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_do_get_test.cc @@ -0,0 +1,329 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_test_fixtures.h" + +namespace arrow::flight::transport::grpc { + +TEST(AsyncGrpcTest, BasicDoGet) { + TestFlightServer flight_server; + AsyncGenericFlightService service(&flight_server, MakeAsyncHelper()); + + int port = 0; + ::grpc::ServerBuilder builder; + builder.AddListeningPort("localhost:0", ::grpc::InsecureServerCredentials(), &port); + builder.RegisterCallbackGenericService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + ASSERT_NE(port, 0); + + // Connect the existing arrow::flight::FlightClient implementation. + std::string uri = "grpc://localhost:" + std::to_string(port); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse(uri)); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(location)); + + // Call DoGet: the ticket must reach TestFlightServer::DoGetAsync. + Ticket ticket{"basic-do-get-ticket"}; + ASSERT_OK_AND_ASSIGN(auto reader, client->DoGet(ticket)); + + // Read schema. + ASSERT_OK_AND_ASSIGN(auto schema, reader->GetSchema()); + ASSERT_EQ(schema->num_fields(), 2); + ASSERT_EQ(schema->field(0)->name(), "a"); + ASSERT_EQ(schema->field(1)->name(), "b"); + + // Read batches. + int batch_count = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(auto chunk, reader->Next()); + if (!chunk.data) break; + ASSERT_EQ(chunk.data->num_rows(), 5); + ASSERT_OK_AND_ASSIGN(auto value, chunk.data->column(0)->GetScalar(0)); + ASSERT_TRUE(value->Equals(*arrow::MakeScalar(int64_t(1)))); + batch_count++; + } + ASSERT_EQ(batch_count, 3); + + // Cleanup. The ticket is asserted only after the server stopped, so the + // DoGet callback cannot race with the read. + ASSERT_OK(client->Close()); + server->Shutdown(); + server->Wait(); + + ASSERT_EQ(flight_server.ticket(), "basic-do-get-ticket"); +} + +TEST(AsyncGrpcTest, SchemaOnlyDoGet) { + TestFlightServer flight_server; + AsyncGenericFlightService service(&flight_server, MakeAsyncHelper()); + + int port = 0; + ::grpc::ServerBuilder builder; + builder.AddListeningPort("localhost:0", ::grpc::InsecureServerCredentials(), &port); + builder.RegisterCallbackGenericService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + ASSERT_NE(port, 0); + + std::string uri = "grpc://localhost:" + std::to_string(port); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse(uri)); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(location)); + + // This ticket is served as a stream with no batches at all. + Ticket ticket{kSchemaOnlyTicket}; + ASSERT_OK_AND_ASSIGN(auto reader, client->DoGet(ticket)); + + // The schema is served... + ASSERT_OK_AND_ASSIGN(auto schema, reader->GetSchema()); + ASSERT_EQ(schema->num_fields(), 2); + ASSERT_EQ(schema->field(0)->name(), "a"); + ASSERT_EQ(schema->field(1)->name(), "b"); + + // ...and the very first read reports the end of stream instead of failing. + ASSERT_OK_AND_ASSIGN(auto chunk, reader->Next()); + ASSERT_EQ(chunk.data, nullptr); + + // Cleanup. The ticket is asserted only after the server stopped, so the + // DoGet callback cannot race with the read. + ASSERT_OK(client->Close()); + server->Shutdown(); + server->Wait(); + + ASSERT_EQ(flight_server.ticket(), kSchemaOnlyTicket); +} + +// A DoGet stream whose first payload never arrives: the producer is parked +// forever, so only a client cancel can ever end the RPC. `close_count` and +// `next_awaits` let the test observe, from the test thread, that the payload +// demand is live and that Close() was delivered exactly once. +class HangingDoGetStream final : public AsyncFlightDataStream { + public: + HangingDoGetStream(std::shared_ptr resolved, + std::atomic* next_awaits, std::atomic* close_count) + : resolved_(std::move(resolved)), + next_awaits_(next_awaits), + close_count_(close_count) {} + + std::shared_ptr schema() override { return resolved_->schema(); } + + arrow::Future GetSchemaPayloadAsync() override { + return resolved_->GetSchemaPayloadAsync(); + } + + arrow::Future> NextAsync() override { + next_awaits_->fetch_add(1); + std::lock_guard lock(mutex_); + pending_ = arrow::Future>::Make(); + return pending_; + } + + // The transport's stop signal: resolve the pending payload so the producer + // does not park forever (mirrors the example's TimedAsyncStream::Close). + Status Close() override { + close_count_->fetch_add(1); + std::lock_guard lock(mutex_); + if (pending_.is_valid() && !pending_.is_finished()) { + pending_.MarkFinished(arrow::Status::Cancelled()); + } + return resolved_->Close(); + } + + private: + std::shared_ptr resolved_; + std::atomic* next_awaits_; + std::atomic* close_count_; + std::mutex mutex_; + arrow::Future> pending_; +}; + +// DoGet server with a configurable hang: kPrepare leaves DoGetAsync() pending +// (no stream ever exists), kStream resolves the stream but leaves the first +// payload pending. +class HangingDoGetServer final : public AsyncGenericFlightServerBase { + public: + enum class HangAt { kPrepare, kStream }; + explicit HangingDoGetServer(HangAt hang) : hang_(hang) {} + + arrow::Future> DoGetAsync( + const ServerCallContext&, const Ticket&) override { + if (hang_ == HangAt::kPrepare) { + return arrow::Future>::Make(); + } + auto reader = arrow::RecordBatchReader::Make( + {}, arrow::schema({arrow::field("value", arrow::int64())})) + .ValueOrDie(); + return arrow::Future>::MakeFinished( + std::make_shared( + std::make_shared(std::move(reader)), &next_awaits, + &stream_closes)); + } + + std::atomic next_awaits{0}; + std::atomic stream_closes{0}; + + private: + HangAt hang_; +}; + +TEST(AsyncGrpcTest, DoGetClientCancelWithPendingPayloadFinishes) { + // The client gives up while the producer's NextAsync() is pending: the + // transport must finish the RPC itself and deliver Close() to the producer, + // instead of waiting for a payload that will not come. + const int64_t before = detail::AsyncReactorBase::live_reactors(); + auto trace = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + HangingDoGetServer flight_server(HangingDoGetServer::HangAt::kStream); + FlightServerOptions options(location); + options.middleware = { + {"recording", std::make_shared(trace)}}; + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + // DoGet returns once the schema payload is written,the payload after it is + // the one that never comes. Wait until the server has actually asked for + // it, so the cancel lands on a live demand (not a race with the schema + // write). + ASSERT_OK_AND_ASSIGN(auto reader, client->DoGet(Ticket{"hang"})); + ASSERT_TRUE(WaitFor([&flight_server] { return flight_server.next_awaits.load() >= 1; })) + << "the server never asked for its first payload"; + + reader->Cancel(); + reader.reset(); + ASSERT_OK(client->Close()); + + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // The producer was told to stop, exactly once. + EXPECT_EQ(1, flight_server.stream_closes.load()); + // And the reactor is gone even though its payload future never completed. + EXPECT_TRUE(WaitFor([&] { + return detail::AsyncReactorBase::live_reactors() == before; + })) << "a reactor leaked: still alive after the server stopped (before=" + << before << ", now=" << detail::AsyncReactorBase::live_reactors() << ")"; + // And the call reported completion exactly once (the CAS in FinishOnce). + std::lock_guard guard(trace->mutex); + EXPECT_EQ(1, trace->call_completed); +} + +TEST(AsyncGrpcTest, DoGetClientDeadlineWhilePreparingFinishes) { + // DoGetAsync() never resolves and the client's deadline expires: the RPC + // must end anyway: the Shutdown()/Wait() below would otherwise block. No + // stream exists here, so Close() has nothing to deliver and the producer + // notices through ServerCallContext::is_cancelled(). + auto trace = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + HangingDoGetServer flight_server(HangingDoGetServer::HangAt::kPrepare); + FlightServerOptions options(location); + options.middleware = { + {"recording", std::make_shared(trace)}}; + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + FlightCallOptions impatient; + impatient.timeout = std::chrono::milliseconds(500); + auto reader = client->DoGet(impatient, Ticket{"hang"}); + EXPECT_FALSE(reader.ok()) << "the deadline must expire while the server prepares"; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + EXPECT_EQ(0, flight_server.stream_closes.load()) + << "no stream was ever prepared, so there is nothing to Close"; + std::lock_guard guard(trace->mutex); + EXPECT_EQ(1, trace->call_completed); +} + +/// A DoGet stream whose first payload is an error: the transport must fail the +/// RPC and still tell the producer the RPC is over through Close(). +class ErrorPayloadDoGetStream final : public AsyncFlightDataStream { + public: + explicit ErrorPayloadDoGetStream(std::atomic* close_count) + : close_count_(close_count) {} + + std::shared_ptr schema() override { + return arrow::schema({arrow::field("value", arrow::int64())}); + } + + arrow::Future GetSchemaPayloadAsync() override { + return resolved_->GetSchemaPayloadAsync(); + } + + arrow::Future> NextAsync() override { + return arrow::Future>::MakeFinished( + arrow::Status::IOError("the producer failed")); + } + + Status Close() override { + close_count_->fetch_add(1); + return Status::OK(); + } + + std::shared_ptr resolved_ = std::make_shared( + arrow::RecordBatchReader::Make( + {}, arrow::schema({arrow::field("value", arrow::int64())})) + .ValueOrDie()); + + private: + std::atomic* close_count_; +}; + +/// Serves ErrorPayloadDoGetStream and counts the Close() it was handed. +class ErrorPayloadDoGetServer final : public AsyncGenericFlightServerBase { + public: + arrow::Future> DoGetAsync( + const ServerCallContext&, const Ticket&) override { + return arrow::Future>::MakeFinished( + std::make_shared(&stream_closes)); + } + + std::atomic stream_closes{0}; +}; + +TEST(AsyncGrpcTest, DoGetPumpErrorStillClosesTheStream) { + // The producer's first NextAsync() fails: the RPC ends with that status, and + // the producer must still be told the RPC is over through Close(), like the + // write-failure and cancel paths already do. + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + ErrorPayloadDoGetServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + // DoGet returns once the schema payload is written, the error payload follows. + ASSERT_OK_AND_ASSIGN(auto reader, client->DoGet(Ticket{"error-payload"})); + while (true) { + auto chunk = reader->Next(); + if (!chunk.ok()) { + EXPECT_EQ(chunk.status().code(), arrow::StatusCode::IOError) << chunk.status(); + break; + } + if (!chunk->data) { + break; + } + } + reader.reset(); + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + EXPECT_EQ(1, flight_server.stream_closes.load()) + << "the producer was never told the RPC is over"; +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_do_put.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_do_put.cc new file mode 100644 index 000000000000..6449912f8e1d --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_do_put.cc @@ -0,0 +1,255 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_service_internal.h" + +#include "arrow/flight/transport/grpc/customize_grpc.h" + +#include +#include +#include + +namespace arrow::flight::transport::grpc::detail { + +namespace { + +/// Serve one DoPut RPC over the generic callback API. +class DoPutReactor : public AsyncReactorBase, + public arrow::flight::internal::FlightDataListenerTransport { + public: + DoPutReactor(AsyncCallContext flight_context, + std::shared_ptr listener) + : AsyncReactorBase(std::move(flight_context)), + listener_(std::move(listener)), + decoder_(listener_) { + arrow::flight::internal::FlightDataListenerTransport::Install(listener_, this); + StartRead(&request_buf_); + } + + void Hold() override { AsyncReactorBase::Hold(); } + + void ReleaseHold() override { AsyncReactorBase::ReleaseHold(); } + + void CancelUpload(Status status) override { + if (status.ok()) { + status = Status::Invalid("Cancel() needs a non-OK status"); + } + AsyncReactorBase::Hold(); + FinishUpload(std::move(status)); + AsyncReactorBase::ReleaseHold(); + } + + Status WriteMetadata(const Buffer& app_metadata) override { + std::lock_guard lock(write_mutex()); + if (finished()) { + return Status::Invalid( + "the upload is over: WriteMetadata() is only usable while the " + "upload's RPC is running"); + } + if (metadata_in_flight_ || ending_) { + // One message in flight at a time: OnWriteDone serializes the PutResults. + return Status::Invalid( + "a PutResult message is already in flight: await it before writing " + "another metadata message"); + } + metadata_in_flight_ = true; + pb::PutResult pb_result; + if (app_metadata.size() > 0) { + pb_result.set_app_metadata(app_metadata.data(), app_metadata.size()); + } + write_buf_ = MakeWriteBuffer(pb_result); + StartWrite(&write_buf_); + return Status::OK(); + } + + void OnReadDone(bool ok) override { + if (!ok) { + // Decide under the lock, act outside it: FinishUpload() runs the + // listener's OnFinish(), which is application code and may itself call + // WriteMetadata(). That takes writer_mutex_ and then write_mutex_, the + // reverse of a caller holding write_mutex_ and waiting for writer_mutex_, + // so invoking it here would deadlock the two threads. + { + std::lock_guard lock(write_mutex()); + if (finished()) { + return; + } + if (read_descriptor_) { + // The client half-closed: this ack ends the upload. It takes the + // same write slot as a mid-upload WriteMetadata, which may be on the + // wire right now: starting a second write would clobber write_buf_ + // and the RPC would never see its OnWriteDone. + ending_ = true; + if (metadata_in_flight_) { + return; // OnWriteDone sends the ending ack once the slot frees up + } + StartEndingAck(); + return; + } + } + FinishUpload( + MakeFlightError(FlightStatusCode::Failed, + "Client sent an empty upload: the first message must carry the " + "descriptor")); + return; + } + std::shared_ptr arrow_buf; + const Status wrap_status = WrapGrpcBuffer(&request_buf_, &arrow_buf); + if (!wrap_status.ok()) { + FinishUpload( + MakeFlightError(FlightStatusCode::Internal, + "Failed to wrap gRPC buffer: " + wrap_status.message())); + return; + } + + if (!read_descriptor_) { + // The first message must carry the descriptor of the upload, exactly as + // the sync transport requires. + // Without this the listener never hears OnDescriptor and the upload is + // acknowledged as if it were well formed. + read_descriptor_ = true; + internal::FlightData data; + const ::grpc::Status deserialized = FlightDataDeserialize(&request_buf_, &data); + if (!deserialized.ok()) { + FinishUpload( + MakeFlightError(FlightStatusCode::Internal, deserialized.error_message())); + return; + } + if (!data.descriptor) { + FinishUpload(Status::IOError("Descriptor missing on first message")); + return; + } + } + + Future<> decode_status = decoder_.Consume(std::move(arrow_buf)); + decode_status.AddCallback([this, token = hold()](arrow::Status status) { + if (finished()) { + // The RPC is over (cancelled while the listener was working): the + // message is abandoned. + return; + } + if (!status.ok()) { + FinishUpload(std::move(status)); + return; + } + // Arm the next read through WriteIfLive(): this continuation runs on the + // decoder's thread, so a cancellation can finish the RPC right after the + // finished() check above. WriteIfLive re-checks under write_mutex_, which + // FinishOnce also holds, so the read is never started on a dead call. + WriteIfLive([this] { StartRead(&request_buf_); }); + }); + } + + void OnWriteDone(bool ok) override { + // Decide under the lock, act outside it: FinishUpload() runs the listener's + // OnFinish(), which is application code and may call back into WriteMetadata(). + Status finish_status; + bool ended = false; + { + std::lock_guard lock(write_mutex()); + if (finished()) { + return; + } + if (!ok) { + finish_status = MakeFlightError(FlightStatusCode::Internal, "Write failed"); + ended = true; + } else if (ending_ack_in_flight_) { + // The ending ack reached the client: that is the clean end of the upload. + ended = true; + } else { + metadata_in_flight_ = false; + if (ending_) { + // The client half-closed while a metadata message was on the wire: + // with the slot free, the ending ack is the next write. + StartEndingAck(); + } + } + } + if (ended) { + FinishUpload(std::move(finish_status)); + } + } + + void OnCancel() override { + // The RPC is cancelled: report it to the listener but do not let a failing + // OnFinish replace the cancellation the client is seeing. The report is + // fire-and-forget here; keeping a reference alive for it is unnecessary as + // FinishOnce() below releases the reactor's hold and the callback only + // touches immutable state. + const auto status = Status::Cancelled("the client cancelled the upload"); + static_cast(arrow::flight::internal::FlightDataListenerTransport::ReportFinish( + listener_, status)); + FinishOnce(status); + } + + void OnDone() override { + arrow::flight::internal::FlightDataListenerTransport::Clear(listener_); + ReleaseHold(); + } + + private: + /// Send the empty ack that ends the upload. Caller holds write_mutex_. + void StartEndingAck() { + ending_ack_in_flight_ = true; + pb::PutResult pb_result; + write_buf_ = MakeWriteBuffer(pb_result); + StartWrite(&write_buf_); + } + + void FinishUpload(Status status) { + // Chain the RPC's completion on the listener's OnFinish() future instead of + // waiting for it: FinishOnce() must not run on a gRPC callback thread while + // the application is still committing, and a listener that fails + // asynchronously must not make the client see success. Keep the reactor + // alive across the wait. + Status reported = status; + arrow::flight::internal::FlightDataListenerTransport::ReportFinish(listener_, + std::move(status)) + .AddCallback([this, token = hold(), + reported = std::move(reported)](const Status& finish_status) { + // A failing OnFinish overrides a successful upload, but must not mask + // a status the upload already failed with. + FinishOnce(reported.ok() ? finish_status : reported); + }); + } + + std::shared_ptr listener_; + AsyncFlightMessageDecoder decoder_; + /// Whether the descriptor of the first message has been checked. + bool read_descriptor_ = false; + // The write side is serialized by AsyncReactorBase::write_mutex_, which + // FinishOnce() also takes: WriteMetadata() runs on the listener's thread + // while OnReadDone()/OnWriteDone() run on gRPC callback threads, and this + // must be the same mutex the base uses or finishing can race a write. + /// Whether a mid-upload metadata message is on the wire. + bool metadata_in_flight_ = false; + /// Whether the client half-closed, so the ending ack is owed. + bool ending_ = false; + /// Whether that ending ack is the write in flight. + bool ending_ack_in_flight_ = false; + ::grpc::ByteBuffer request_buf_; + ::grpc::ByteBuffer write_buf_; +}; + +} // namespace + +::grpc::ServerGenericBidiReactor* MakeDoPutReactor( + AsyncCallContext flight_context, std::shared_ptr listener) { + return new DoPutReactor(std::move(flight_context), std::move(listener)); +} + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_do_put_test.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_do_put_test.cc new file mode 100644 index 000000000000..ceccdb91ed6f --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_do_put_test.cc @@ -0,0 +1,762 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// DoPut over the async generic service: the push decoder feeding the per-RPC +// listener, the acknowledgement that makes the client's Close() return, the +// PutResult metadata channel, and every way an upload ends -- clean, rejected by +// the listener, rejected on its descriptor, cancelled from the server side, +// client deadline, and the client vanishing while the listener is parked. + +#include "arrow/flight/transport/grpc/async_grpc_test_fixtures.h" + +namespace arrow::flight::transport::grpc { + +TEST(AsyncGrpcTest, BasicDoPut) { + // One listener for the whole test: shared with the service's factory and + // kept here so it can be asserted on after the RPC. + auto listener = std::make_shared(); + TestFlightServer flight_server; + flight_server.set_listener(listener); + AsyncGenericFlightService service(&flight_server, MakeAsyncHelper()); + + int port = 0; + ::grpc::ServerBuilder builder; + builder.AddListeningPort("localhost:0", ::grpc::InsecureServerCredentials(), &port); + builder.RegisterCallbackGenericService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + ASSERT_NE(port, 0); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(port)); + + auto schema = TwoColumnSchema(); + auto batch = TwoColumnBatch(schema); + + // Call DoPut and upload the same batch twice. Close() returns the server's + // final status, i.e. it only succeeds once the acknowledgement written on + // end of stream has been read. + FlightDescriptor descriptor = FlightDescriptor::Path({"test"}); + ASSERT_OK_AND_ASSIGN(auto result, client->DoPut(descriptor, schema)); + ASSERT_OK(result.writer->WriteRecordBatch(*batch)); + ASSERT_OK(result.writer->WriteRecordBatch(*batch)); + ASSERT_OK(result.writer->Close()); + + // Cleanup. The listener is asserted only after the server stopped, so the + // DoPut callback cannot race with the reads. + ASSERT_OK(client->Close()); + server->Shutdown(); + server->Wait(); + + // The listener saw the schema once and both batches, in order. + ASSERT_EQ(listener->schema_count(), 1); + AssertTwoColumnSchema(listener->schema()); + ASSERT_EQ(listener->batches().size(), 2); + for (const auto& chunk : listener->batches()) { + ASSERT_EQ(chunk->num_rows(), 3); + ASSERT_EQ(chunk->num_columns(), 2); + } + // No metadata-only (error/app-metadata) chunks, and nothing the consumer + // rejected: a non-OK OnNext status would have failed DoPut with no ack. + ASSERT_EQ(listener->metadata_only_count(), 0); + ASSERT_OK(listener->last_status()); +} + +TEST(AsyncGrpcTest, EmptyDoPut) { + TestFlightServer flight_server; + auto listener = std::make_shared(); + flight_server.set_listener(listener); + AsyncGenericFlightService service(&flight_server, MakeAsyncHelper()); + + int port = 0; + ::grpc::ServerBuilder builder; + builder.AddListeningPort("localhost:0", ::grpc::InsecureServerCredentials(), &port); + builder.RegisterCallbackGenericService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + ASSERT_NE(port, 0); + + std::string uri = "grpc://localhost:" + std::to_string(port); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse(uri)); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(location)); + + auto schema = arrow::schema( + {arrow::field("a", arrow::int64()), arrow::field("b", arrow::int64())}); + + // Upload nothing at all: the acknowledgement the server writes back on the + // client's half-close is what makes Close() return. + FlightDescriptor descriptor = FlightDescriptor::Path({"empty-put"}); + ASSERT_OK_AND_ASSIGN(auto result, client->DoPut(descriptor, schema)); + ASSERT_OK(result.writer->Close()); + + // Cleanup before looking at the listener, so the DoPut callback cannot race + // with the reads. + ASSERT_OK(client->Close()); + server->Shutdown(); + server->Wait(); + + // No data of any kind reached the listener. + ASSERT_EQ(listener->batches().size(), 0); + ASSERT_EQ(listener->metadata_only_count(), 0); + ASSERT_OK(listener->last_status()); + // The sync client sends the schema payload eagerly, even when no batch + // follows, so the listener does see the schema here. Deliberately not + // asserted: that is the client's business, not the transport's. +} + +TEST(AsyncGrpcTest, MetadataOnlyPutChunk) { + TestFlightServer flight_server; + auto listener = std::make_shared(); + flight_server.set_listener(listener); + AsyncGenericFlightService service(&flight_server, MakeAsyncHelper()); + + int port = 0; + ::grpc::ServerBuilder builder; + builder.AddListeningPort("localhost:0", ::grpc::InsecureServerCredentials(), &port); + builder.RegisterCallbackGenericService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + ASSERT_NE(port, 0); + + std::string uri = "grpc://localhost:" + std::to_string(port); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse(uri)); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(location)); + + auto schema = arrow::schema( + {arrow::field("a", arrow::int64()), arrow::field("b", arrow::int64())}); + + FlightDescriptor descriptor = FlightDescriptor::Path({"metadata-only-put"}); + ASSERT_OK_AND_ASSIGN(auto result, client->DoPut(descriptor, schema)); + auto app_metadata = arrow::Buffer::FromString("metadata-only-upload"); + ASSERT_OK(result.writer->WriteMetadata(app_metadata)); + ASSERT_OK(result.writer->Close()); + + ASSERT_OK(client->Close()); + server->Shutdown(); + server->Wait(); + + // Exactly one chunk reached the listener: no data, that exact metadata. + ASSERT_EQ(listener->batches().size(), 0); + ASSERT_EQ(listener->metadata_only_count(), 1); + ASSERT_EQ(listener->metadata_chunks().size(), 1); + ASSERT_NE(listener->metadata_chunks()[0], nullptr); + ASSERT_EQ(listener->metadata_chunks()[0]->ToString(), "metadata-only-upload"); +} + +TEST(AsyncGrpcTest, DoPutRejectedByListener) { + TestFlightServer flight_server; + auto listener = std::make_shared(); + // The consumer rejects the upload. + listener->set_next_status(arrow::Status::Invalid("listener rejected this upload")); + flight_server.set_listener(listener); + AsyncGenericFlightService service(&flight_server, MakeAsyncHelper()); + + int port = 0; + ::grpc::ServerBuilder builder; + builder.AddListeningPort("localhost:0", ::grpc::InsecureServerCredentials(), &port); + builder.RegisterCallbackGenericService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + ASSERT_NE(port, 0); + + std::string uri = "grpc://localhost:" + std::to_string(port); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse(uri)); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(location)); + + auto schema = TwoColumnSchema(); + auto batch = TwoColumnBatch(schema); + + // The rejection must reach the client as the upload's status, whichever of + // the three surfaces reports it first (in practice the writer's Close()). + FlightDescriptor descriptor = FlightDescriptor::Path({"rejected-put"}); + auto [surface, status] = UploadOneBatch(client.get(), descriptor, schema, batch); + ASSERT_FALSE(status.ok()) + << "the listener's rejection never reached the client, surfaced on " << surface; + ASSERT_EQ(status.code(), arrow::StatusCode::Invalid); + ASSERT_NE(status.message().find("listener rejected this upload"), std::string::npos) + << "unexpected status: " << status; + // The rejection must also be what OnFinish reports. + ASSERT_EQ(listener->finish_status().code(), arrow::StatusCode::Invalid) + << "OnFinish got: " << listener->finish_status(); + + ASSERT_OK(client->Close()); + server->Shutdown(); + server->Wait(); +} + +TEST(AsyncGrpcTest, CreateDoPutListenerServesUpload) { + // The upload half of the server class API: CreateDoPutListener() hands out + // the per-RPC listener, with no options-level factory. + TestUploadServer flight_server; + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_GT(flight_server.port(), 0); + + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = TwoColumnSchema(); + auto batch = TwoColumnBatch(schema); + + const FlightDescriptor descriptor = FlightDescriptor::Path({"virtual", "h2"}); + auto [surface, status] = UploadOneBatch(client.get(), descriptor, schema, batch); + ASSERT_TRUE(status.ok()) << "upload failed on " << surface << ": " << status; + + // Cleanup before reading the recordings: the listener's callbacks run on a + // gRPC thread. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // The virtual was consulted exactly once, and its listener saw the upload. + ASSERT_EQ(flight_server.calls(), 1); + const auto& listener = flight_server.listener(0); + ASSERT_EQ(listener->descriptor_count(), 1); + ASSERT_EQ(listener->descriptors()[0], descriptor); + ASSERT_EQ(listener->schema_count(), 1); + ASSERT_EQ(listener->batches().size(), 1); + ASSERT_EQ(listener->batches()[0]->num_rows(), 3); +} + +TEST(AsyncGrpcTest, DoPutWithoutListenerIsUnimplemented) { + // A server class that does not hand out a listener (CreateDoPutListener() + // returns nullptr, the base default) does not accept uploads: the RPC is + // answered UNIMPLEMENTED. There is no options-level fallback. + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = arrow::schema({arrow::field("a", arrow::int64())}); + auto batch = arrow::RecordBatch::Make(schema, 2, + {arrow::ArrayFromJSON(arrow::int64(), "[7, 8]")}); + auto [surface, status] = UploadOneBatch( + client.get(), FlightDescriptor::Path({"no-listener"}), schema, batch); + ASSERT_FALSE(status.ok()) << "the upload succeeded on " << surface + << " with no listener to serve it"; + ASSERT_EQ(status.code(), arrow::StatusCode::NotImplemented) << status; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, DoPutListenerSeesCleanFinish) { + // A client that ends the upload normally (DoneWriting + Close) makes the + // listener's OnFinish fire once, with OK: that is how a consumer knows the + // upload ended and it can commit what it accumulated. + TestUploadServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = OneColumnSchema(); + auto batch = OneColumnBatch(schema); + auto [surface, status] = + UploadOneBatch(client.get(), FlightDescriptor::Path({"finish"}), schema, batch); + ASSERT_TRUE(status.ok()) << "upload failed on " << surface << ": " << status; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // OnFinish ran exactly once, and said the upload ended cleanly. + ASSERT_EQ(flight_server.listener(0)->finish_count(), 1); + ASSERT_TRUE(flight_server.listener(0)->finish_status().ok()) + << "unexpected finish status: " << flight_server.listener(0)->finish_status(); +} + +TEST(AsyncGrpcTest, OnFinishMayWriteMetadataWithoutDeadlocking) { + // OnFinish is application code and may call back into the listener's writer. + // The transport must not hold its write lock across that call: doing so + // deadlocks this upload (verified: the RPC never finishes). + RecordingUploadServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = OneColumnSchema(); + ASSERT_OK_AND_ASSIGN(auto put, + client->DoPut(FlightDescriptor::Path({"onfinish"}), schema)); + ASSERT_OK(put.writer->WriteRecordBatch(*OneColumnBatch(schema))); + ARROW_UNUSED(put.writer->Close()); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + auto listener = flight_server.WaitForListener(0); + ASSERT_NE(listener, nullptr) << "the server never handed out a DoPut listener"; + // The write is refused (the upload is ending) rather than deadlocking. + EXPECT_TRUE(listener->finish_write.IsInvalid()) << listener->finish_write; + ASSERT_EQ(1, listener->finish_count()); +} + +TEST(AsyncGrpcTest, DoPutWritesMetadataBackToTheClient) { + // The listener's writer() sends PutResult messages during the upload, the way + // the sync FlightMetadataWriter does: the client reads them back with + // ReadMetadata() between writes. + MetadataReplyingUploadServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = OneColumnSchema(); + auto batch = OneColumnBatch(schema); + + ASSERT_OK_AND_ASSIGN(auto put, + client->DoPut(FlightDescriptor::Path({"metadata"}), schema)); + auto& writer = *put.writer; + auto& reader = *put.reader; + + // Write one batch, then read the acknowledgement the server sent for it. + ASSERT_OK(writer.WriteRecordBatch(*batch)); + std::shared_ptr metadata; + ASSERT_OK(reader.ReadMetadata(&metadata)); + ASSERT_NE(nullptr, metadata) << "the server's metadata never reached the client"; + ASSERT_EQ("ack 2 rows", metadata->ToString()); + + // A second batch and its acknowledgement prove the channel survives a round. + ASSERT_OK(writer.WriteRecordBatch(*batch)); + ASSERT_OK(reader.ReadMetadata(&metadata)); + ASSERT_NE(nullptr, metadata); + ASSERT_EQ("ack 2 rows", metadata->ToString()); + + ASSERT_OK(writer.Close()); + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + auto listener = flight_server.WaitForListener(); + ASSERT_NE(listener, nullptr) << "the server never handed out a DoPut listener"; + ASSERT_EQ(2, listener->wrote()) << "every batch must have been acknowledged"; + ASSERT_OK(listener->last_write()) << listener->last_write(); +} + +TEST(AsyncGrpcTest, DoPutCancelFromServerSideFinishesTheUpload) { + // The server aborts the upload without waiting for the client: Cancel() + // finishes the RPC with the given status, the client's write or Close + // reports it, and the listener hears it as its OnFinish status. + TestUploadServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = arrow::schema({arrow::field("a", arrow::int64())}); + auto batch = arrow::RecordBatch::Make(schema, 2, + {arrow::ArrayFromJSON(arrow::int64(), "[3, 4]")}); + + // Open the upload and write one batch, but do not end it. + auto put = client->DoPut(FlightDescriptor::Path({"cancel"}), schema); + ASSERT_TRUE(put.ok()) << put.status(); + auto& writer = *put->writer; + ASSERT_OK(writer.WriteRecordBatch(*batch)); + + // Wait for the listener the virtual handed out, then cancel it from this + // (non-gRPC) thread: that is the mid-upload server-side abort. + auto listener = flight_server.WaitForListener(0); + ASSERT_NE(listener, nullptr) << "the server never handed out a DoPut listener"; + listener->CancelWith(arrow::Status::Cancelled("server said stop")); + + // The client sees the cancellation on the surface that observes the RPC. + const Status close_status = writer.Close(); + ASSERT_FALSE(close_status.ok()) << "the server-side cancel never reached the client"; + ASSERT_EQ(close_status.code(), arrow::StatusCode::Cancelled); + ASSERT_NE(close_status.message().find("server said stop"), std::string::npos) + << "unexpected status: " << close_status; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // Cancel() reached the transport, and the listener was told the upload ended + // with that status exactly once. + ASSERT_EQ(listener->cancel_count(), 1); + ASSERT_OK(listener->cancel_status()) << listener->cancel_status(); + ASSERT_EQ(listener->finish_count(), 1); + // OnFinish must carry the ending that caused it, not a generic status. + ASSERT_EQ(listener->finish_status().code(), arrow::StatusCode::Cancelled) + << "OnFinish got: " << listener->finish_status(); + + // After the upload, Cancel() has nothing to cancel. + listener->CancelWith(arrow::Status::Cancelled("too late")); + ASSERT_EQ(listener->cancel_status().code(), arrow::StatusCode::Invalid); +} + +/// A listener stuck inside OnNext(): the upload is parked in the application +/// while the RPC underneath it dies. +class HangingDoPutListener final : public FinishRecordingListener { + public: + arrow::Future<> OnNext(FlightStreamChunk) override { + next_awaits_.fetch_add(1); + return arrow::Future<>::Make(); + } + + std::atomic next_awaits_{0}; +}; + +/// A DoPut server that hands out a hanging listener, one per upload. +using HangingDoPutServer = RecordingUploadServer; + +TEST(AsyncGrpcTest, DoPutClientDeadlineWithPendingHandlerFinishes) { + // The client's deadline expires while the listener is still processing a + // chunk (its OnNext() future never resolves): the transport must finish the + // RPC itself instead of waiting for the listener: Shutdown()/Wait() below + // would otherwise block, since gRPC keeps the call alive until Finish() + // runs. + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + HangingDoPutServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = OneColumnSchema(); + auto batch = OneColumnBatch(schema); + + FlightCallOptions impatient; + impatient.timeout = std::chrono::milliseconds(1000); + auto put = client->DoPut(impatient, FlightDescriptor::Path({"hang"}), schema); + ASSERT_TRUE(put.ok()) << put.status(); + (void)put->writer->WriteRecordBatch(*batch); + + // From here the listener is inside OnNext() and the deadline does the rest. + ASSERT_TRUE(WaitFor([&flight_server] { + auto listener = flight_server.WaitForListener(0); + return listener != nullptr && listener->next_awaits_.load() >= 1; + })) << "the listener never saw a chunk"; + + ARROW_UNUSED(put->writer->Close()); + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // The upload ended anyway, and the listener heard about it exactly once. + ASSERT_EQ(1, flight_server.listener(0)->finish_count()); +} + +TEST(AsyncGrpcTest, DoPutClientDeadlineWhileIdleIsNotACleanFinish) { + // The deadline expires while the server is idle (the batch was consumed and + // a read is armed): the upload was cut short by the client, so OnFinish() + // must not report it as a clean end, only the acknowledgement write + // completing cleanly may say OK. + TestUploadServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = OneColumnSchema(); + auto batch = OneColumnBatch(schema); + + FlightCallOptions impatient; + impatient.timeout = std::chrono::milliseconds(600); + auto put = client->DoPut(impatient, FlightDescriptor::Path({"idle"}), schema); + ASSERT_TRUE(put.ok()) << put.status(); + ASSERT_OK(put->writer->WriteRecordBatch(*batch)); + + auto listener = flight_server.WaitForListener(0); + ASSERT_NE(listener, nullptr) << "the server never handed out a DoPut listener"; + ASSERT_TRUE(WaitFor([&listener] { return listener->batches().size() >= 1; })) + << "the server never consumed the batch"; + + // Wait for the deadline to end the upload instead of sleeping out its + // duration: the listener's OnFinish runs on the transport thread that ends + // the RPC, so the wait is on the observable event, not on wall-clock time. + // DoPut writes the schema eagerly, so the client's DoPut() above already + // showed the channel was alive when its deadline was armed. + ASSERT_TRUE(WaitFor([&listener] { return listener->finish_count() >= 1; })) + << "the deadline never fired; inconclusive"; + + const auto close_status = put->writer->Close(); + ASSERT_FALSE(close_status.ok()) << "the deadline never fired; inconclusive"; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ(1, listener->finish_count()); + EXPECT_FALSE(listener->finish_status().ok()) + << "a deadline-cancelled upload was reported as a clean finish: " + << listener->finish_status(); +} + +/// A listener whose OnNext() future the test resolves on demand: the upload +/// parks in the application until the late-continuation test releases it. +class DeferredDoPutListener final : public FinishRecordingListener { + public: + arrow::Future<> OnNext(FlightStreamChunk) override { + next_awaits_.fetch_add(1); + std::lock_guard lock(mutex_); + pending_ = arrow::Future<>::Make(); + return pending_; + } + + /// Resolve the parked OnNext, OK releases the decode continuation. + void ResolvePending() { + std::lock_guard lock(mutex_); + if (pending_.is_valid() && !pending_.is_finished()) { + pending_.MarkFinished(arrow::Status::OK()); + } + } + + std::atomic next_awaits_{0}; + + private: + std::mutex mutex_; + arrow::Future<> pending_; +}; + +using DeferredDoPutServer = RecordingUploadServer; + +TEST(AsyncGrpcTest, DoPutClientCancelWhileListenerParkedStaysSafe) { + // The listener's OnNext() resolves after the deadline already finished the + // RPC: the parked decode continuation must stand down (release its hold, + // never touch the finished call) rather than reach into a destroyed reactor. + // The test resolves the parked future from its own thread, and the server + // keeps its shared_ptr alive, so the resolution lands on the real listener. + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + DeferredDoPutServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = OneColumnSchema(); + auto batch = OneColumnBatch(schema); + + // The deadline is what ends this RPC (the listener stays parked, so only the + // deadline can finish it), and DoPut writes the schema eagerly, so it must be + // long enough to survive the setup on a slow machine. The wait below must + // outlast it, or the poll loop gives up at the moment the deadline fires: + // that race is what made this test fail on the macOS runner. + constexpr auto kDeadline = std::chrono::milliseconds(2000); + FlightCallOptions impatient; + impatient.timeout = kDeadline; + auto put = client->DoPut(impatient, FlightDescriptor::Path({"late"}), schema); + ASSERT_TRUE(put.ok()) << put.status(); + ASSERT_OK(put->writer->WriteRecordBatch(*batch)); + + auto listener = flight_server.WaitForListener(0); + ASSERT_NE(listener, nullptr) << "the server never handed out a DoPut listener"; + ASSERT_TRUE(WaitFor([&listener] { return listener->next_awaits_.load() >= 1; })) + << "the listener never saw a chunk"; + + // Walk away: the client stops reading, so the server's write has nowhere to + // go and the RPC ends while the listener is still parked. Wait for that to + // have happened rather than sleeping blindly, and fail clearly instead of + // reporting a vacuous pass if it did not. + ASSERT_OK(client->Close()); + ASSERT_TRUE(WaitFor([&listener] { return listener->finish_count() >= 1; }, + kDeadline + std::chrono::seconds(10))) + << "the RPC never ended; the listener is not the parked one from this upload"; + + // Now resolve the parked listener: the continuation runs on a finished RPC, + // which is the case the guards and holds are there to make safe. + listener->ResolvePending(); + + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // The upload had ended once, and the late resolution added no second report. + ASSERT_EQ(1, listener->finish_count()); + + // Closing the writer after the RPC is gone must report a failure, not a + // silent success. + ASSERT_FALSE(put->writer->Close().ok()); +} + +TEST(AsyncGrpcTest, DoPutWithoutDescriptorOnFirstMessageIsRejected) { + // A client that starts an upload without a descriptor must be refused with + // the same status as the sync transport ("Descriptor missing on first + // message"): the listener must not hear a descriptor, and the upload must not + // be acknowledged as well formed. + TestFlightServer flight_server; + auto listener = std::make_shared(); + flight_server.set_listener(listener); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_GT(flight_server.port(), 0); + + // Hand-build the first FlightData of a DoPut: the schema IPC message with no + // descriptor field at all. + auto schema = OneColumnSchema(); + auto batch = OneColumnBatch(schema); + auto reader = arrow::RecordBatchReader::Make({batch}, schema).ValueOrDie(); + RecordBatchStream stream(std::move(reader)); + ASSERT_OK_AND_ASSIGN(auto schema_payload, stream.GetSchemaPayload()); + schema_payload.descriptor = nullptr; + + ::grpc::ByteBuffer request; + bool own_buffer = false; + ASSERT_TRUE(FlightDataSerialize(schema_payload, &request, &own_buffer).ok()); + + // Drive the bidi call through the raw generic stub: the regular FlightClient + // always sets the descriptor, which is exactly what this case lacks. + std::shared_ptr<::grpc::Channel> channel = + ::grpc::CreateChannel("127.0.0.1:" + std::to_string(flight_server.port()), + ::grpc::InsecureChannelCredentials()); + ::grpc::GenericStub stub(channel); + ::grpc::ClientContext context; + // Bound the call: were the descriptor ever accepted, the client would wait + // forever for an acknowledgement, turning a logic error into a hung suite. + context.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(10)); + ::grpc::CompletionQueue cq; + void* tag = nullptr; + bool ok = false; + + auto call = stub.Call(&context, "/arrow.flight.protocol.FlightService/DoPut", &cq, + reinterpret_cast(1)); + ASSERT_NE(call, nullptr); + ASSERT_TRUE(cq.Next(&tag, &ok)) << "the call never started"; + + call->Write(request, reinterpret_cast(2)); + ASSERT_TRUE(cq.Next(&tag, &ok)) << "the message was never written"; + + call->WritesDone(reinterpret_cast(3)); + ASSERT_TRUE(cq.Next(&tag, &ok)) << "WritesDone never completed"; + + ::grpc::Status status; + call->Finish(&status, reinterpret_cast(4)); + ASSERT_TRUE(cq.Next(&tag, &ok)) << "the call never finished"; + + EXPECT_FALSE(status.ok()) << "an upload without a descriptor must be refused"; + EXPECT_THAT(status.error_message(), ::testing::HasSubstr("Descriptor missing")); + + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // The listener never saw a descriptor, and never saw the message as valid. + EXPECT_EQ(0, listener->descriptor_count()); + EXPECT_EQ(0, listener->batches().size()); +} + +TEST(AsyncGrpcTest, DoPutReportsAFailingOnFinishToTheClient) { + // The listener's OnFinish() is where an application commits the upload. When + // it fails asynchronously, the client must not see a successful Close(): the + // RPC's status is chained on that future instead of being reported OK as soon + // as the upload's data arrived. + TestFlightServer flight_server; + auto listener = std::make_shared(); + flight_server.set_listener(listener); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = OneColumnSchema(); + auto batch = OneColumnBatch(schema); + auto put = client->DoPut(FlightDescriptor::Path({"fail-finish"}), schema); + ASSERT_TRUE(put.ok()) << put.status(); + ASSERT_OK(put->writer->WriteRecordBatch(*batch)); + + // Close() acknowledges the ending and then sees the RPC's final status, which + // is the listener's failure. + const auto close_status = put->writer->Close(); + EXPECT_FALSE(close_status.ok()) + << "a failing OnFinish must not be reported to the client as success"; + EXPECT_THAT(close_status.ToString(), ::testing::HasSubstr("could not be committed")); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + EXPECT_EQ(1, listener->finish_count()); +} + +TEST(AsyncGrpcTest, DoPutHalfCloseWithoutAnyMessageIsRejected) { + // A client that half-closes without sending a single message must not be + // acknowledged: there is no descriptor, so there is no upload. The sync + // transport rejects this ("Stream finished before first message sent"); + // without the check the reactor sent the clean-ending ack and the listener + // heard OnFinish(OK) for a request that was never well formed. + TestFlightServer flight_server; + auto listener = std::make_shared(); + flight_server.set_listener(listener); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_GT(flight_server.port(), 0); + + std::shared_ptr<::grpc::Channel> channel = + ::grpc::CreateChannel("127.0.0.1:" + std::to_string(flight_server.port()), + ::grpc::InsecureChannelCredentials()); + ::grpc::GenericStub stub(channel); + ::grpc::ClientContext context; + context.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(10)); + ::grpc::CompletionQueue cq; + void* tag = nullptr; + bool ok = false; + + auto call = stub.Call(&context, "/arrow.flight.protocol.FlightService/DoPut", &cq, + reinterpret_cast(1)); + ASSERT_NE(call, nullptr); + ASSERT_TRUE(cq.Next(&tag, &ok)) << "the call never started"; + + // Half-close with no message written at all. + call->WritesDone(reinterpret_cast(2)); + ASSERT_TRUE(cq.Next(&tag, &ok)) << "WritesDone never completed"; + + ::grpc::Status status; + call->Finish(&status, reinterpret_cast(3)); + ASSERT_TRUE(cq.Next(&tag, &ok)) << "the call never finished"; + + EXPECT_FALSE(status.ok()) << "an upload with no message at all must be refused"; + EXPECT_THAT(status.error_message(), ::testing::HasSubstr("empty upload")); + + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // The listener must not have been told the upload ended cleanly. + EXPECT_EQ(0, listener->descriptor_count()); + EXPECT_EQ(0, listener->batches().size()); +} + +TEST(AsyncGrpcTest, UseAsyncGrpcFlagUploadRejectedOnDescriptor) { + // A handler that rejects the descriptor rejects the whole upload through + // the same surface as an OnNext rejection: the writer's status. + TestFlightServer flight_server; + auto listener = std::make_shared(); + listener->set_descriptor_status( + arrow::Status::Invalid("listener rejected this descriptor")); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + flight_server.set_listener(listener); + ASSERT_OK(flight_server.Init(options)); + ASSERT_GT(flight_server.port(), 0); + + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = TwoColumnSchema(); + auto batch = TwoColumnBatch(schema); + + const FlightDescriptor descriptor = FlightDescriptor::Path({"rejected-h2"}); + auto [surface, status] = UploadOneBatch(client.get(), descriptor, schema, batch); + ASSERT_FALSE(status.ok()) + << "the descriptor rejection never reached the client, surfaced on " << surface; + ASSERT_EQ(status.code(), arrow::StatusCode::Invalid); + ASSERT_NE(status.message().find("listener rejected this descriptor"), std::string::npos) + << "unexpected status: " << status; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // The upload was rejected on its descriptor: the handler saw it, and no + // schema or data followed it. + ASSERT_EQ(listener->descriptor_count(), 1); + ASSERT_EQ(listener->descriptors()[0], descriptor); + ASSERT_EQ(listener->schema_count(), 0); + ASSERT_EQ(listener->batches().size(), 0); +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange.cc new file mode 100644 index 000000000000..16ae8f3ac723 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange.cc @@ -0,0 +1,123 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_exchange_internal.h" + +namespace arrow::flight::transport::grpc::detail { + +ExchangeReactor::ExchangeReactor(AsyncCallContext flight_context, + AsyncGenericFlightServerBase* base) + : AsyncReactorBase(std::move(flight_context)), + base_(base), + reader_(std::make_shared()), + writer_(std::make_shared(this)) { + reader_->SetReactor(this); + // The first message carries the request descriptor. + StartRead(&read_buf_); +} + +void ExchangeReactor::OnReadDone(bool ok) { + if (finished()) { + return; + } + if (!ok) { + if (!started_) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, "Failed to read request")); + return; + } + reader_->OnReadClosed(); + return; + } + const ::grpc::Status deserialized = FlightDataDeserialize(&read_buf_, &read_data_); + if (!deserialized.ok()) { + if (!started_) { + FinishOnce( + MakeFlightError(FlightStatusCode::Internal, deserialized.error_message())); + return; + } + reader_->OnReadFailed( + MakeFlightError(FlightStatusCode::Internal, deserialized.error_message())); + return; + } + if (!started_) { + started_ = true; + if (read_data_.descriptor == nullptr) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, + "Descriptor missing on first message")); + return; + } + reader_->SetDescriptor(*read_data_.descriptor); + // The first message may carry data,the reader consumes it on its first + // NextAsync, so hand it over before the handler can ask. + reader_->OfferFirstRead(std::move(read_data_)); + arrow::Future<> exchange = base_->DoExchangeAsync(flight_context(), reader_, writer_); + exchange.AddCallback([this, token = hold()](const arrow::Status& status) { + if (!finished()) { + FinishOnce(status); + } + }); + return; + } + // A message read because the reader asked for one. + reader_->OnMessage(std::move(read_data_)); +} + +void ExchangeReactor::OnWriteDone(bool ok) { + if (finished()) { + return; + } + writer_->OnWriteDone(ok); +} + +void ExchangeReactor::OnCancel() { + // Finish first (the handler's future may never resolve), then unwind the + // reader and writer so a handler parked on their futures wakes up. + FinishOnce(arrow::Status::Cancelled()); + reader_->OnCancelled(); + writer_->OnCancelled(); +} + +void ExchangeReactor::OnDone() { + reader_->DetachReactor(); + writer_->DetachReactor(); + AsyncReactorBase::OnDone(); +} + +void ExchangeReactor::StartReadNext() { + WriteIfLive([this] { StartRead(&read_buf_); }); +} + +arrow::Status ExchangeReactor::WritePayload(FlightPayload payload) { + bool own_buffer = false; + const ::grpc::Status serialize = FlightDataSerialize(payload, &write_buf_, &own_buffer); + (void)own_buffer; + if (!serialize.ok()) { + return MakeFlightError(FlightStatusCode::Internal, serialize.error_message()); + } + const bool written = WriteIfLive([this] { StartWrite(&write_buf_); }); + if (!written) { + return arrow::Status::Cancelled("the exchange is over"); + } + return arrow::Status::OK(); +} + +::grpc::ServerGenericBidiReactor* MakeExchangeReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) { + return new ExchangeReactor(std::move(flight_context), base); +} + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_internal.h b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_internal.h new file mode 100644 index 000000000000..e64edf4d4e1f --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_internal.h @@ -0,0 +1,218 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/flight/transport/grpc/async_grpc_service_internal.h" +#include "arrow/ipc/writer.h" + +namespace arrow::flight::transport::grpc::detail { + +// DoExchange: the bidirectional RPC. +// +// Its handler is asynchronous user code: it receives an +// AsyncFlightMessageReader/AsyncFlightMessageWriter pair and returns a future +// the reactor finishes the RPC with. The three classes below are mutually +// coupled on purpose -- the reader asks the reactor for the next read, the +// writer asks it to write a payload, and the reactor drives both from the +// gRPC callbacks -- so they are declared together here and implemented one +// per translation unit (async_grpc_exchange_reader.cc, +// async_grpc_exchange_writer.cc, async_grpc_exchange_reactor.cc). + +class ExchangeReactor; + +/// \brief The AsyncFlightMessageReader of one DoExchange RPC. +/// +/// Decoding reuses the push decoder DoPut uses, so the wire shapes behave +/// identically: the schema and dictionary messages are read through (they are +/// not chunks), a metadata-only message is a chunk with a null data member, +/// and the end of the exchange is the chunk whose members are both null. +/// Messages are pulled one demand at a time: NextAsync() arms the next read, +/// and the chunk (or the end, or the failure) resolves the demand. +/// +/// The reactor drives the I/O (it owns the read buffer and calls OnMessage / +/// OnReadClosed / OnReadFailed), this class owns the decode state machine. +class ExchangeReader final : public AsyncFlightMessageReader, + public std::enable_shared_from_this { + public: + ExchangeReader(); + + const FlightDescriptor& descriptor() const override; + + arrow::Future NextAsync() override; + + /// \brief The reactor wires itself before the handler starts. + void SetReactor(ExchangeReactor* reactor); + /// \brief Called by the reactor once the exchange is over: from then on the + /// handler's reader must refuse to start reads instead of calling into a + /// freed reactor. + void DetachReactor(); + /// \brief The reactor sets the descriptor from the first message, before the + /// handler starts. + void SetDescriptor(const FlightDescriptor& descriptor); + /// \brief The reactor hands over the already-read first message, the reader + /// consumes it on its first NextAsync(). + void OfferFirstRead(internal::FlightData data); + + /// \brief The reactor's OnReadDone: a message was read. + void OnMessage(internal::FlightData data); + /// \brief The reactor's OnReadDone with ok == false: the read side is over. + void OnReadClosed(); + /// \brief The reactor read a message that could not be deserialized. + void OnReadFailed(Status status); + /// \brief The RPC was cancelled: fail the demand so the handler unwinds. + void OnCancelled(); + + private: + /// The decoder's listener: chunks are handed to the reader. + class ChunkListener; + + void ResolvePending(FlightStreamChunk chunk); + void ResolvePending(const Status& status); + + /// \brief Decode one message and resolve the demand it satisfies, or keep + /// reading when the message was read through (a schema or dictionary + /// message). + void Consume(internal::FlightData data); + + /// \brief Arm the next read through the reactor, defined in + /// async_grpc_exchange_reader.cc, which has the reactor's definition. + void RequestRead(); + + std::atomic reactor_{nullptr}; + std::shared_ptr listener_; + AsyncFlightMessageDecoder decoder_; + FlightDescriptor descriptor_; + std::optional first_read_; + std::optional buffered_; + arrow::Future pending_; + Status read_error_; + bool read_in_flight_ = false; + bool read_closed_ = false; + std::mutex mutex_; +}; + +/// \brief The AsyncFlightMessageWriter of one DoExchange RPC. +/// +/// Record batches are serialized through the same IPC writer stack the +/// synchronous transport uses (dictionary messages included) into an +/// IpcPayloadWriter sink that queues FlightPayloads, the reactor then writes +/// them one at a time. Each method's future resolves when its messages have +/// been handed to the reactor, or with the first failure. +class ExchangeWriter final : public AsyncFlightMessageWriter { + public: + explicit ExchangeWriter(ExchangeReactor* reactor); + + /// \brief Called by the reactor once the exchange is over: from then on the + /// handler's writer must not call into a freed reactor. + void DetachReactor(); + + arrow::Future<> BeginAsync(std::shared_ptr schema) override; + arrow::Future<> WriteRecordBatchAsync(const RecordBatch& batch) override; + arrow::Future<> WriteWithMetadataAsync(const RecordBatch& batch, + std::shared_ptr app_metadata) override; + arrow::Future<> WriteMetadataAsync(std::shared_ptr app_metadata) override; + arrow::Future<> CloseAsync() override; + + /// \brief The reactor's OnWriteDone: continue flushing the queue. + void OnWriteDone(bool ok); + /// \brief The RPC was cancelled: fail the pending write so the handler + /// unwinds. + void OnCancelled(); + + private: + /// The IpcPayloadWriter that queues the IPC payloads the writer stack + /// produces, the schema, dictionary and record-batch messages all travel + /// this way, as on the synchronous path. + class PayloadSink; + + arrow::Future<> WriteBatch(const RecordBatch& batch); + + /// \brief Write the queued payloads one at a time, the returned future + /// resolves when the queue is empty, or with the first failure. + arrow::Future<> Flush(); + + /// \brief Write queued payloads through the reactor, defined in + /// async_grpc_exchange_writer.cc, which has the reactor's definition. + void Pump(); + + void FailFlush(Status status); + + std::atomic reactor_; + std::unique_ptr batch_writer_; + ipc::IpcWriteOptions options_ = ipc::IpcWriteOptions::Defaults(); + std::shared_ptr pending_app_metadata_; + std::deque queue_; + arrow::Future<> flush_; + bool write_in_flight_ = false; + std::mutex mutex_; +}; + +/// \brief Serve one DoExchange RPC. +/// +/// The handler receives an AsyncFlightMessageReader/Writer pair and returns a +/// future, the RPC then runs entirely on continuations. The reader pulls one +/// message per NextAsync() demand (the reactor arms the read), the writer +/// serializes through the IPC writer stack and the reactor writes one message +/// per StartWrite, and the handler's future completing finishes the RPC with +/// its status. No callback thread is ever held by the handler. +class ExchangeReactor final : public AsyncReactorBase { + public: + ExchangeReactor(AsyncCallContext flight_context, AsyncGenericFlightServerBase* base); + + void OnReadDone(bool ok) override; + void OnWriteDone(bool ok) override; + void OnCancel() override; + void OnDone() override; + + /// \brief The reader's demand: arm the next read. The demand guarantees at + /// most one read is outstanding (ExchangeReader::RequestRead is the only + /// caller). + void StartReadNext(); + + /// \brief Serialize one payload into the write buffer and start the write. + /// Called by ExchangeWriter, OnWriteDone() resumes it. + arrow::Status WritePayload(FlightPayload payload); + + /// \brief Take a reference across a call made through a back-pointer. + /// + /// OnDone() detaches the reader/writer back-pointers and then releases the + /// reactor's last reference, so a reader/writer that loads the raw pointer and + /// calls in can reach freed memory. These re-expose the base's protected + /// Hold()/ReleaseHold() publicly (the same way DoPutReactor does for its + /// transport interface); callers pair them around the call. + void Hold() { AsyncReactorBase::Hold(); } + void ReleaseHold() { AsyncReactorBase::ReleaseHold(); } + + private: + AsyncGenericFlightServerBase* base_; + std::shared_ptr reader_; + std::shared_ptr writer_; + ::grpc::ByteBuffer read_buf_; + ::grpc::ByteBuffer write_buf_; + internal::FlightData read_data_; + bool started_ = false; +}; + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_reader.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_reader.cc new file mode 100644 index 000000000000..fda5dd3a6f2c --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_reader.cc @@ -0,0 +1,229 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_exchange_internal.h" + +namespace arrow::flight::transport::grpc::detail { + +/// The decoder's listener: chunks are handed to the reader, the descriptor and +/// the schema are consumed (the reactor sets the descriptor before the handler +/// starts). +class ExchangeReader::ChunkListener final : public AsyncFlightDataListener { + public: + explicit ChunkListener(ExchangeReader* reader) : reader_(reader) {} + + arrow::Future<> OnNext(FlightStreamChunk chunk) override { + bool has_demand; + { + std::lock_guard lock(reader_->mutex_); + has_demand = reader_->pending_.is_valid() && !reader_->pending_.is_finished(); + if (!has_demand) { + reader_->buffered_ = std::move(chunk); + } + } + if (has_demand) { + reader_->ResolvePending(std::move(chunk)); + } + return arrow::Future<>::MakeFinished(); + } + + arrow::Future<> OnDescriptor(const FlightDescriptor&) override { + return arrow::Future<>::MakeFinished(); + } + + arrow::Status OnSchemaDecoded(std::shared_ptr) override { + return arrow::Status::OK(); + } + + private: + ExchangeReader* reader_; +}; + +ExchangeReader::ExchangeReader() + : listener_(std::make_shared(this)), decoder_(listener_) {} + +const FlightDescriptor& ExchangeReader::descriptor() const { return descriptor_; } + +arrow::Future ExchangeReader::NextAsync() { + using FSCFuture = arrow::Future; + + arrow::Future returned; + { + std::lock_guard lock(mutex_); + if (pending_.is_valid() && !pending_.is_finished()) { + return FSCFuture::MakeFinished(arrow::Status::Invalid("one NextAsync at a time")); + } + if (buffered_.has_value()) { + auto out = FSCFuture::MakeFinished(std::move(*buffered_)); + buffered_.reset(); + return out; + } + + if (!read_error_.ok()) { + return FSCFuture::MakeFinished(read_error_); + } + + if (read_closed_ || reactor_.load() == nullptr) { + // End of exchange: the reactor either saw the read side close or the RPC + // is over (and the reactor is gone), so no further read can be started. + return FSCFuture::MakeFinished(FlightStreamChunk{}); + } + + pending_ = FSCFuture::Make(); + returned = pending_; + } + // Consume()/RequestRead() take the lock themselves and may resolve pending_ + // from a callback, so they must run unlocked. + if (first_read_.has_value()) { + internal::FlightData data; + { + std::lock_guard lock(mutex_); + data = std::move(*first_read_); + first_read_.reset(); + } + Consume(std::move(data)); + } else { + RequestRead(); + } + return returned; +} + +void ExchangeReader::SetReactor(ExchangeReactor* reactor) { + std::lock_guard lock(mutex_); + reactor_.store(reactor); +} + +void ExchangeReader::DetachReactor() { + std::lock_guard lock(mutex_); + reactor_.store(nullptr); +} + +void ExchangeReader::SetDescriptor(const FlightDescriptor& descriptor) { + descriptor_ = descriptor; +} + +void ExchangeReader::OfferFirstRead(internal::FlightData data) { + // Set during reactor setup, but NextAsync() reads it under mutex_. + std::lock_guard lock(mutex_); + first_read_ = std::move(data); +} + +void ExchangeReader::OnMessage(internal::FlightData data) { + { + std::lock_guard lock(mutex_); + read_in_flight_ = false; + } + Consume(std::move(data)); +} + +void ExchangeReader::OnReadClosed() { + { + std::lock_guard lock(mutex_); + read_in_flight_ = false; + read_closed_ = true; + } + ResolvePending(FlightStreamChunk{}); +} + +void ExchangeReader::OnReadFailed(Status status) { + { + std::lock_guard lock(mutex_); + read_in_flight_ = false; + read_error_ = status; + } + ResolvePending(status); +} + +void ExchangeReader::OnCancelled() { + { + std::lock_guard lock(mutex_); + read_closed_ = true; + } + ResolvePending(arrow::Status::Cancelled("the client cancelled the exchange")); +} + +void ExchangeReader::ResolvePending(FlightStreamChunk chunk) { + arrow::Future pending; + { + std::lock_guard lock(mutex_); + if (!pending_.is_valid() || pending_.is_finished()) { + return; + } + pending = std::move(pending_); + pending_ = arrow::Future(); + } + pending.MarkFinished(std::move(chunk)); +} + +void ExchangeReader::ResolvePending(const Status& status) { + arrow::Future pending; + { + std::lock_guard lock(mutex_); + if (!pending_.is_valid() || pending_.is_finished()) { + return; + } + pending = std::move(pending_); + pending_ = arrow::Future(); + } + pending.MarkFinished(status); +} + +void ExchangeReader::Consume(internal::FlightData data) { + { + std::lock_guard lock(mutex_); + if (read_closed_) { + return; + } + } + auto self = shared_from_this(); + decoder_.Consume(std::move(data)) + .AddCallback([this, self](const arrow::Status& status) { + if (!status.ok()) { + { + std::lock_guard lock(mutex_); + read_error_ = status; + } + ResolvePending(status); + } else { + bool has_demand; + { + std::lock_guard lock(mutex_); + has_demand = pending_.is_valid() && !pending_.is_finished(); + } + if (has_demand) { + RequestRead(); + } + } + }); +} + +void ExchangeReader::RequestRead() { + ExchangeReactor* reactor = nullptr; + { + std::lock_guard lock(mutex_); + reactor = reactor_.load(); + if (read_in_flight_ || read_closed_ || reactor == nullptr) { + return; + } + reactor->Hold(); + read_in_flight_ = true; + } + reactor->StartReadNext(); + reactor->ReleaseHold(); +} + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_test.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_test.cc new file mode 100644 index 000000000000..f1480cea72e5 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_test.cc @@ -0,0 +1,206 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// DoExchange over the async generic service: the bidirectional RPC, whose handler +// awaits the reader's demand and each write in turn, so reads and writes are +// interleaved without blocking a callback thread. Covers the echo path, a +// metadata-only exchange, the client vanishing mid-exchange, and a handler whose +// future never resolves. + +#include "arrow/flight/transport/grpc/async_grpc_test_fixtures.h" + +namespace arrow::flight::transport::grpc { + +// --------------------------------------------------------------------------- +// DoExchange: the bidirectional RPC. Its handler is DoExchangeAsync, which +// awaits the reader's demand and each write before the next one, so reads and +// writes are interleaved without blocking the callback thread. +// --------------------------------------------------------------------------- + +TEST(AsyncGrpcTest, DoExchangeEchoesBatches) { + // Two batches up, two batches back, in order: this is the case that needs a + // read and a write outstanding at the same time. + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + EchoExchangeTestServer flight_server; + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN( + auto client_location, + Location::Parse("grpc://localhost:" + std::to_string(flight_server.port()))); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(client_location)); + + ASSERT_OK_AND_ASSIGN(auto exchange, + client->DoExchange(FlightDescriptor::Command("echo"))); + + auto schema = arrow::schema({arrow::field("a", arrow::int64())}); + std::vector> sent = { + arrow::RecordBatch::Make(schema, 2, + {arrow::ArrayFromJSON(arrow::int64(), "[1, 2]")}), + arrow::RecordBatch::Make(schema, 3, + {arrow::ArrayFromJSON(arrow::int64(), "[3, 4, 5]")})}; + ASSERT_OK(exchange.writer->Begin(schema)); + for (const auto& batch : sent) { + ASSERT_OK(exchange.writer->WriteRecordBatch(*batch)); + } + ASSERT_OK(exchange.writer->DoneWriting()); + + ASSERT_OK_AND_ASSIGN(auto server_schema, exchange.reader->GetSchema()); + AssertSchemaEqual(*schema, *server_schema); + for (const auto& batch : sent) { + ASSERT_OK_AND_ASSIGN(auto chunk, exchange.reader->Next()); + ASSERT_NE(nullptr, chunk.data); + ASSERT_BATCHES_EQUAL(*batch, *chunk.data); + } + ASSERT_OK_AND_ASSIGN(auto end, exchange.reader->Next()); + ASSERT_EQ(nullptr, end.data); + ASSERT_OK(exchange.writer->Close()); + + // Cleanup before looking at the recording: it was written on a gRPC thread. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ(1, flight_server.exchange_count()); + ASSERT_EQ(0, flight_server.exchanges_in_flight()); +} + +TEST(AsyncGrpcTest, DoExchangeMetadataOnlyDoesNotHang) { + // A client that writes metadata and nothing else, then reads: the server's + // handler must see the end of the exchange and reply, not wait for a schema + // message that will never come (the case DataTest::TestDoExchangeNoData + // covers). A hang here is the failure mode this test exists for. + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + EchoExchangeTestServer flight_server; + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN( + auto client_location, + Location::Parse("grpc://localhost:" + std::to_string(flight_server.port()))); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(client_location)); + + ASSERT_OK_AND_ASSIGN(auto exchange, + client->DoExchange(FlightDescriptor::Command("echo"))); + + ASSERT_OK(exchange.writer->WriteMetadata(arrow::Buffer::FromString("only-metadata"))); + ASSERT_OK(exchange.writer->DoneWriting()); + + // The echo server reflects the metadata chunk back with no data. + ASSERT_OK_AND_ASSIGN(auto chunk, exchange.reader->Next()); + ASSERT_EQ(nullptr, chunk.data); + ASSERT_NE(nullptr, chunk.app_metadata); + ASSERT_EQ("only-metadata", chunk.app_metadata->ToString()); + ASSERT_OK_AND_ASSIGN(auto end, exchange.reader->Next()); + ASSERT_EQ(nullptr, end.data); + ASSERT_OK(exchange.writer->Close()); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ(1, flight_server.exchange_count()); + ASSERT_EQ(0, flight_server.exchanges_in_flight()); +} + +TEST(AsyncGrpcTest, DoExchangeClientDisconnectFinishes) { + // A client that vanishes mid-exchange: the server's read fails, the reactor + // must let the handler unwind and finish the RPC, and the exchange counter + // added to the test server must come back to zero -- a hang or a leaked + // handler leaves it above zero. + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + EchoExchangeTestServer flight_server; + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN( + auto client_location, + Location::Parse("grpc://localhost:" + std::to_string(flight_server.port()))); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(client_location)); + + ASSERT_OK_AND_ASSIGN(auto exchange, + client->DoExchange(FlightDescriptor::Command("echo"))); + + auto schema = arrow::schema({arrow::field("a", arrow::int64())}); + ASSERT_OK(exchange.writer->Begin(schema)); + ASSERT_OK(exchange.writer->WriteRecordBatch(*arrow::RecordBatch::Make( + schema, 2, {arrow::ArrayFromJSON(arrow::int64(), "[1, 2]")}))); + // The server has entered its handler by now (it answered the schema), the + // exchange is live. + ASSERT_TRUE(WaitFor([&flight_server] { return flight_server.exchange_count() == 1; })); + + // Walk away without finishing: no DoneWriting, no Close, no reader drain. + // The client's destructor closes the channel, which cancels the RPC. + exchange.reader->Cancel(); + exchange.writer.reset(); + exchange.reader.reset(); + ASSERT_OK(client->Close()); + + // The handler must leave, whichever status path it takes. + const bool drained = + WaitFor([&flight_server] { return flight_server.exchanges_in_flight() == 0; }); + ASSERT_TRUE(drained) << "the exchange handler never left (in flight: " + << flight_server.exchanges_in_flight() << ")"; + + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +// A DoExchange handler whose future never resolves. +class HangingExchangeServer final : public AsyncGenericFlightServerBase { + public: + arrow::Future<> DoExchangeAsync(const ServerCallContext&, + std::shared_ptr, + std::shared_ptr) override { + handler_started.fetch_add(1); + return arrow::Future<>::Make(); // never resolved + } + + std::atomic handler_started{0}; +}; + +TEST(AsyncGrpcTest, DoExchangeClientCancelWithPendingHandlerFinishes) { + // The client walks away with the handler still parked: the transport must + // finish the RPC (and unwind reader/writer), not wait for the handler. + const int64_t before = detail::AsyncReactorBase::live_reactors(); + + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + HangingExchangeServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + ASSERT_OK_AND_ASSIGN(auto exchange, + client->DoExchange(FlightDescriptor::Command("hang"))); + // Wait until the handler is in (the descriptor reached it), so the cancel + // lands on a live handler rather than a race with dispatch. + ASSERT_TRUE(WaitFor([&flight_server] { + return flight_server.handler_started.load() >= 1; + })) << "the exchange handler never started"; + + exchange.reader->Cancel(); + exchange.writer.reset(); + exchange.reader.reset(); + ASSERT_OK(client->Close()); + + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + EXPECT_TRUE(WaitFor([&] { + return detail::AsyncReactorBase::live_reactors() == before; + })) << "a reactor leaked: still alive after the server stopped (before=" + << before << ", now=" << detail::AsyncReactorBase::live_reactors() << ")"; +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_writer.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_writer.cc new file mode 100644 index 000000000000..ce9ccf04caff --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_writer.cc @@ -0,0 +1,234 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_exchange_internal.h" + +namespace arrow::flight::transport::grpc::detail { + +/// The IpcPayloadWriter that queues the IPC payloads the writer stack +/// produces, the schema, dictionary and record-batch messages all travel this +/// way, as on the synchronous path. +class ExchangeWriter::PayloadSink final : public ipc::internal::IpcPayloadWriter { + public: + explicit PayloadSink(ExchangeWriter* writer) : writer_(writer) {} + + arrow::Status Start() override { return arrow::Status::OK(); } + + arrow::Status WritePayload(const ipc::IpcPayload& ipc_payload) override { + FlightPayload payload; + payload.ipc_message = ipc_payload; + if (ipc_payload.type == ipc::MessageType::RECORD_BATCH && + writer_->pending_app_metadata_ != nullptr) { + payload.app_metadata = std::move(writer_->pending_app_metadata_); + } + writer_->queue_.push_back(std::move(payload)); + return arrow::Status::OK(); + } + + arrow::Status Close() override { return arrow::Status::OK(); } + + private: + ExchangeWriter* writer_; +}; + +ExchangeWriter::ExchangeWriter(ExchangeReactor* reactor) : reactor_(reactor) {} + +void ExchangeWriter::DetachReactor() { + std::lock_guard lock(mutex_); + reactor_.store(nullptr); +} + +arrow::Future<> ExchangeWriter::BeginAsync(std::shared_ptr schema) { + if (reactor_.load() == nullptr) { + return arrow::Future<>::MakeFinished(arrow::Status::Cancelled( + "the exchange is over: the client disconnected or cancelled")); + } + if (batch_writer_ != nullptr) { + return arrow::Future<>::MakeFinished( + arrow::Status::Invalid("This writer has already been started.")); + } + auto sink = std::make_unique(this); + ARROW_ASSIGN_OR_RAISE(batch_writer_, ipc::internal::OpenRecordBatchWriter( + std::move(sink), schema, options_)); + return Flush(); +} + +arrow::Future<> ExchangeWriter::WriteRecordBatchAsync(const RecordBatch& batch) { + pending_app_metadata_ = nullptr; + return WriteBatch(batch); +} + +arrow::Future<> ExchangeWriter::WriteWithMetadataAsync( + const RecordBatch& batch, std::shared_ptr app_metadata) { + pending_app_metadata_ = std::move(app_metadata); + return WriteBatch(batch); +} + +arrow::Future<> ExchangeWriter::WriteMetadataAsync(std::shared_ptr app_metadata) { + FlightPayload payload; + payload.app_metadata = std::move(app_metadata); + queue_.push_back(std::move(payload)); + return Flush(); +} + +arrow::Future<> ExchangeWriter::CloseAsync() { + if (batch_writer_ != nullptr) { + auto status = batch_writer_->Close(); + batch_writer_ = nullptr; + if (!status.ok()) { + return arrow::Future<>::MakeFinished(std::move(status)); + } + } + return arrow::Future<>::MakeFinished(); +} + +void ExchangeWriter::OnWriteDone(bool ok) { + if (!ok) { + { + std::lock_guard lock(mutex_); + write_in_flight_ = false; + } + FailFlush(MakeFlightError(FlightStatusCode::Internal, "Failed to write response")); + return; + } + { + std::lock_guard lock(mutex_); + write_in_flight_ = false; + } + Pump(); +} + +void ExchangeWriter::OnCancelled() { + bool had_write; + { + std::lock_guard lock(mutex_); + had_write = write_in_flight_; + write_in_flight_ = false; + } + if (had_write) { + FailFlush(arrow::Status::Cancelled("the client cancelled the exchange")); + } +} + +arrow::Future<> ExchangeWriter::WriteBatch(const RecordBatch& batch) { + if (batch_writer_ == nullptr) { + return arrow::Future<>::MakeFinished(arrow::Status::Invalid( + "This writer is not started. Call BeginAsync() with a schema")); + } + auto status = batch_writer_->WriteRecordBatch(batch); + if (!status.ok()) { + return arrow::Future<>::MakeFinished(std::move(status)); + } + return Flush(); +} + +arrow::Future<> ExchangeWriter::Flush() { + arrow::Future<> flush; + { + std::lock_guard lock(mutex_); + if (reactor_.load() == nullptr) { + // The exchange is over: report it rather than pumping into a dead reactor. + return arrow::Future<>::MakeFinished(arrow::Status::Cancelled( + "the exchange is over: the client disconnected or cancelled")); + } + if (!flush_.is_valid() || flush_.is_finished()) { + flush_ = arrow::Future<>::Make(); + } + flush = flush_; + } + // Pump() takes the lock itself and may complete flush_ from a callback. + Pump(); + return flush; +} + +void ExchangeWriter::Pump() { + FlightPayload payload; + ExchangeReactor* reactor = nullptr; + enum class Action { kIdle, kDrained, kNoReactor, kWrite } action = Action::kIdle; + arrow::Future<> drained; + { + std::lock_guard lock(mutex_); + if (write_in_flight_) { + return; + } + reactor = reactor_.load(); + if (reactor != nullptr) { + reactor->Hold(); + } + if (queue_.empty()) { + drained = std::move(flush_); + flush_ = arrow::Future<>(); + action = Action::kDrained; + } else if (reactor == nullptr) { + write_in_flight_ = false; + action = Action::kNoReactor; + } else { + payload = std::move(queue_.front()); + queue_.pop_front(); + write_in_flight_ = true; + action = Action::kWrite; + } + } + // Everything below runs without mutex_. Every path that pinned above must + // unpin exactly once, including the ones that never reach WritePayload(). + switch (action) { + case Action::kDrained: + if (reactor != nullptr) { + reactor->ReleaseHold(); + } + if (drained.is_valid() && !drained.is_finished()) { + drained.MarkFinished(); + } + return; + case Action::kNoReactor: + // reactor is null here, so nothing was pinned. + FailFlush(arrow::Status::Cancelled( + "the exchange is over: the client disconnected or cancelled")); + return; + case Action::kWrite: { + const auto status = reactor->WritePayload(std::move(payload)); + reactor->ReleaseHold(); + if (!status.ok()) { + { + std::lock_guard lock(mutex_); + write_in_flight_ = false; + } + FailFlush(std::move(status)); + } + return; + } + case Action::kIdle: + return; + } +} + +void ExchangeWriter::FailFlush(Status status) { + // Claim the future under the lock, complete it outside: MarkFinished() runs + // continuations inline, which may re-enter the writer. + arrow::Future<> flush; + { + std::lock_guard lock(mutex_); + if (!flush_.is_valid() || flush_.is_finished()) { + return; + } + flush = std::move(flush_); + flush_ = arrow::Future<>(); + } + flush.MarkFinished(std::move(status)); +} + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_handshake.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_handshake.cc new file mode 100644 index 000000000000..c64fc3d3efea --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_handshake.cc @@ -0,0 +1,82 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_service_internal.h" + +#include +#include + +namespace arrow::flight::transport::grpc::detail { + +namespace { + +/// Serve the Handshake RPC over the generic callback API. The handler is +/// synchronous, so every path finishes the RPC from the callback that started +/// it: no callback outlives its read or write, and the single reference the base +/// class holds is enough. +class HandshakeReactor final : public AsyncReactorBase { + public: + HandshakeReactor(AsyncCallContext flight_context, HandshakeFn handshake_handler) + : AsyncReactorBase(std::move(flight_context)), + handshake_handler_(std::move(handshake_handler)) { + StartRead(&request_buf_); + } + + void OnReadDone(bool ok) override { + if (!ok) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, "Failed to read request")); + return; + } + auto request = + ParseProtoRequest(request_buf_, "HandshakeRequest"); + if (!request.ok()) { + FinishOnce(std::move(request).status()); + return; + } + std::string response; + const auto status = + handshake_handler_(flight_context(), request->payload(), &response); + if (!status.ok()) { + FinishOnce(status); + return; + } + pb::HandshakeResponse pb_response; + pb_response.set_payload(std::move(response)); + write_buf_ = MakeWriteBuffer(pb_response); + WriteIfLive([this] { StartWrite(&write_buf_); }); + } + + void OnWriteDone(bool ok) override { + FinishOnce( + ok ? arrow::Status::OK() + : MakeFlightError(FlightStatusCode::Internal, "Failed to write response")); + } + + private: + HandshakeFn handshake_handler_; + ::grpc::ByteBuffer request_buf_; + ::grpc::ByteBuffer write_buf_; +}; + +} // namespace + +::grpc::ServerGenericBidiReactor* MakeHandshakeReactor(AsyncCallContext flight_context, + HandshakeFn handshake_handler) { + return new HandshakeReactor(std::move(flight_context), std::move(handshake_handler)); +} + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_handshake_test.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_handshake_test.cc new file mode 100644 index 000000000000..a68e0abc5de9 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_handshake_test.cc @@ -0,0 +1,509 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Handshake and per-RPC token auth on the async generic service: the server +// class's Handshake()/ValidateToken() hooks, the token being required on data +// calls, payload fidelity (binary, empty), an abandoned or late handshake, and +// concurrency on fresh and reused channels. + +#include "arrow/flight/transport/grpc/async_grpc_test_fixtures.h" + +namespace arrow::flight::transport::grpc { + +// --------------------------------------------------------------------------- +// Handshake + per-RPC token auth on the async generic service. +// --------------------------------------------------------------------------- + +// Counters shared between the server callbacks and the test thread. +struct AsyncAuthState { + std::mutex mutex; + int handshakes = 0; + int validations = 0; + std::string peer_identity; +}; + +// Async-transport twin of TestServerAuthHandler: the handshake answers the +// client's password with the username, later RPCs must carry that password as +// the token, which is what TestClientAuthHandler::GetToken() puts in the +// `auth-token-bin` header (the sync twin validates it the same way in +// TestServerAuthHandler::IsValid). It is a server class: the two hooks below +// are the only configuration the transport needs. +class AsyncAuthTestServer : public AsyncGenericFlightServerBase { + public: + AsyncAuthTestServer(std::string username, std::string password, std::string ticket, + std::shared_ptr state = nullptr) + : username_(std::move(username)), + password_(std::move(password)), + ticket_(std::move(ticket)), + state_(state ? std::move(state) : std::make_shared()) {} + + const std::shared_ptr& state() const { return state_; } + + arrow::Future> DoGetAsync( + const ServerCallContext& context, const Ticket& request) override { + { + std::lock_guard guard(state_->mutex); + state_->peer_identity = context.peer_identity(); + } + if (request.ticket != ticket_) { + return arrow::Future>::MakeFinished( + Status::KeyError("No such ticket: ", request.ticket)); + } + auto schema = arrow::schema({arrow::field("a", arrow::int64())}); + auto batch = arrow::RecordBatch::Make( + schema, 2, {arrow::ArrayFromJSON(arrow::int64(), "[1, 2]")}); + ARROW_ASSIGN_OR_RAISE(auto reader, arrow::RecordBatchReader::Make({batch})); + return arrow::Future>::MakeFinished( + std::make_shared(std::move(reader))); + } + + // The hook: the transport read the client's password, answer with the + // username, or reject the RPC the way the sync twin does. + Status Handshake(const ServerCallContext& /*context*/, const std::string& password, + std::string* response) override { + { + std::lock_guard guard(state_->mutex); + state_->handshakes++; + } + if (password != password_) { + return MakeFlightError(FlightStatusCode::Unauthenticated, "Invalid token"); + } + *response = username_; + return Status::OK(); + } + + Status ValidateToken(const ServerCallContext& /*context*/, const std::string& token, + std::string* peer_identity) override { + std::lock_guard guard(state_->mutex); + state_->validations++; + if (token != password_) { + return MakeFlightError(FlightStatusCode::Unauthenticated, "Invalid token"); + } + *peer_identity = username_; + return Status::OK(); + } + + private: + std::string username_; + std::string password_; + std::string ticket_; + std::shared_ptr state_; +}; + +// A server with the async auth hooks installed, and a client connected to it. +// `status` carries the setup result, construction never fails a test. +struct AuthenticatedHarness { + AsyncAuthTestServer server{"user", "p4ssw0rd", "auth-ticket"}; + std::unique_ptr client; + Status status; + + AuthenticatedHarness() : status(Start()) {} + + ~AuthenticatedHarness() { + if (client) { + ARROW_WARN_NOT_OK(client->Close(), "Close()"); + } + ARROW_WARN_NOT_OK(server.Shutdown(), "Shutdown()"); + } + + private: + Status Start() { + ARROW_ASSIGN_OR_RAISE(auto location, Location::Parse("grpc://localhost:0")); + // The server class alone selects the async transport: deriving from it and + // overriding its virtuals are the whole configuration. + FlightServerOptions options(location); + RETURN_NOT_OK(server.Init(options)); + ARROW_ASSIGN_OR_RAISE(auto client_location, + Location::ForScheme("grpc", "127.0.0.1", server.port())); + ARROW_ASSIGN_OR_RAISE(client, FlightClient::Connect(client_location)); + return Status::OK(); + } +}; + +// A minimal server that mixes in the async hooks: the server class alone +// selects the async generic transport. +class AsyncGenericTestServer : public AsyncGenericFlightServerBase {}; + +TEST(AsyncGrpcTest, HandshakeWithoutHookIsUnimplemented) { + // A bare async server: Handshake is answered by the class's default virtual + // with UNIMPLEMENTED, and the client's Authenticate() surfaces it. + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + AsyncGenericTestServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto status = client->Authenticate( + {}, std::make_unique("user", "p4ssw0rd")); + ASSERT_FALSE(status.ok()) << "handshake must not succeed without a hook"; + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("authentication mechanism")); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, AsyncGenericFlightServerBaseImpliesAsyncTransport) { + // Deriving from the async class is enough: no hooks and no option to set, + // yet the RPC is served by the generic callback service. DoPut is the + // discriminator: an async server with no listener factory answers with its own + // message rather than "NYI". + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + AsyncGenericTestServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + // DoPut is the discriminator: an async server with no listener factory + // answers with its own message rather than "NYI". The refusal can surface on + // either call: FlightClient::DoPut + // writes the schema payload eagerly, so when the server's UNIMPLEMENTED + // arrives first the DoPut call itself fails, otherwise the writer's Close() + // reports it. + auto put = client->DoPut(FlightDescriptor::Path({"x"}), arrow::schema({})); + auto status = put.status(); + if (status.ok()) { + status = put->writer->Close(); + } + EXPECT_FALSE(status.ok()); + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("no listener available")) + << "the async generic service must have served this call, not the sync path"; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, HandshakeAuthenticatesAndTokenIsRequired) { + AuthenticatedHarness harness; + ASSERT_OK(harness.status); + + // 1. Without a handshake the data call is rejected before it is served. + ASSERT_OK_AND_ASSIGN(auto unauthenticated_location, + Location::ForScheme("grpc", "127.0.0.1", harness.server.port())); + ASSERT_OK_AND_ASSIGN(auto unauthenticated_client, + FlightClient::Connect(unauthenticated_location)); + auto status = unauthenticated_client->DoGet(Ticket{"auth-ticket"}).status(); + EXPECT_FALSE(status.ok()); + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("Invalid token")); + ASSERT_OK(unauthenticated_client->Close()); + + // 2. After the handshake the same call succeeds and the server sees the + // validated identity. + ASSERT_OK(harness.client->Authenticate( + {}, std::make_unique("user", "p4ssw0rd"))); + ASSERT_OK_AND_ASSIGN(auto reader, harness.client->DoGet(Ticket{"auth-ticket"})); + int batches = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(auto chunk, reader->Next()); + if (!chunk.data) break; + batches++; + } + EXPECT_EQ(1, batches); + { + std::lock_guard guard(harness.server.state()->mutex); + EXPECT_EQ(1, harness.server.state()->handshakes); + EXPECT_EQ("user", harness.server.state()->peer_identity); + EXPECT_GE(harness.server.state()->validations, 1); + } +} + +TEST(AsyncGrpcTest, BadPasswordIsRejectedAtHandshake) { + AuthenticatedHarness harness; + ASSERT_OK(harness.status); + auto status = harness.client->Authenticate( + {}, std::make_unique("user", "wrong")); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("Invalid token")); +} + +TEST(AsyncGrpcTest, UnauthenticatedUnknownMethodIsRejectedBeforeUnimplemented) { + AuthenticatedHarness harness; + ASSERT_OK(harness.status); + // GetFlightInfo is not served by the async generic service, but auth runs + // first: the client must see the auth failure, not UNIMPLEMENTED. + auto status = harness.client->GetFlightInfo(FlightDescriptor::Path({"x"})).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("Invalid token")); + EXPECT_EQ(arrow::StatusCode::IOError, status.code()); +} + +// --------------------------------------------------------------------------- +// The synchronous request/response Handshake hook: payload fidelity, edge +// cases, and reactor cleanup (the reactor finishes on the strength of "one +// terminal path per direction", so the concurrency test is its real check). +// --------------------------------------------------------------------------- + +// Send `password`, require the server to echo `expect` back. +class EchoClientAuthHandler : public ClientAuthHandler { + public: + EchoClientAuthHandler(std::string password, std::string expect) + : password_(std::move(password)), expect_(std::move(expect)) {} + + Status Authenticate(ClientAuthSender* outgoing, ClientAuthReader* incoming) override { + ARROW_RETURN_NOT_OK(outgoing->Write(password_)); + ARROW_RETURN_NOT_OK(incoming->Read(&response_)); + if (response_ != expect_) { + return MakeFlightError(FlightStatusCode::Unauthenticated, + "unexpected handshake response"); + } + return Status::OK(); + } + + Status GetToken(std::string* token) override { + *token = password_; + return Status::OK(); + } + + const std::string& response() const { return response_; } + + private: + std::string password_; + std::string expect_; + std::string response_; +}; + +TEST(AsyncGrpcTest, HandshakePayloadIsBinarySafe) { + // The hook's request is a std::string taken from the protobuf and its + // response is echoed back through the client's reader: an embedded NUL and + // non-ASCII bytes must survive both directions unchanged. + const std::string binary("p4ss\0w0rd\xC3\xA9", 11); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + AsyncAuthTestServer server(binary, binary, "bin-ticket"); + FlightServerOptions options(location); + ASSERT_OK(server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(server.port())); + + auto handler = std::make_unique(binary, binary); + const auto* observed = handler.get(); + ASSERT_OK(client->Authenticate({}, std::move(handler))); + EXPECT_EQ(binary, observed->response()); + + // The token is the password, so an authenticated call must also work. + ASSERT_OK_AND_ASSIGN(auto reader, client->DoGet(Ticket{"bin-ticket"})); + ASSERT_OK_AND_ASSIGN(auto chunk, reader->Next()); + EXPECT_NE(nullptr, chunk.data); + + ASSERT_OK(client->Close()); + ASSERT_OK(server.Shutdown()); + ASSERT_OK(server.Wait()); +} + +TEST(AsyncGrpcTest, EmptyHandshakeRequestIsRejectedNotHung) { + // An empty HandshakeRequest is a valid protobuf message: the read succeeds, + // the hook sees an empty string, and its rejection must reach the client. + class SilentClientAuthHandler : public ClientAuthHandler { + public: + Status Authenticate(ClientAuthSender* outgoing, ClientAuthReader*) override { + return outgoing->Write(""); + } + Status GetToken(std::string*) override { return Status::OK(); } + }; + + AuthenticatedHarness harness; + ASSERT_OK(harness.status); + auto status = + harness.client->Authenticate({}, std::make_unique()); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("Invalid token")); + std::lock_guard guard(harness.server.state()->mutex); + EXPECT_EQ(1, harness.server.state()->handshakes) + << "the hook must run on an empty request, not fail the read"; +} + +TEST(AsyncGrpcTest, AbandonedHandshakeLeavesServerUsable) { + // A client that never sends a message: the server's read completes with + // ok=false, the reactor must still finish, and the server must keep serving. + class AbandoningClientAuthHandler : public ClientAuthHandler { + public: + Status Authenticate(ClientAuthSender*, ClientAuthReader*) override { + return Status::Invalid("client never speaks"); + } + Status GetToken(std::string*) override { return Status::OK(); } + }; + + AuthenticatedHarness harness; + ASSERT_OK(harness.status); + EXPECT_FALSE( + harness.client->Authenticate({}, std::make_unique()) + .ok()); + + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(harness.server.port())); + ASSERT_OK(client->Authenticate( + {}, std::make_unique("user", "p4ssw0rd"))); + ASSERT_OK_AND_ASSIGN(auto reader, client->DoGet(Ticket{"auth-ticket"})); + ASSERT_OK_AND_ASSIGN(auto chunk, reader->Next()); + EXPECT_NE(nullptr, chunk.data); + ASSERT_OK(client->Close()); + + std::lock_guard guard(harness.server.state()->mutex); + EXPECT_EQ(1, harness.server.state()->handshakes) + << "the abandoned call must never have reached the hook"; +} + +TEST(AsyncGrpcTest, ConcurrentHandshakesAllComplete) { + constexpr int kClients = 32; + AuthenticatedHarness harness; + ASSERT_OK(harness.status); + + std::vector threads; + std::vector statuses(kClients); + for (int i = 0; i < kClients; i++) { + threads.emplace_back([&harness, &statuses, i] { + auto location = Location::ForScheme("grpc", "127.0.0.1", harness.server.port()); + if (!location.ok()) { + statuses[i] = location.status(); + return; + } + auto client = FlightClient::Connect(*location); + if (!client.ok()) { + statuses[i] = client.status(); + return; + } + auto authenticated = (*client)->Authenticate( + {}, std::make_unique("user", "p4ssw0rd")); + if (!authenticated.ok()) { + statuses[i] = authenticated; + return; + } + auto reader = (*client)->DoGet(Ticket{"auth-ticket"}); + if (!reader.ok()) { + statuses[i] = reader.status(); + return; + } + auto chunk = (*reader)->Next(); + if (!chunk.ok()) { + statuses[i] = chunk.status(); + return; + } + statuses[i] = (*client)->Close(); + }); + } + for (auto& thread : threads) thread.join(); + for (int i = 0; i < kClients; i++) { + EXPECT_TRUE(statuses[i].ok()) << "client " << i << ": " << statuses[i]; + } + + std::lock_guard guard(harness.server.state()->mutex); + EXPECT_EQ(kClients, harness.server.state()->handshakes); +} + +TEST(AsyncGrpcTest, RepeatedHandshakesOnFreshChannels) { + constexpr int kRounds = 50; + AuthenticatedHarness harness; + ASSERT_OK(harness.status); + for (int i = 0; i < kRounds; i++) { + ASSERT_OK_AND_ASSIGN(auto location, + Location::ForScheme("grpc", "127.0.0.1", harness.server.port())); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(location)); + ASSERT_OK(client->Authenticate( + {}, std::make_unique("user", "p4ssw0rd"))); + ASSERT_OK(client->Close()); + } + std::lock_guard guard(harness.server.state()->mutex); + EXPECT_EQ(kRounds, harness.server.state()->handshakes); +} + +TEST(AsyncGrpcTest, LateHandshakeResponseLeavesServerUsable) { + // The client gives up while the hook is still running, so the server's write + // fails: the reactor's write-side terminal path must still Finish, and the + // server must keep serving. + class SlowAnsweringServer : public AsyncGenericFlightServerBase { + public: + Status Handshake(const ServerCallContext&, const std::string&, + std::string* response) override { + handshakes.fetch_add(1); + std::this_thread::sleep_for(std::chrono::milliseconds(300)); + *response = "user"; + return Status::OK(); + } + std::atomic handshakes{0}; + }; + + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + SlowAnsweringServer server; + FlightServerOptions options(location); + ASSERT_OK(server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(server.port())); + + FlightCallOptions impatient; + impatient.timeout = std::chrono::milliseconds(50); + auto status = client->Authenticate( + impatient, std::make_unique("user", "p4ssw0rd")); + EXPECT_FALSE(status.ok()) << "the client's deadline must expire during the hook"; + EXPECT_NE(std::string::npos, status.ToString().find("Deadline")) + << "expected a deadline error, got: " << status.ToString(); + EXPECT_EQ(1, server.handshakes.load()) + << "the first call must have reached the hook, so its write had to finish"; + + ASSERT_OK_AND_ASSIGN(auto second_location, + Location::ForScheme("grpc", "127.0.0.1", server.port())); + ASSERT_OK_AND_ASSIGN(auto second_client, FlightClient::Connect(second_location)); + ASSERT_OK(second_client->Authenticate( + {}, std::make_unique("user", "p4ssw0rd"))); + + ASSERT_OK(second_client->Close()); + ASSERT_OK(client->Close()); + ASSERT_OK(server.Shutdown()); + ASSERT_OK(server.Wait()); +} + +TEST(AsyncGrpcTest, EmptyHandshakeResponseIsDelivered) { + // A hook that answers nothing still produces exactly one (empty) response + // message: the client's Read must return an empty string rather than hang. + class NoResponseServer : public AsyncGenericFlightServerBase { + public: + Status Handshake(const ServerCallContext&, const std::string&, + std::string*) override { + return Status::OK(); + } + }; + + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + NoResponseServer server; + FlightServerOptions options(location); + ASSERT_OK(server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(server.port())); + + auto handler = std::make_unique("anything", std::string()); + const auto* observed = handler.get(); + ASSERT_OK(client->Authenticate({}, std::move(handler))); + EXPECT_EQ(std::string(), observed->response()); + + ASSERT_OK(client->Close()); + ASSERT_OK(server.Shutdown()); + ASSERT_OK(server.Wait()); +} + +TEST(AsyncGrpcTest, RehandshakeOnTheSameChannelWorks) { + // The Handshake RPC is per-call, so authenticating twice on one channel must + // work: the second handshake reaches the hook and the channel keeps serving. + AuthenticatedHarness harness; + ASSERT_OK(harness.status); + ASSERT_OK(harness.client->Authenticate( + {}, std::make_unique("user", "p4ssw0rd"))); + ASSERT_OK(harness.client->Authenticate( + {}, std::make_unique("user", "p4ssw0rd"))); + ASSERT_OK_AND_ASSIGN(auto reader, harness.client->DoGet(Ticket{"auth-ticket"})); + ASSERT_OK_AND_ASSIGN(auto chunk, reader->Next()); + EXPECT_NE(nullptr, chunk.data); + + std::lock_guard guard(harness.server.state()->mutex); + EXPECT_EQ(2, harness.server.state()->handshakes); +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_service.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_service.cc new file mode 100644 index 000000000000..e4daa60a85d8 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_service.cc @@ -0,0 +1,150 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_service.h" + +#include +#include +#include + +#include "arrow/flight/transport/grpc/async_grpc_service_internal.h" + +namespace arrow::flight::transport::grpc { + +namespace { + +namespace detail = arrow::flight::transport::grpc::detail; + +constexpr std::string_view kPrefix = "/arrow.flight.protocol.FlightService/"; +constexpr std::string_view kHandshakeMethod = "Handshake"; +constexpr std::string_view kDoGetMethod = "DoGet"; +constexpr std::string_view kDoPutMethod = "DoPut"; +constexpr std::string_view kDoActionMethod = "DoAction"; +constexpr std::string_view kListActionsMethod = "ListActions"; +constexpr std::string_view kDoExchangeMethod = "DoExchange"; +constexpr std::string_view kPollFlightInfoMethod = "PollFlightInfo"; +constexpr std::string_view kListFlightsMethod = "ListFlights"; +constexpr std::string_view kGetFlightInfoMethod = "GetFlightInfo"; +constexpr std::string_view kGetSchemaMethod = "GetSchema"; + +static const std::unordered_map kMethodMap = { + {kHandshakeMethod, FlightMethod::Handshake}, + {kListFlightsMethod, FlightMethod::ListFlights}, + {kGetFlightInfoMethod, FlightMethod::GetFlightInfo}, + {kGetSchemaMethod, FlightMethod::GetSchema}, + {kDoGetMethod, FlightMethod::DoGet}, + {kDoPutMethod, FlightMethod::DoPut}, + {kDoActionMethod, FlightMethod::DoAction}, + {kListActionsMethod, FlightMethod::ListActions}, + {kDoExchangeMethod, FlightMethod::DoExchange}, + {kPollFlightInfoMethod, FlightMethod::PollFlightInfo}, +}; + +FlightMethod MethodFromName(std::string_view method) { + if (!method.starts_with(kPrefix)) { + return FlightMethod::Invalid; + } + method.remove_prefix(kPrefix.size()); + const auto it = kMethodMap.find(method); + return it != kMethodMap.end() ? it->second : FlightMethod::Invalid; +} + +class Unimplemented : public ::grpc::ServerGenericBidiReactor { + public: + Unimplemented(detail::AsyncCallContext& flight_context, ::grpc::Status status) { + Finish(flight_context.FinishRequest(status)); + } + + explicit Unimplemented(::grpc::Status status) { Finish(status); } + + void OnDone() override { delete this; } +}; + +} // namespace + +AsyncGenericFlightService::AsyncGenericFlightService( + AsyncGenericFlightServerBase* async_base, + std::shared_ptr> helper, + HandshakeFn handshake_handler) + : base_(async_base), + helper_(std::move(helper)), + handshake_handler_(std::move(handshake_handler)) {} + +::grpc::ServerGenericBidiReactor* AsyncGenericFlightService::CreateReactor( + ::grpc::GenericCallbackServerContext* context) { + const std::string_view full_method = context->method(); + const std::string_view method = + full_method.starts_with(kPrefix) ? full_method.substr(kPrefix.size()) : full_method; + const FlightMethod flight_method = MethodFromName(full_method); + detail::AsyncCallContext flight_context(context); + + const auto prepare_status = + method == kHandshakeMethod + ? helper_->MakeCallContext(flight_method, context, &flight_context) + : helper_->CheckAuth(flight_method, context, &flight_context); + if (!prepare_status.ok()) { + // CheckAuth()/MakeCallContext() already ran FinishRequest() for this call, + // so only the gRPC status is left to send. + return new Unimplemented(prepare_status); + } + + switch (flight_method) { + case FlightMethod::Handshake: + if (!handshake_handler_) { + return new Unimplemented( + flight_context, + ::grpc::Status( + ::grpc::StatusCode::UNIMPLEMENTED, + "This service does not have an authentication mechanism enabled.")); + } + return detail::MakeHandshakeReactor(std::move(flight_context), handshake_handler_); + case FlightMethod::GetFlightInfo: + return detail::MakeGetFlightInfoReactor(std::move(flight_context), base_); + case FlightMethod::GetSchema: + return detail::MakeGetSchemaReactor(std::move(flight_context), base_); + case FlightMethod::PollFlightInfo: + return detail::MakePollFlightInfoReactor(std::move(flight_context), base_); + case FlightMethod::ListActions: + return detail::MakeListActionsReactor(std::move(flight_context), base_); + case FlightMethod::DoAction: + return detail::MakeDoActionReactor(std::move(flight_context), base_); + case FlightMethod::ListFlights: + return detail::MakeListFlightsReactor(std::move(flight_context), base_); + case FlightMethod::DoExchange: + return detail::MakeExchangeReactor(std::move(flight_context), base_); + case FlightMethod::DoGet: + return detail::MakeDoGetReactor(std::move(flight_context), base_); + case FlightMethod::DoPut: { + std::shared_ptr listener = + base_->CreateDoPutListener(flight_context); + if (!listener) { + return new Unimplemented( + flight_context, + ::grpc::Status(::grpc::StatusCode::UNIMPLEMENTED, + "DoPut is not implemented: no listener available")); + } + return detail::MakeDoPutReactor(std::move(flight_context), std::move(listener)); + } + case FlightMethod::Invalid: + break; + } + return new Unimplemented( + flight_context, + ::grpc::Status(::grpc::StatusCode::UNIMPLEMENTED, "Unknown method")); +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_service.h b/cpp/src/arrow/flight/transport/grpc/async_grpc_service.h new file mode 100644 index 000000000000..ae4f13a5fe02 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_service.h @@ -0,0 +1,81 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include + +#include "arrow/flight/transport/grpc/async_grpc_compat.h" +#include "arrow/flight/transport/grpc/grpc_server_internal.h" + +#include "arrow/flight/server_async.h" + +namespace arrow::flight::transport::grpc { + +#if ARROW_FLIGHT_HAS_ASYNC_SERVER + +using HandshakeFn = + GrpcServerCallContextHelper<::grpc::CallbackServerContext>::HandshakeFn; + +/// \brief Generic callback service dispatching Flight methods to reactors. +/// +/// Constructed with the AsyncGenericFlightServerBase whose handlers are served. +/// The server's middleware runs for every call, through the shared context +/// helper, which also authenticates every method except Handshake. Handshake, +/// DoGet, DoPut and the synchronous Flight methods (GetFlightInfo, GetSchema, +/// PollFlightInfo, ListActions, DoAction, ListFlights, DoExchange) have a +/// reactor. Every other method is answered UNIMPLEMENTED. + +// Silence warning +// "non dll-interface class CallbackGenericService used as base for dll-interface class" +# ifdef _MSC_VER +# pragma warning(push) +# pragma warning(disable : 4275) +# endif +class ARROW_FLIGHT_EXPORT AsyncGenericFlightService + : public ::grpc::CallbackGenericService { + public: + /// \param async_base is the async server whose handlers are served, it must + /// outlive the service. + /// \param helper runs middleware and auth and builds the call context. + /// \param handshake_handler is the server's Handshake hook, an empty hook + /// means the server has no authentication mechanism. + AsyncGenericFlightService( + AsyncGenericFlightServerBase* async_base, + std::shared_ptr> helper, + HandshakeFn handshake_handler = {}); + + /// \brief Creates a reactor for the incoming RPC. + /// \param context is the gRPC callback server context for the RPC. + /// \return a new reactor handling the RPC, or an Unimplemented reactor if the method is + /// unknown. + ::grpc::ServerGenericBidiReactor* CreateReactor( + ::grpc::GenericCallbackServerContext* context) override; + + private: + AsyncGenericFlightServerBase* base_; + std::shared_ptr> helper_; + /// The server class's Handshake hook. Empty when the server has none. + HandshakeFn handshake_handler_; +}; +# ifdef _MSC_VER +# pragma warning(pop) +# endif + +#endif // ARROW_FLIGHT_HAS_ASYNC_SERVER + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_service_internal.h b/cpp/src/arrow/flight/transport/grpc/async_grpc_service_internal.h new file mode 100644 index 000000000000..394a61b4512a --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_service_internal.h @@ -0,0 +1,232 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "arrow/flight/flight_data_decoder.h" +#include "arrow/flight/serialization_internal.h" +#include "arrow/flight/server_async.h" +#include "arrow/flight/transport/grpc/async_grpc_compat.h" +#include "arrow/flight/transport/grpc/async_grpc_service.h" +#include "arrow/flight/transport/grpc/serialization_internal.h" + +namespace arrow::flight::transport::grpc::detail { + +namespace pb = arrow::flight::protocol; + +#if ARROW_FLIGHT_HAS_ASYNC_SERVER + +using AsyncCallContext = GrpcServerCallContext<::grpc::CallbackServerContext>; + +template +::grpc::ByteBuffer MakeWriteBuffer(const ProtoT& message) { + const std::string bytes = message.SerializeAsString(); + ::grpc::Slice slice(bytes); + return {&slice, 1}; +} + +template +arrow::Result ParseProtoRequest(const ::grpc::ByteBuffer& buf, + std::string_view what) { + ARROW_ASSIGN_OR_RAISE(std::string bytes, + BytesFromBuffer(buf, "Failed to read request")); + PbT pb; + if (!pb.ParseFromArray(bytes.data(), static_cast(bytes.size()))) { + return arrow::Status::Invalid("Failed to parse ", what); + } + return pb; +} + +template +arrow::Result ParseProtoRequest(const ::grpc::ByteBuffer& buf, std::string_view what) { + ARROW_ASSIGN_OR_RAISE(auto pb, ParseProtoRequest(buf, what)); + T out; + ARROW_RETURN_NOT_OK(arrow::flight::internal::FromProto(pb, &out)); + return out; +} + +template +arrow::Status SerializeOrNotFound(const std::shared_ptr& value, PbT* out) { + if (value == nullptr) { + return arrow::Status::KeyError("Flight not found"); + } + return arrow::flight::internal::ToProto(*value, out); +} + +/// \brief Lifetime and termination handling shared by the callback reactors. +/// +/// A reactor's callbacks run on arbitrary gRPC threads, so it cannot be +/// destroyed while one is outstanding: the RPC's own lifetime holds the initial +/// reference (released in OnDone(), after which gRPC never touches the reactor), +/// and every callback that may outlive its StartRead()/StartWrite() takes +/// another one. The last ReleaseHold() deletes the reactor. +class AsyncReactorBase : public ::grpc::ServerGenericBidiReactor { + public: + explicit AsyncReactorBase(AsyncCallContext flight_context) + : flight_context_(std::move(flight_context)) { + live_reactors_.fetch_add(1, std::memory_order_relaxed); + } + + ~AsyncReactorBase() override { live_reactors_.fetch_sub(1, std::memory_order_relaxed); } + + /// \brief How many reactors are alive right now. + /// + /// A reactor is deleted by the last ReleaseHold(), so this returns to its + /// starting value once every RPC (including the ones whose handler abandoned + /// its future) has been torn down. + static int64_t live_reactors() { + return live_reactors_.load(std::memory_order_relaxed); + } + + /// The default cancellation: report the RPC as cancelled and finish it. + /// Reactors with their own teardown override this and call it first. + void OnCancel() override { FinishOnce(arrow::Status::Cancelled()); } + + void OnDone() override { ReleaseHold(); } + + protected: + /// \brief Finish the RPC with this status, later calls are no-ops. + /// + /// The status is the one the client sees, so the first one wins: the other + /// paths (cancel after error, error after cancel) must not overwrite it. + void FinishOnce(arrow::Status status) { + bool expected = false; + if (finished_.compare_exchange_strong(expected, true)) { + std::lock_guard lock(write_mutex_); + Finish(flight_context_.FinishRequest(status)); + } + } + + /// Whether FinishOnce() has already run (the RPC's status is settled). + bool finished() const { return finished_.load(); } + + /// The call context: handlers and the teardown paths need it. + AsyncCallContext& flight_context() { return flight_context_; } + + /// \brief Run `start_write` unless the RPC has already been finished. + /// + /// A handler continuation runs on the application's thread and may pass a + /// bare finished() check just before another path (a cancel, a deadline) + /// calls Finish() and gRPC tears the call down. Running the write under + /// write_mutex_ closes that window: FinishOnce() takes the same mutex, so + /// either the write wins and Finish() waits for it to be issued, or + /// finished() is already set and no write is started. A HoldToken keeps the + /// reactor alive but not the underlying gRPC call, which is why the check + /// alone is not enough. + /// \return true if `start_write` ran. + template + bool WriteIfLive(StartWriteFn&& start_write) { + std::lock_guard lock(write_mutex_); + if (finished()) { + return false; + } + start_write(); + return true; + } + + /// \brief The mutex guarding write initiation, see WriteIfLive(). + std::mutex& write_mutex() { return write_mutex_; } + + /// Account for work that may outlive the callback that started it. + void Hold() { refs_.fetch_add(1, std::memory_order_relaxed); } + + /// Drop a reference, deleting the reactor once the last one is gone. + void ReleaseHold() { + if (refs_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + delete this; + } + } + + /// A reference the reactor owns, to be handed to a callback. + class HoldToken { + public: + HoldToken() = default; + explicit HoldToken(AsyncReactorBase* reactor) : reactor_(reactor) { + reactor_->Hold(); + } + HoldToken(const HoldToken&) = delete; + HoldToken& operator=(const HoldToken&) = delete; + HoldToken(HoldToken&& other) noexcept + : reactor_(std::exchange(other.reactor_, nullptr)) {} + HoldToken& operator=(HoldToken&& other) noexcept { + if (this != &other) { + Release(); + reactor_ = std::exchange(other.reactor_, nullptr); + } + return *this; + } + ~HoldToken() { Release(); } + + /// \brief The reference, to be captured by the callback and moved out of. + HoldToken&& release() { return std::move(*this); } + + private: + void Release() { + if (reactor_ != nullptr) { + reactor_->ReleaseHold(); + reactor_ = nullptr; + } + } + AsyncReactorBase* reactor_ = nullptr; + }; + + /// Take a reference and return the token that holds it. + HoldToken hold() { return HoldToken(this); } + + private: + AsyncCallContext flight_context_; + std::mutex write_mutex_; + std::atomic finished_{false}; + std::atomic refs_{1}; + inline static std::atomic live_reactors_{0}; +}; + +::grpc::ServerGenericBidiReactor* MakeHandshakeReactor(AsyncCallContext flight_context, + HandshakeFn handshake_handler); +::grpc::ServerGenericBidiReactor* MakeGetFlightInfoReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base); +::grpc::ServerGenericBidiReactor* MakeGetSchemaReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base); +::grpc::ServerGenericBidiReactor* MakePollFlightInfoReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base); +::grpc::ServerGenericBidiReactor* MakeDoGetReactor(AsyncCallContext flight_context, + AsyncGenericFlightServerBase* base); +::grpc::ServerGenericBidiReactor* MakeDoPutReactor( + AsyncCallContext flight_context, std::shared_ptr listener); +::grpc::ServerGenericBidiReactor* MakeListActionsReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base); +::grpc::ServerGenericBidiReactor* MakeDoActionReactor(AsyncCallContext flight_context, + AsyncGenericFlightServerBase* base); +::grpc::ServerGenericBidiReactor* MakeListFlightsReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base); +::grpc::ServerGenericBidiReactor* MakeExchangeReactor(AsyncCallContext flight_context, + AsyncGenericFlightServerBase* base); + +#endif // ARROW_FLIGHT_HAS_ASYNC_SERVER + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_service_test.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_service_test.cc new file mode 100644 index 000000000000..086ab00938ed --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_service_test.cc @@ -0,0 +1,268 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Dispatch-level tests of the async generic gRPC Flight service: the +// UNIMPLEMENTED fallback reactor, middleware on every call (including the calls +// that never reach a handler), the options an async server accepts, and the two +// ways it can be started -- its own Init() or a hand-built grpc::ServerBuilder. +// +// Per-RPC behaviour lives in async_grpc_do_get_test.cc, async_grpc_do_put_test.cc, +// async_grpc_unary_test.cc, async_grpc_streaming_test.cc and +// async_grpc_exchange_test.cc. + +#include "arrow/flight/transport/grpc/async_grpc_test_fixtures.h" + +namespace arrow::flight::transport::grpc { + +TEST(AsyncGrpcTest, MiddlewareCallCompletedRunsOnUnknownMethod) { + // Middleware that started for a call must also see CallCompleted when the + // call is answered by the UNIMPLEMENTED fallback reactor: otherwise tracing + // middleware leaks a span per unknown/unimplemented RPC. + auto trace = std::make_shared(); + TestFlightServer flight_server; + auto helper = + std::make_shared>( + /*auth_handler=*/nullptr, + MiddlewareFactoryList{ + {"recording", std::make_shared(trace)}}); + AsyncGenericFlightService service(&flight_server, helper); + + int port = 0; + ::grpc::ServerBuilder builder; + builder.AddListeningPort("localhost:0", ::grpc::InsecureServerCredentials(), &port); + builder.RegisterCallbackGenericService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + ASSERT_NE(port, 0); + + const ::grpc::Status status = CallUnknownMethod(port); + ASSERT_EQ(::grpc::StatusCode::UNIMPLEMENTED, status.error_code()) + << status.error_message(); + + server->Shutdown(); + server->Wait(); + + std::lock_guard guard(trace->mutex); + EXPECT_EQ(1, trace->start_call) << "the middleware must have started for the call"; + EXPECT_EQ(1, trace->call_completed) + << "CallCompleted must run for an UNIMPLEMENTED reply too"; +} + +TEST(AsyncGrpcTest, OtherMethodsAreUnimplemented) { + TestFlightServer flight_server; + auto listener = std::make_shared(); + flight_server.set_listener(listener); + AsyncGenericFlightService service(&flight_server, MakeAsyncHelper()); + + int port = 0; + ::grpc::ServerBuilder builder; + builder.AddListeningPort("localhost:0", ::grpc::InsecureServerCredentials(), &port); + builder.RegisterCallbackGenericService(&service); + auto server = builder.BuildAndStart(); + ASSERT_NE(server, nullptr); + ASSERT_NE(port, 0); + + std::string uri = "grpc://localhost:" + std::to_string(port); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse(uri)); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(location)); + + // Every typed FlightService method is served now, the fallback reactor is + // still reachable through a method name no stub knows (which maps to + // FlightMethod::Invalid in MethodFromName), and that is the only pin left + // that an unknown method cannot be answered by a typed path. + const ::grpc::Status status = CallUnknownMethod(port); + ASSERT_FALSE(status.ok()) << "an unknown method must not be served"; + ASSERT_EQ(::grpc::StatusCode::UNIMPLEMENTED, status.error_code()) + << status.error_message(); + + ASSERT_OK(client->Close()); + server->Shutdown(); + server->Wait(); +} + +TEST(AsyncGrpcTest, UseAsyncGrpcFlag) { + ASSERT_OK_AND_ASSIGN(auto default_location, Location::Parse("grpc://localhost:0")); + FlightServerOptions default_options(default_location); + + // Same adapter as the tests above, but now started through the async + // server's own Init instead of a hand-built grpc::ServerBuilder. + TestFlightServer flight_server; + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_GT(flight_server.port(), 0); + + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + // (i) DoGet went through the generic reactor: the ticket reached + // TestFlightServer::DoGetAsync and the whole stream came back. + Ticket ticket{"flag-path-do-get-ticket"}; + ASSERT_OK_AND_ASSIGN(auto reader, client->DoGet(ticket)); + ASSERT_OK_AND_ASSIGN(auto schema, reader->GetSchema()); + ASSERT_EQ(schema->num_fields(), 2); + ASSERT_EQ(schema->field(0)->name(), "a"); + ASSERT_EQ(schema->field(1)->name(), "b"); + int batch_count = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(auto chunk, reader->Next()); + if (!chunk.data) break; + ASSERT_EQ(chunk.data->num_rows(), 5); + ASSERT_OK_AND_ASSIGN(auto value, chunk.data->column(0)->GetScalar(0)); + ASSERT_TRUE(value->Equals(*arrow::MakeScalar(int64_t(1)))); + batch_count++; + } + ASSERT_EQ(batch_count, 3); + + // (ii) Every typed method is served on this path too, the fallback reactor + // still answers a method name no stub knows, which is what pins "the typed + // sync service is NOT registered alongside the generic one". + const ::grpc::Status unknown_status = CallUnknownMethod(flight_server.port()); + ASSERT_FALSE(unknown_status.ok()) << "an unknown method must not be served"; + ASSERT_EQ(::grpc::StatusCode::UNIMPLEMENTED, unknown_status.error_code()) + << unknown_status.error_message(); + + // (iii) The async server takes no listener factory here, so uploads are + // refused the same way. + FlightDescriptor descriptor = FlightDescriptor::Path({"not-served"}); + auto put_schema = arrow::schema({arrow::field("a", arrow::int64())}); + auto batch = arrow::RecordBatch::Make(put_schema, 1, + {arrow::ArrayFromJSON(arrow::int64(), "[1]")}); + auto [surface, status] = UploadOneBatch(client.get(), descriptor, put_schema, batch); + ASSERT_FALSE(status.ok()) << "DoPut must be refused, surfaced on " << surface; + ASSERT_EQ(status.code(), arrow::StatusCode::NotImplemented); + + // Cleanup. The ticket is asserted only after the server stopped, so the + // DoGet callback cannot race with the read. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ(flight_server.ticket(), "flag-path-do-get-ticket"); +} + +TEST(AsyncGrpcTest, UseAsyncGrpcFlagServesDoPut) { + // The async server serves uploads because the server class hands out the + // per-RPC listener through CreateDoPutListener(), instead of the RPC falling + // through to UNIMPLEMENTED. + TestFlightServer flight_server; + auto listener = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + FlightServerOptions options(location); + flight_server.set_listener(listener); + ASSERT_OK(flight_server.Init(options)); + ASSERT_GT(flight_server.port(), 0); + + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto schema = arrow::schema( + {arrow::field("a", arrow::int64()), arrow::field("b", arrow::int64())}); + auto batch = + arrow::RecordBatch::Make(schema, 3, + {arrow::ArrayFromJSON(arrow::int64(), "[1, 2, 3]"), + arrow::ArrayFromJSON(arrow::int64(), "[10, 20, 30]")}); + + const FlightDescriptor descriptor = FlightDescriptor::Path({"wave", "h2"}); + auto [surface, status] = UploadOneBatch(client.get(), descriptor, schema, batch); + ASSERT_TRUE(status.ok()) << "upload failed on " << surface << ": " << status; + + // Cleanup before looking at the listener: its callbacks run on a gRPC + // thread, and nothing may read the recordings while that thread writes. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + // The listener saw the descriptor the client sent, the schema, and the batch. + ASSERT_EQ(listener->descriptor_count(), 1); + ASSERT_EQ(listener->descriptors()[0], descriptor); + ASSERT_EQ(listener->schema_count(), 1); + ASSERT_EQ(listener->schema()->num_fields(), 2); + ASSERT_EQ(listener->schema()->field(0)->name(), "a"); + ASSERT_EQ(listener->schema()->field(1)->name(), "b"); + ASSERT_EQ(listener->batches().size(), 1); + const std::shared_ptr chunk = listener->batches()[0]; + ASSERT_EQ(chunk->num_rows(), 3); + ASSERT_EQ(chunk->num_columns(), 2); + ASSERT_OK_AND_ASSIGN(auto value, chunk->column(0)->GetScalar(0)); + ASSERT_TRUE(value->Equals(*arrow::MakeScalar(int64_t(1)))); +} + +TEST(AsyncGrpcTest, UseAsyncGrpcFlagRefusesBlockingAuthHandler) { + // The callback path cannot run the blocking ServerAuthHandler (it would run + // on callback threads), so a server configured with one must refuse to start + // instead of silently serving unauthenticated requests. Middleware *is* + // supported now - see UseAsyncGrpcFlagAllowsMiddleware. + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + + TestFlightServer auth_server; + FlightServerOptions with_auth(location); + with_auth.auth_handler = std::make_shared("user", "pass"); + auto auth_status = auth_server.Init(with_auth); + ASSERT_FALSE(auth_status.ok()) << "an auth handler must not be silently skipped"; + ASSERT_EQ(auth_status.code(), arrow::StatusCode::NotImplemented); + ASSERT_NE(auth_status.message().find("AsyncGenericFlightServerBase"), std::string::npos) + << auth_status; +} + +TEST(AsyncGrpcTest, UseAsyncGrpcFlagAllowsMiddleware) { + // Middleware configured through FlightServerOptions must run on the async + // path: StartCall at call creation, SendingHeaders and CallCompleted at the + // end - success or failure alike. + auto trace = std::make_shared(); + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + + TestFlightServer flight_server; + FlightServerOptions options(location); + options.middleware = { + {"recording", std::make_shared(trace)}}; + ASSERT_OK(flight_server.Init(options)); + + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + ASSERT_OK_AND_ASSIGN(auto reader, client->DoGet(Ticket{"middleware-ticket"})); + while (true) { + ASSERT_OK_AND_ASSIGN(auto chunk, reader->Next()); + if (!chunk.data) break; + } + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + std::lock_guard guard(trace->mutex); + EXPECT_EQ(1, trace->start_call); + EXPECT_EQ(1, trace->sending_headers); + EXPECT_EQ(1, trace->call_completed); +} + +TEST(AsyncGrpcTest, UseAsyncGrpcFlagMiddlewareCanRejectCall) { + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + TestFlightServer flight_server; + FlightServerOptions options(location); + options.middleware = { + {"rejecting", std::make_shared()}}; + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto status = client->DoGet(Ticket{"rejected"}).status(); + EXPECT_FALSE(status.ok()); + EXPECT_THAT(status.ToString(), ::testing::HasSubstr("rejected by middleware")); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_streaming.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_streaming.cc new file mode 100644 index 000000000000..61757c0474ab --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_streaming.cc @@ -0,0 +1,254 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_service_internal.h" + +#include +#include +#include +#include +#include + +namespace arrow::flight::transport::grpc::detail { + +namespace { + +/** + * A generic streaming reactor for handling async gRPC calls in the Flight service. + * One request message is read, followed by N response messages, then the stream finishes. + */ +class StreamingReactor : public AsyncReactorBase { + public: + StreamingReactor(AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) + : AsyncReactorBase(std::move(flight_context)), base_(base) { + StartRead(&request_buf_); + } + + void OnReadDone(bool ok) override { + if (!ok) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, "Failed to read request")); + return; + } + Start().AddCallback([this, token = hold()](const arrow::Status& status) { + if (!finished()) { + if (!status.ok()) { + FinishOnce(status); + } else { + WriteNextMessage(); + } + } + }); + } + + void OnWriteDone(bool ok) override { + if (finished()) { + return; + } + if (!ok) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, "Failed to write response")); + return; + } + WriteNextMessage(); + } + + protected: + virtual arrow::Future<> Start() = 0; + virtual arrow::Future NextMessage() = 0; + + void WriteNextMessage() { + if (finished()) { + return; + } + NextMessage().AddCallback( + [this, token = hold()](const arrow::Result& has_next) { + if (!finished()) { + if (!has_next.ok()) { + FinishOnce(has_next.status()); + } else if (!*has_next) { + FinishOnce(arrow::Status::OK()); + } else { + WriteIfLive([this] { StartWrite(&write_buf_); }); + } + } + }); + } + + /// \brief Guards the subclass state that Start()'s continuation publishes and + /// NextMessage() consumes. NextMessage() runs on a gRPC callback thread + /// (OnWriteDone) while Start()'s continuation runs on the handler's, so both + /// share this one lock rather than each reactor declaring its own. + std::mutex& state_mutex() { return state_mutex_; } + + protected: + AsyncGenericFlightServerBase* base_; + ::grpc::ByteBuffer request_buf_; + ::grpc::ByteBuffer write_buf_; + + private: + std::mutex state_mutex_; +}; + +class ListActionsReactor final : public StreamingReactor { + public: + using StreamingReactor::StreamingReactor; + + protected: + arrow::Future<> Start() override { + return base_->ListActionsAsync(flight_context()) + .Then([this](const std::vector& actions) -> arrow::Status { + // Published on the handler's thread, read from NextMessage() which + // OnWriteDone() also calls from a gRPC thread: guard both. + std::lock_guard lock(state_mutex()); + actions_ = actions; + next_ = 0; + return arrow::Status::OK(); + }); + } + + arrow::Future NextMessage() override { + pb::ActionType pb_type; + { + std::lock_guard lock(state_mutex()); + if (next_ >= actions_.size()) { + return arrow::Future::MakeFinished(false); + } + const auto status = internal::ToProto(actions_[next_++], &pb_type); + if (!status.ok()) { + return arrow::Future::MakeFinished(status); + } + } + write_buf_ = MakeWriteBuffer(pb_type); + return arrow::Future::MakeFinished(true); + } + + private: + std::vector actions_; + size_t next_ = 0; +}; + +class DoActionReactor final : public StreamingReactor { + public: + using StreamingReactor::StreamingReactor; + + protected: + arrow::Future<> Start() override { + auto action = ParseProtoRequest(request_buf_, "Action"); + if (!action.ok()) { + return arrow::Future<>::MakeFinished(action.status()); + } + return base_->DoActionAsync(flight_context(), *action) + .Then([this](const std::shared_ptr& stream) -> arrow::Status { + if (stream == nullptr) { + return arrow::Status::Cancelled(); + } + { + std::lock_guard lock(state_mutex()); + stream_ = stream; + } + return arrow::Status::OK(); + }); + } + + arrow::Future NextMessage() override { + std::shared_ptr stream; + { + std::lock_guard lock(state_mutex()); + stream = stream_; + } + if (stream == nullptr) { + return arrow::Future::MakeFinished(false); + } + return stream->NextAsync().Then( + [this](const std::shared_ptr& result) -> arrow::Result { + if (result == nullptr) { + return false; + } + pb::Result pb_result; + ARROW_RETURN_NOT_OK(internal::ToProto(*result, &pb_result)); + write_buf_ = MakeWriteBuffer(pb_result); + return true; + }); + } + + private: + std::shared_ptr stream_; +}; + +class ListFlightsReactor final : public StreamingReactor { + public: + using StreamingReactor::StreamingReactor; + + protected: + arrow::Future<> Start() override { + auto criteria = ParseProtoRequest(request_buf_, "Criteria"); + if (!criteria.ok()) { + return arrow::Future<>::MakeFinished(criteria.status()); + } + criteria_ = *criteria; + return base_->ListFlightsAsync(flight_context(), &criteria_) + .Then( + [this](const std::shared_ptr& listing) -> arrow::Status { + std::lock_guard lock(state_mutex()); + listing_ = listing; + return arrow::Status::OK(); + }); + } + + arrow::Future NextMessage() override { + std::shared_ptr listing; + { + std::lock_guard lock(state_mutex()); + listing = listing_; + } + if (listing == nullptr) { + return arrow::Future::MakeFinished(false); + } + return listing->NextAsync().Then( + [this](const std::shared_ptr& info) -> arrow::Result { + if (info == nullptr) { + return false; + } + pb::FlightInfo pb_info; + ARROW_RETURN_NOT_OK(internal::ToProto(*info, &pb_info)); + write_buf_ = MakeWriteBuffer(pb_info); + return true; + }); + } + + private: + Criteria criteria_; + std::shared_ptr listing_; +}; + +} // namespace + +::grpc::ServerGenericBidiReactor* MakeListActionsReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) { + return new ListActionsReactor(std::move(flight_context), base); +} + +::grpc::ServerGenericBidiReactor* MakeDoActionReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) { + return new DoActionReactor(std::move(flight_context), base); +} + +::grpc::ServerGenericBidiReactor* MakeListFlightsReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) { + return new ListFlightsReactor(std::move(flight_context), base); +} + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_streaming_test.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_streaming_test.cc new file mode 100644 index 000000000000..9b0a8cadd2be --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_streaming_test.cc @@ -0,0 +1,304 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// The streaming RPCs served by StreamingReactor (one request in, N responses +// out, driven by OnWriteDone): ListActions (a vector), DoAction (an +// AsyncResultStream) and ListFlights (an AsyncFlightListing). Covers the +// empty-stream case (no write is ever armed), the mid-stream error rule, a +// listing completed on another thread, and a handler deadline. + +#include "arrow/flight/transport/grpc/async_grpc_test_fixtures.h" + +namespace arrow::flight::transport::grpc { + +// --------------------------------------------------------------------------- +// The streaming trio: ListActions (a vector), DoAction (an AsyncResultStream) +// and ListFlights (an AsyncFlightListing) share StreamingReactor (one request +// in, N responses out, driven by OnWriteDone -> WriteNext). Each RPC has a +// happy-path test, the reactor never writing a message must still finish the +// RPC (the empty tests), and the two iterator-backed ones pin the mid-stream +// error rule. +// --------------------------------------------------------------------------- + +TEST(AsyncGrpcTest, ListActionsIsServedAsync) { + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + ASSERT_OK_AND_ASSIGN(auto actions, client->ListActions()); + ASSERT_EQ(actions.size(), 2); + ASSERT_EQ(actions[0].type, "echo"); + ASSERT_EQ(actions[0].description, "echo back"); + ASSERT_EQ(actions[1].type, "fail"); + ASSERT_EQ(actions[1].description, "always fails"); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, ListActionsEmptyFinishesCleanly) { + // Zero response messages: WriteNext() finds the vector exhausted before it + // ever calls StartWrite(), so the RPC must be finished with OK rather than + // left waiting for a write completion that will never arrive. + TestFlightServer flight_server; + flight_server.set_empty_actions(); + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + ASSERT_OK_AND_ASSIGN(auto actions, client->ListActions()); + ASSERT_TRUE(actions.empty()); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, DoActionIsServedAsync) { + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + const Action action{"ping", Buffer::FromString("request-body")}; + ASSERT_OK_AND_ASSIGN(auto results, client->DoAction(action)); + ASSERT_OK_AND_ASSIGN(auto result, results->Next()); + ASSERT_NE(result, nullptr); + ASSERT_EQ(result->body->ToString(), "one"); + // The stream ends with a null Result, not with an error. + ASSERT_OK_AND_ASSIGN(auto end, results->Next()); + ASSERT_EQ(end, nullptr); + + // Cleanup before looking at the recording: it was written on a gRPC thread. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ("ping", flight_server.last_action_type()); +} + +TEST(AsyncGrpcTest, DoActionStreamsMultipleResults) { + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + const Action action{"multi", Buffer::FromString("")}; + ASSERT_OK_AND_ASSIGN(auto results, client->DoAction(action)); + + ASSERT_OK_AND_ASSIGN(auto bodies, DrainResultBodies(results.get())); + ASSERT_EQ(bodies, (std::vector{"one", "two", "three"})); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, DoActionMidStreamErrorFinishesTheRpc) { + // The stream yields one Result and then fails. The error must become the + // RPC status -- the client must not see a truncated success. + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + const Action action{"fail-mid-stream", Buffer::FromString("")}; + ASSERT_OK_AND_ASSIGN(auto results, client->DoAction(action)); + ASSERT_OK_AND_ASSIGN(auto first, results->Next()); + ASSERT_NE(first, nullptr); + ASSERT_EQ(first->body->ToString(), "first"); + + auto second = results->Next(); + ASSERT_FALSE(second.ok()) << "the mid-stream error must reach the client"; + ASSERT_EQ(second.status().code(), arrow::StatusCode::IOError); + ASSERT_THAT(second.status().message(), + ::testing::HasSubstr("mid-stream action failure")); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, DoActionWithNullResultStreamIsCancelled) { + // A handler that answers OK without a ResultStream answers CANCELLED on the + // sync transport (grpc_server.cc:404-406), before any message is written, the + // async reactor must answer the same. The client may see it on DoAction() or + // on the first Next(), so both surfaces are accepted. + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + const Action action{"null-result-stream", Buffer::FromString("")}; + auto results = client->DoAction(action); + if (results.ok()) { + // The stream was handed out before the status arrived: CANCELLED is then + // reported by the first read. + ASSERT_RAISES(Cancelled, (*results)->Next()); + } else { + ASSERT_RAISES(Cancelled, results.status()); + } + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, ListFlightsIsServedAsync) { + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + const Criteria criteria("filter-me"); + ASSERT_OK_AND_ASSIGN(auto listing, client->ListFlights({}, criteria)); + + ASSERT_OK_AND_ASSIGN(auto paths, DrainListingPaths(listing.get())); + ASSERT_EQ(paths, (std::vector>{{"first-listing"}, + {"second-listing"}})); + + // Cleanup before looking at the recording: it was written on a gRPC thread. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ("filter-me", flight_server.last_criteria().expression); +} + +TEST(AsyncGrpcTest, ListFlightsEmptyFinishesCleanly) { + // A listing with no flights: the server answers OK with zero messages, and + // the client sees an exhausted listing rather than a hang. + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + ASSERT_OK_AND_ASSIGN(auto listing, client->ListFlights({}, Criteria(""))); + ASSERT_OK_AND_ASSIGN(auto info, listing->Next()); + ASSERT_EQ(info, nullptr); + // The client sent no criteria, the server must still have received the + // (empty) Criteria, exactly as the sync handler does. + ASSERT_EQ("", flight_server.last_criteria().expression); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, ListFlightsMidStreamErrorFinishesTheRpc) { + // A listing that fails on its second Next(): the error must become the RPC + // status, not a silently truncated success. + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + auto listing = client->ListFlights({}, Criteria("fail-mid-stream")); + ASSERT_FALSE(listing.ok()) << "the mid-stream error must reach the client"; + ASSERT_EQ(listing.status().code(), arrow::StatusCode::IOError); + ASSERT_THAT(listing.status().message(), + ::testing::HasSubstr("mid-stream listing failure")); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, ListFlightsAsyncCompletesOnAnotherThread) { + // The listing object itself arrives on a helper thread: the transport must + // serve the response from wherever the future resolves. + class DeferredListingServer : public AsyncGenericFlightServerBase { + public: + arrow::Future> ListFlightsAsync( + const ServerCallContext&, const Criteria*) override { + auto out = arrow::Future>::Make(); + std::thread([out]() mutable { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + std::vector flights; + auto info = MakeTestFlightInfo(FlightDescriptor::Path({"deferred"})); + if (!info.ok()) { + out.MarkFinished(info.status()); + return; + } + flights.push_back(std::move(**info)); + out.MarkFinished(std::make_shared(std::move(flights))); + }).detach(); + return out; + } + }; + + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + DeferredListingServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN( + auto client_location, + Location::Parse("grpc://localhost:" + std::to_string(flight_server.port()))); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(client_location)); + + ASSERT_OK_AND_ASSIGN(auto listing, client->ListFlights()); + int seen = 0; + while (true) { + ASSERT_OK_AND_ASSIGN(auto info, listing->Next()); + if (info == nullptr) break; + ++seen; + } + ASSERT_EQ(1, seen); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +// A ListFlights handler whose future never resolves. +class HangingListingServer final : public AsyncGenericFlightServerBase { + public: + arrow::Future> ListFlightsAsync( + const ServerCallContext&, const Criteria*) override { + return arrow::Future>::Make(); + } +}; + +TEST(AsyncGrpcTest, ListFlightsDeadlineWithPendingHandlerFinishes) { + // The handler's future never resolves and the client's deadline ends the RPC: + // the reactor must be released even though the handler never completes. + const int64_t before = detail::AsyncReactorBase::live_reactors(); + + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + HangingListingServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + FlightCallOptions impatient; + impatient.timeout = std::chrono::milliseconds(500); + auto listing = client->ListFlights(impatient, Criteria("hang")); + EXPECT_FALSE(listing.ok()) << "the deadline must expire while the handler hangs"; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + EXPECT_TRUE(WaitFor([&] { + return detail::AsyncReactorBase::live_reactors() == before; + })) << "a reactor leaked: still alive after the server stopped (before=" + << before << ", now=" << detail::AsyncReactorBase::live_reactors() << ")"; +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_test_fixtures.h b/cpp/src/arrow/flight/transport/grpc/async_grpc_test_fixtures.h new file mode 100644 index 000000000000..e08960ad7119 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_test_fixtures.h @@ -0,0 +1,865 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Shared fixtures for the async generic gRPC Flight transport tests. +// +// One translation unit per part of the transport would otherwise duplicate the +// adapters, the canned test servers and the gRPC call helpers, so they live +// here. This header is internal (the `_fixtures` suffix keeps it out of the +// installed headers) and is only included by the async_grpc_*_test.cc files. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "arrow/flight/client.h" +#include "arrow/flight/server.h" +#include "arrow/flight/server_async.h" +#include "arrow/flight/server_middleware.h" +#include "arrow/flight/test_auth_handlers.h" +#include "arrow/flight/transport/grpc/async_grpc_service.h" +#include "arrow/flight/transport/grpc/async_grpc_service_internal.h" +#include "arrow/ipc/writer.h" +#include "arrow/scalar.h" +#include "arrow/testing/gtest_util.h" + +namespace arrow::flight::transport::grpc { + +// The ticket the test server answers with a stream that has a schema and no +// batches at all, every other ticket gets the 3-batch stream. +constexpr const char* kSchemaOnlyTicket = "schema-only-do-get-ticket"; + +// One canned FlightInfo, `descriptor` names it for the assertions. +inline arrow::Result> MakeTestFlightInfo( + FlightDescriptor descriptor) { + ARROW_ASSIGN_OR_RAISE( + auto made, FlightInfo::Make(*arrow::schema({arrow::field("value", arrow::int64())}), + std::move(descriptor), std::vector{}, + /*total_records=*/-1, /*total_bytes=*/-1)); + return std::make_unique(std::move(made)); +} + +// An AsyncFlightDataStream over an in-memory RecordBatchReader: every payload is +// already available, so each future is finished when it is created. +class ResolvedStream : public AsyncFlightDataStream { + public: + explicit ResolvedStream(std::shared_ptr reader) + : reader_(std::move(reader)), mapper_(*reader_->schema()) {} + + std::shared_ptr schema() override { return reader_->schema(); } + + arrow::Future GetSchemaPayloadAsync() override { + FlightPayload payload; + auto status = + ipc::GetSchemaPayload(*schema(), ipc_options_, mapper_, &payload.ipc_message); + if (!status.ok()) { + return arrow::Future::MakeFinished(status); + } + return arrow::Future::MakeFinished(std::move(payload)); + } + + arrow::Future> NextAsync() override { + auto batch = reader_->Next(); + if (!batch.ok()) { + return arrow::Future>::MakeFinished(batch.status()); + } + if (*batch == nullptr) { + // End of stream: the async interface reports it as a nullopt. + return arrow::Future>::MakeFinished(std::nullopt); + } + FlightPayload payload; + auto status = ipc::GetRecordBatchPayload(**batch, ipc_options_, &payload.ipc_message); + if (!status.ok()) { + return arrow::Future>::MakeFinished(status); + } + return arrow::Future>::MakeFinished(std::move(payload)); + } + + Status Close() override { return reader_->Close(); } + + private: + std::shared_ptr reader_; + ipc::DictionaryFieldMapper mapper_; + ipc::IpcWriteOptions ipc_options_ = ipc::IpcWriteOptions::Defaults(); +}; + +// A vector-backed async iterator: every NextAsync() resolves immediately, and +// exhaustion is reported as a null item. Instantiated below for both iterator +// interfaces the streaming RPCs use. +template +class VectorAsyncIterator final : public Base { + public: + explicit VectorAsyncIterator(std::vector items) : items_(std::move(items)) {} + + arrow::Future> NextAsync() override { + if (next_ >= items_.size()) { + return arrow::Future>::MakeFinished(std::shared_ptr{}); + } + return arrow::Future>::MakeFinished( + std::make_shared(std::move(items_[next_++]))); + } + + private: + std::vector items_; + size_t next_ = 0; +}; + +using SimpleListingAdapter = VectorAsyncIterator; +using VectorAsyncResultStream = VectorAsyncIterator; + +// An async iterator that yields one item and then fails: the error must become +// the RPC status, not truncate the stream into a success. `make_item` produces +// that first item, or the status that fails the RPC before it. +template +class FailAfterFirstAsyncIterator final : public Base { + public: + FailAfterFirstAsyncIterator(std::function()> make_item, + std::string error_message) + : make_item_(std::move(make_item)), error_message_(std::move(error_message)) {} + + arrow::Future> NextAsync() override { + if (served_) { + return arrow::Future>::MakeFinished( + arrow::Status::IOError(error_message_)); + } + served_ = true; + auto item = make_item_(); + if (!item.ok()) { + return arrow::Future>::MakeFinished(item.status()); + } + return arrow::Future>::MakeFinished( + std::make_shared(std::move(*item))); + } + + private: + std::function()> make_item_; + std::string error_message_; + bool served_ = false; +}; + +using MidStreamErrorAsyncResultStream = + FailAfterFirstAsyncIterator; +using MidStreamErrorAsyncListing = + FailAfterFirstAsyncIterator; + +// Serves the canned answers the async-transport tests drive: 3 batches of 5 +// rows for any ticket except kSchemaOnlyTicket, uploads through the listener the +// test installs, and one canned answer per unary/streaming RPC. The handlers +// run on gRPC threads, the tests read the recordings only after the server has +// shut down. +class TestFlightServer final : public AsyncGenericFlightServerBase { + public: + // The listener the next DoPut RPC is handed, nullptr (the default) refuses + // uploads, which is the base class behavior. + std::shared_ptr CreateDoPutListener( + const ServerCallContext&) override { + return listener_; + } + + void set_listener(std::shared_ptr listener) { + listener_ = std::move(listener); + } + + arrow::Future> DoGetAsync( + const ServerCallContext&, const Ticket& request) override { + ticket_ = request.ticket; + auto schema = arrow::schema( + {arrow::field("a", arrow::int64()), arrow::field("b", arrow::int64())}); + if (request.ticket == kSchemaOnlyTicket) { + // An empty reader still carries the schema, NextAsync() reports the end + // of stream right away. + auto reader = arrow::RecordBatchReader::Make({}, schema); + if (!reader.ok()) { + return arrow::Future>::MakeFinished( + reader.status()); + } + return arrow::Future>::MakeFinished( + std::make_shared(std::move(*reader))); + } + auto batch = arrow::RecordBatch::Make( + schema, 5, + {arrow::ArrayFromJSON(arrow::int64(), "[1, 2, 3, 4, 5]"), + arrow::ArrayFromJSON(arrow::int64(), "[10, 20, 30, 40, 50]")}); + auto reader = arrow::RecordBatchReader::Make({batch, batch, batch}); + if (!reader.ok()) { + return arrow::Future>::MakeFinished( + reader.status()); + } + return arrow::Future>::MakeFinished( + std::make_shared(std::move(*reader))); + } + + // The unary trio: canned responses. `last_descriptor_` records the request + // of the last one served (it is written on a gRPC thread, so tests read it + // only after the server has shut down), Command("fail") is the sentinel that + // drives the error-path test. + arrow::Future> GetFlightInfoAsync( + const ServerCallContext&, const FlightDescriptor& request) override { + last_descriptor_ = request; + if (request == FlightDescriptor::Command("fail")) { + return arrow::Future>::MakeFinished( + arrow::Status::KeyError("no such flight")); + } + if (request == FlightDescriptor::Command("null-info")) { + // Answered with a null FlightInfo: the transport must answer what the + // sync transport does for that case (grpc_server.cc: "Flight not found"). + return arrow::Future>::MakeFinished( + std::shared_ptr{}); + } + auto made = + FlightInfo::Make(*arrow::schema({arrow::field("value", arrow::int64())}), + FlightDescriptor{}, std::vector{}, -1, -1); + if (!made.ok()) { + return arrow::Future>::MakeFinished(made.status()); + } + return arrow::Future>::MakeFinished( + std::make_shared(std::move(*made))); + } + + arrow::Future> GetSchemaAsync( + const ServerCallContext&, const FlightDescriptor& request) override { + last_descriptor_ = request; + auto made = + SchemaResult::Make(*arrow::schema({arrow::field("value", arrow::int64())})); + if (!made.ok()) { + return arrow::Future>::MakeFinished(made.status()); + } + return arrow::Future>::MakeFinished( + std::shared_ptr(std::move(*made))); + } + + arrow::Future> PollFlightInfoAsync( + const ServerCallContext&, const FlightDescriptor& request) override { + last_descriptor_ = request; + auto poll_info = std::make_shared(); + poll_info->descriptor = FlightDescriptor::Command("poll-next"); + poll_info->progress = 0.5; + return arrow::Future>::MakeFinished(std::move(poll_info)); + } + + // The streaming trio: ListActions (a vector), DoAction (an AsyncResultStream) + // and ListFlights (an AsyncFlightListing). `last_criteria_` records the + // filter the client sent (read only after the server stopped), the + // "fail-mid-stream" sentinels drive the error-path tests. + arrow::Future> ListActionsAsync( + const ServerCallContext&) override { + std::vector actions = {ActionType("echo", "echo back"), + ActionType("fail", "always fails")}; + if (empty_actions_) { + actions.clear(); + } + return arrow::Future>::MakeFinished(std::move(actions)); + } + + arrow::Future> DoActionAsync( + const ServerCallContext&, const Action& action) override { + last_action_type_ = action.type; + if (action.type == "fail-mid-stream") { + return arrow::Future>::MakeFinished( + std::make_shared( + [] { return Result{arrow::Buffer::FromString("first")}; }, + "mid-stream action failure")); + } + if (action.type == "null-result-stream") { + // Answered with a null stream: the transport answers CANCELLED for this + // case (grpc_server.cc:404-406), before writing anything. + return arrow::Future>::MakeFinished( + std::shared_ptr{}); + } + std::vector made; + made.emplace_back(arrow::Buffer::FromString("one")); + if (action.type == "multi") { + made.emplace_back(arrow::Buffer::FromString("two")); + made.emplace_back(arrow::Buffer::FromString("three")); + } + return arrow::Future>::MakeFinished( + std::make_shared(std::move(made))); + } + + arrow::Future> ListFlightsAsync( + const ServerCallContext&, const Criteria* criteria) override { + if (criteria != nullptr) { + last_criteria_ = *criteria; + } + const std::string expression = criteria == nullptr ? "" : criteria->expression; + if (expression == "fail-mid-stream") { + return arrow::Future>::MakeFinished( + std::make_shared( + []() -> arrow::Result { + ARROW_ASSIGN_OR_RAISE( + auto info, + MakeTestFlightInfo(FlightDescriptor::Path({"first-listing"}))); + return std::move(*info); + }, + "mid-stream listing failure")); + } + if (expression.empty()) { + // No criteria: the sync handler passes an empty (never null) Criteria + // to the handler and the async reactor must do the same, so an empty + // expression is reachable here. + return arrow::Future>::MakeFinished( + std::make_shared(std::vector{})); + } + std::vector flights; + for (const auto& name : {"first-listing", "second-listing"}) { + auto info = MakeTestFlightInfo(FlightDescriptor::Path({name})); + if (!info.ok()) { + return arrow::Future>::MakeFinished( + info.status()); + } + flights.push_back(std::move(**info)); + } + return arrow::Future>::MakeFinished( + std::make_shared(std::move(flights))); + } + + void set_empty_actions() { empty_actions_ = true; } + + const std::string& ticket() const { return ticket_; } + const FlightDescriptor& last_descriptor() const { return last_descriptor_; } + const std::string& last_action_type() const { return last_action_type_; } + const Criteria& last_criteria() const { return last_criteria_; } + + private: + std::shared_ptr listener_; + std::string ticket_; + FlightDescriptor last_descriptor_; + std::string last_action_type_; + Criteria last_criteria_; + bool empty_actions_ = false; +}; + +// Echoes every chunk of a DoExchange back, in order, on the async API: the +// reader's demand drives the loop, each write is awaited, and the all-null +// chunk ends the exchange. `exchanges_in_flight` proves an abandoned +// exchange still leaves the handler: the state object is freed when the last +// continuation dies and decrements the counter on destruction. +class EchoExchangeTestServer : public AsyncGenericFlightServerBase { + public: + arrow::Future<> DoExchangeAsync( + const ServerCallContext&, std::shared_ptr reader, + std::shared_ptr writer) override { + exchange_count_.fetch_add(1); + exchanges_in_flight_.fetch_add(1); + auto state = std::make_shared(); + state->reader = std::move(reader); + state->writer = std::move(writer); + state->in_flight = &exchanges_in_flight_; + return EchoNext(state); + } + + int exchanges_in_flight() const { return exchanges_in_flight_.load(); } + int exchange_count() const { return exchange_count_.load(); } + + private: + struct EchoState { + std::shared_ptr reader; + std::shared_ptr writer; + std::atomic* in_flight = nullptr; + bool begun = false; + ~EchoState() { in_flight->fetch_sub(1); } + }; + + arrow::Future<> EchoNext(const std::shared_ptr& state) { + // The recursion depth follows the batch count of one exchange. Flatten + // into a loop if a single exchange ever carries thousands. + return state->reader->NextAsync().Then( + [this, state](const FlightStreamChunk& chunk) -> arrow::Future<> { + if (chunk.data == nullptr && chunk.app_metadata == nullptr) { + // End of the exchange. + return arrow::Future<>::MakeFinished(); + } + arrow::Future<> written = arrow::Future<>::MakeFinished(); + if (chunk.data != nullptr && !state->begun) { + state->begun = true; + written = state->writer->BeginAsync(chunk.data->schema()); + } + if (chunk.data != nullptr && chunk.app_metadata != nullptr) { + written = written.Then([state, chunk]() { + return state->writer->WriteWithMetadataAsync(*chunk.data, + chunk.app_metadata); + }); + } else if (chunk.data != nullptr) { + written = written.Then([state, chunk]() { + return state->writer->WriteRecordBatchAsync(*chunk.data); + }); + } else { + written = written.Then([state, chunk]() { + return state->writer->WriteMetadataAsync(chunk.app_metadata); + }); + } + return written.Then([this, state]() { return EchoNext(state); }); + }); + } + + std::atomic exchanges_in_flight_{0}; + std::atomic exchange_count_{0}; +}; + +// Records what the decoder hands to the application: the schema, every chunk, +// and whether the consumer accepted it. OnNext returns whatever status the +// test configured, so the upload's error path is reachable from here. +class RecordingListener : public AsyncFlightDataListener { + public: + arrow::Status OnSchemaDecoded(std::shared_ptr schema) override { + std::lock_guard lock(mutex_); + ++schema_count_; + schema_ = std::move(schema); + return arrow::Status::OK(); + } + + arrow::Future<> OnNext(FlightStreamChunk chunk) override { + std::lock_guard lock(mutex_); + if (chunk.data) { + batches_.push_back(std::move(chunk.data)); + } else { + ++metadata_only_count_; + metadata_chunks_.push_back(std::move(chunk.app_metadata)); + } + last_status_ = next_status_; + return arrow::Future<>::MakeFinished(next_status_); + } + + // Records the upload's descriptor and reports whatever status the test + // configured, so the descriptor rejection path is reachable from here too. + arrow::Future<> OnDescriptor(const FlightDescriptor& descriptor) override { + std::lock_guard lock(mutex_); + descriptors_.push_back(descriptor); + return arrow::Future<>::MakeFinished(descriptor_status_); + } + + // Consumer-side rejection of the upload, as the decoder tests do. + void set_next_status(arrow::Status status) { + std::lock_guard lock(mutex_); + next_status_ = std::move(status); + } + // Rejection before any schema or data arrives. + void set_descriptor_status(arrow::Status status) { + std::lock_guard lock(mutex_); + descriptor_status_ = std::move(status); + } + + int schema_count() const { + std::lock_guard lock(mutex_); + return schema_count_; + } + int descriptor_count() const { + std::lock_guard lock(mutex_); + return static_cast(descriptors_.size()); + } + int metadata_only_count() const { + std::lock_guard lock(mutex_); + return metadata_only_count_; + } + std::shared_ptr schema() const { + std::lock_guard lock(mutex_); + return schema_; + } + std::vector descriptors() const { + std::lock_guard lock(mutex_); + return descriptors_; + } + std::vector> batches() const { + std::lock_guard lock(mutex_); + return batches_; + } + std::vector> metadata_chunks() const { + std::lock_guard lock(mutex_); + return metadata_chunks_; + } + arrow::Status last_status() const { + std::lock_guard lock(mutex_); + return last_status_; + } + + protected: + /// The callbacks run on gRPC threads while the test thread reads the + /// accessors, so everything they share goes through this. + mutable std::mutex mutex_; + int schema_count_ = 0; + int metadata_only_count_ = 0; + std::shared_ptr schema_; + std::vector descriptors_; + std::vector> batches_; + std::vector> metadata_chunks_; + arrow::Status descriptor_status_; + arrow::Status next_status_; + arrow::Status last_status_; + + protected: + int finish_count_ = 0; + int cancel_count_ = 0; + Status finish_status_; + Status cancel_status_; +}; + +// A listener that records the terminal events on top of RecordingListener: +// OnFinish (the upload ended, whichever way) and what Cancel() reported, so +// the tests can assert on an upload cancelled from the server side. +class FinishRecordingListener : public RecordingListener { + public: + arrow::Future<> OnFinish(Status status) override { + std::lock_guard lock(mutex_); + finish_count_++; + finish_status_ = std::move(status); + return arrow::Future<>::MakeFinished(arrow::Status::OK()); + } + + /// Cancel from whatever thread the test drives, records what Cancel() said. + void CancelWith(Status status) { + const arrow::Status cancel_status = Cancel(std::move(status)).status(); + std::lock_guard lock(mutex_); + cancel_status_ = cancel_status; + ++cancel_count_; + } + + int finish_count() const { + std::lock_guard lock(mutex_); + return finish_count_; + } + int cancel_count() const { + std::lock_guard lock(mutex_); + return cancel_count_; + } + Status finish_status() const { + std::lock_guard lock(mutex_); + return finish_status_; + } + Status cancel_status() const { + std::lock_guard lock(mutex_); + return cancel_status_; + } +}; + +// A server that accepts uploads: the transport asks the server for the per-RPC +// listener through CreateDoPutListener(). One fresh listener per RPC, as +// overlapping uploads require. The test either reads them after the server has +// stopped, or waits for them with WaitForListener when it must act mid-upload. +// +// The listener type is the fixture's parameter, so a test can plug in one that +// parks, defers or cancels an upload without re-implementing the seam. +template +class RecordingUploadServer : public AsyncGenericFlightServerBase { + public: + std::shared_ptr CreateDoPutListener( + const ServerCallContext&) override { + auto listener = std::make_shared(); + std::lock_guard lock(mutex_); + listeners_.push_back(listener); + return listener; + } + + int calls() const { + std::lock_guard lock(mutex_); + return static_cast(listeners_.size()); + } + + std::shared_ptr listener(size_t i) const { + std::lock_guard lock(mutex_); + return listeners_[i]; + } + + /// \brief Wait (bounded) for the listener of the i-th upload to exist. + /// The factory runs on a gRPC thread, so a test that must act mid-upload + /// waits for it here instead of racing the read. + std::shared_ptr WaitForListener(size_t i = 0) const { + for (int attempt = 0; attempt < 500; attempt++) { + { + std::lock_guard lock(mutex_); + if (listeners_.size() > i) return listeners_[i]; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + return nullptr; + } + + private: + mutable std::mutex mutex_; + std::vector> listeners_; +}; + +using TestUploadServer = RecordingUploadServer; + +/// A listener that acknowledges every batch it is handed with one application +/// metadata message, the way an application reports per-chunk progress. The +/// batch recording itself is RecordingListener's. +class MetadataReplyingListener final : public RecordingListener { + public: + arrow::Future<> OnNext(FlightStreamChunk chunk) override { + if (chunk.data) { + const Status status = writer()->WriteMetadata(*arrow::Buffer::FromString( + "ack " + std::to_string(chunk.data->num_rows()) + " rows")); + std::lock_guard lock(mutex_); + last_write_ = status; + ++wrote_; + } + return RecordingListener::OnNext(std::move(chunk)); + } + + int wrote() const { + std::lock_guard lock(mutex_); + return wrote_; + } + Status last_write() const { + std::lock_guard lock(mutex_); + return last_write_; + } + + private: + int wrote_ = 0; + Status last_write_; +}; + +/// A listener whose OnFinish fails asynchronously: the commit it performs is +/// what decides whether the upload succeeded, so the client must not be told +/// the upload is OK. +class FailingFinishListener final : public FinishRecordingListener { + public: + arrow::Future<> OnFinish(Status status) override { + static_cast(FinishRecordingListener::OnFinish(std::move(status))); + return arrow::Future<>::MakeFinished( + arrow::Status::IOError("the upload could not be committed")); + } +}; + +/// A listener that calls back into the writer from OnFinish: OnFinish is +/// application code running on a transport thread, so the transport must not +/// hold its own write lock across it. +class FinishingAckListener final : public FinishRecordingListener { + public: + arrow::Future<> OnFinish(Status status) override { + finish_write = writer()->WriteMetadata(*arrow::Buffer::FromString("from-onfinish")); + return FinishRecordingListener::OnFinish(std::move(status)); + } + + Status finish_write; +}; + +/// Hands out a MetadataReplyingListener for every upload. +using MetadataReplyingUploadServer = RecordingUploadServer; + +// Middleware execution counts, shared between the factory, the middleware and +// the test thread (all three touch it). +struct MiddlewareTrace { + std::mutex mutex; + int start_call = 0; + int sending_headers = 0; + int call_completed = 0; +}; + +// Records middleware execution so a test can assert StartCall/CallCompleted +// ran on the async path (the sync path is covered by TestMiddleware). +class RecordingServerMiddleware : public ServerMiddleware { + public: + explicit RecordingServerMiddleware(std::shared_ptr trace) + : trace_(std::move(trace)) {} + + std::string name() const override { return "recording"; } + void SendingHeaders(AddCallHeaders* /*outgoing_headers*/) override { + std::lock_guard guard(trace_->mutex); + trace_->sending_headers++; + } + void CallCompleted(const Status& /*status*/) override { + std::lock_guard guard(trace_->mutex); + trace_->call_completed++; + } + + private: + std::shared_ptr trace_; +}; + +class RecordingServerMiddlewareFactory : public ServerMiddlewareFactory { + public: + explicit RecordingServerMiddlewareFactory(std::shared_ptr trace) + : trace_(std::move(trace)) {} + + Status StartCall(const CallInfo& /*info*/, const ServerCallContext& /*context*/, + std::shared_ptr* middleware) override { + { + std::lock_guard guard(trace_->mutex); + trace_->start_call++; + } + *middleware = std::make_shared(trace_); + return Status::OK(); + } + + private: + std::shared_ptr trace_; +}; + +class RejectingServerMiddlewareFactory : public ServerMiddlewareFactory { + public: + Status StartCall(const CallInfo&, const ServerCallContext&, + std::shared_ptr*) override { + return Status::Invalid("rejected by middleware"); + } +}; + +// The gRPC transport builds the generic service with the shared context +// helper (which runs middleware and auth). Tests that register the service +// on a hand-built grpc::ServerBuilder must pass it explicitly. +inline std::shared_ptr> +MakeAsyncHelper() { + return std::make_shared>( + /*auth_handler=*/nullptr, MiddlewareFactoryList{}); +} + +// The method name no FlightService stub knows: the transport must answer it +// with the UNIMPLEMENTED fallback reactor rather than any typed path. It is +// called over a raw generic stub, because only a generic client can name a +// method the service does not have. +constexpr const char* kUnknownMethod = + "/arrow.flight.protocol.FlightService/NoSuchMethod"; + +// Call `kUnknownMethod` on the server listening on `port` with the sync-style +// generic stub API, and return the server's status. +inline ::grpc::Status CallUnknownMethod(int port) { + std::shared_ptr<::grpc::Channel> channel = ::grpc::CreateChannel( + "127.0.0.1:" + std::to_string(port), ::grpc::InsecureChannelCredentials()); + ::grpc::GenericStub stub(channel); + ::grpc::ClientContext context; + ::grpc::ByteBuffer request; + ::grpc::ByteBuffer response; + ::grpc::CompletionQueue cq; + std::unique_ptr<::grpc::ClientAsyncResponseReader<::grpc::ByteBuffer>> call = + stub.PrepareUnaryCall(&context, kUnknownMethod, request, &cq); + call->StartCall(); + ::grpc::Status status; + call->Finish(&response, &status, reinterpret_cast(1)); + void* tag = nullptr; + bool ok = false; + cq.Next(&tag, &ok); + return status; +} + +// The single int64 column ("a") the upload tests use. +inline std::shared_ptr OneColumnSchema() { + return arrow::schema({arrow::field("a", arrow::int64())}); +} + +// One 2-row batch of OneColumnSchema(), rows [1, 2]. +inline std::shared_ptr OneColumnBatch( + const std::shared_ptr& schema = OneColumnSchema()) { + return arrow::RecordBatch::Make(schema, 2, + {arrow::ArrayFromJSON(arrow::int64(), "[1, 2]")}); +} + +// The listen address the tests start their servers on: an ephemeral port. +inline const Location& LocalhostLocation() { + static const Location location = *Location::Parse("grpc://localhost:0"); + return location; +} + +// Connect a FlightClient to the server listening on `port`. +inline arrow::Result> ConnectTo(int port) { + ARROW_ASSIGN_OR_RAISE(auto location, Location::ForScheme("grpc", "127.0.0.1", port)); + return FlightClient::Connect(location); +} + +// The "a", "b" int64 schema the canned streams serve. +inline std::shared_ptr TwoColumnSchema() { + return arrow::schema( + {arrow::field("a", arrow::int64()), arrow::field("b", arrow::int64())}); +} + +// One 3-row batch of TwoColumnSchema(), the payload the DoPut tests upload. +inline std::shared_ptr TwoColumnBatch( + const std::shared_ptr& schema = TwoColumnSchema()) { + return arrow::RecordBatch::Make(schema, 3, + {arrow::ArrayFromJSON(arrow::int64(), "[1, 2, 3]"), + arrow::ArrayFromJSON(arrow::int64(), "[10, 20, 30]")}); +} + +inline void AssertTwoColumnSchema(const std::shared_ptr& schema) { + ASSERT_NE(schema, nullptr); + ASSERT_EQ(schema->num_fields(), 2); + ASSERT_EQ(schema->field(0)->name(), "a"); + ASSERT_EQ(schema->field(1)->name(), "b"); +} + +// Drain a DoAction result stream, projecting each Result to a string. +inline arrow::Result> DrainResultBodies(ResultStream* results) { + std::vector out; + while (true) { + ARROW_ASSIGN_OR_RAISE(auto result, results->Next()); + if (result == nullptr) { + return out; + } + out.push_back(result->body->ToString()); + } +} + +// Drain a ListFlights listing, projecting each FlightInfo to its path. +inline arrow::Result>> DrainListingPaths( + FlightListing* listing) { + std::vector> out; + while (true) { + ARROW_ASSIGN_OR_RAISE(auto info, listing->Next()); + if (info == nullptr) { + return out; + } + out.push_back(info->descriptor().path); + } +} + +// Poll `predicate` until it holds or the timeout expires, the DoExchange tests +// use it to bound server-side cleanup instead of sleeping a fixed time. +template +inline bool WaitFor(Predicate predicate, + std::chrono::milliseconds timeout = std::chrono::seconds(10)) { + const auto deadline = std::chrono::steady_clock::now() + timeout; + while (std::chrono::steady_clock::now() < deadline) { + if (predicate()) { + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + return predicate(); +} + +// Upload one batch and report whichever surface first reports an error: the +// DoPut call, the write, or Close. +inline std::pair UploadOneBatch( + FlightClient* client, const FlightDescriptor& descriptor, + const std::shared_ptr& schema, + const std::shared_ptr& batch) { + auto result = client->DoPut(descriptor, schema); + if (!result.ok()) { + return {"DoPut", result.status()}; + } + auto write_status = result->writer->WriteRecordBatch(*batch); + if (!write_status.ok()) { + return {"WriteRecordBatch", write_status}; + } + return {"Close", result->writer->Close()}; +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_unary.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_unary.cc new file mode 100644 index 000000000000..fcba29fe5a98 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_unary.cc @@ -0,0 +1,124 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/flight/transport/grpc/async_grpc_service_internal.h" + +#include +#include +#include + +namespace arrow::flight::transport::grpc::detail { + +namespace { + +/** + * A generic unary reactor for handling async gRPC calls in the Flight service. + * T is the type of the response object, and PbT is the corresponding protobuf type. + */ +template +class UnaryReactor final : public AsyncReactorBase { + public: + using HandlerFn = std::function>( + const ServerCallContext&, const FlightDescriptor&)>; + + UnaryReactor(AsyncCallContext flight_context, HandlerFn handler) + : AsyncReactorBase(std::move(flight_context)), handler_(std::move(handler)) { + StartRead(&request_buf_); + } + + void OnReadDone(bool ok) override { + if (!ok) { + FinishOnce(MakeFlightError(FlightStatusCode::Internal, "Failed to read request")); + return; + } + const auto descriptor = ParseProtoRequest( + request_buf_, "FlightDescriptor"); + if (!descriptor.ok()) { + FinishOnce(descriptor.status()); + return; + } + arrow::Future> future = handler_(flight_context(), *descriptor); + future.AddCallback( + [this, token = hold()](const arrow::Result>& result) mutable { + if (!finished()) { + if (!result.ok()) { + FinishOnce(result.status()); + } else { + PbT response; + const auto status = SerializeOrNotFound(*result, &response); + if (!status.ok()) { + FinishOnce(status); + } else { + response_buf_ = MakeWriteBuffer(response); + WriteIfLive([this] { StartWrite(&response_buf_); }); + } + } + } + }); + } + + void OnWriteDone(bool ok) override { + if (finished()) { + return; + } + FinishOnce( + ok ? arrow::Status::OK() + : MakeFlightError(FlightStatusCode::Internal, "Failed to write response")); + } + + private: + HandlerFn handler_; + ::grpc::ByteBuffer request_buf_; + ::grpc::ByteBuffer response_buf_; +}; + +template +::grpc::ServerGenericBidiReactor* MakeUnaryReactor(AsyncCallContext flight_context, + HandlerFn handler) { + return new UnaryReactor(std::move(flight_context), std::move(handler)); +} + +} // namespace + +::grpc::ServerGenericBidiReactor* MakeGetFlightInfoReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) { + return MakeUnaryReactor( + std::move(flight_context), + [base](const ServerCallContext& context, const FlightDescriptor& descriptor) { + return base->GetFlightInfoAsync(context, descriptor); + }); +} + +::grpc::ServerGenericBidiReactor* MakeGetSchemaReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) { + return MakeUnaryReactor( + std::move(flight_context), + [base](const ServerCallContext& context, const FlightDescriptor& descriptor) { + return base->GetSchemaAsync(context, descriptor); + }); +} + +::grpc::ServerGenericBidiReactor* MakePollFlightInfoReactor( + AsyncCallContext flight_context, AsyncGenericFlightServerBase* base) { + return MakeUnaryReactor( + std::move(flight_context), + [base](const ServerCallContext& context, const FlightDescriptor& descriptor) { + return base->PollFlightInfoAsync(context, descriptor); + }); +} + +} // namespace arrow::flight::transport::grpc::detail diff --git a/cpp/src/arrow/flight/transport/grpc/async_grpc_unary_test.cc b/cpp/src/arrow/flight/transport/grpc/async_grpc_unary_test.cc new file mode 100644 index 000000000000..204c4b527b79 --- /dev/null +++ b/cpp/src/arrow/flight/transport/grpc/async_grpc_unary_test.cc @@ -0,0 +1,218 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// The unary RPCs served by UnaryReactor (one request in, one response out): +// GetFlightInfo, GetSchema and PollFlightInfo, plus the rules they share -- a +// null answer maps to the sync transport's "Flight not found", a server error +// propagates, an answer may complete on another thread, and a handler future +// that never resolves must not leak its reactor. + +#include "arrow/flight/transport/grpc/async_grpc_test_fixtures.h" + +namespace arrow::flight::transport::grpc { + +// --------------------------------------------------------------------------- +// The unary trio: GetFlightInfo, GetSchema and PollFlightInfo share the +// UnaryReactor (one request in, one response out), so the trio has one +// happy-path test each plus one error-path test that pins the base class. +// --------------------------------------------------------------------------- + +TEST(AsyncGrpcTest, GetFlightInfoIsServedAsync) { + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + ASSERT_OK_AND_ASSIGN(auto info, + client->GetFlightInfo(FlightDescriptor::Command("ping"))); + ipc::DictionaryMemo dict_memo; + ASSERT_OK_AND_ASSIGN(auto schema, info->GetSchema(&dict_memo)); + ASSERT_EQ(1, schema->num_fields()); + ASSERT_EQ("value", schema->field(0)->name()); + + // Cleanup before looking at the recording: it was written on a gRPC thread. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ("ping", flight_server.last_descriptor().cmd); +} + +TEST(AsyncGrpcTest, GetFlightInfoWithNullResponseIsNotFound) { + // A handler that answers OK without setting the response: the sync transport + // answers NOT_FOUND with "Flight not found" (grpc_server.cc:282-286), and the + // async reactor must answer the same, not INTERNAL. + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + const auto status = + client->GetFlightInfo(FlightDescriptor::Command("null-info")).status(); + ASSERT_RAISES(KeyError, status); + ASSERT_THAT(status.message(), ::testing::HasSubstr("Flight not found")); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, GetSchemaIsServedAsync) { + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + ASSERT_OK_AND_ASSIGN(auto schema_result, + client->GetSchema(FlightDescriptor::Command("ping"))); + ipc::DictionaryMemo dict_memo; + ASSERT_OK_AND_ASSIGN(auto schema, schema_result->GetSchema(&dict_memo)); + ASSERT_EQ(1, schema->num_fields()); + ASSERT_EQ("value", schema->field(0)->name()); + + // Cleanup before looking at the recording: it was written on a gRPC thread. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ("ping", flight_server.last_descriptor().cmd); +} + +TEST(AsyncGrpcTest, PollFlightInfoIsServedAsync) { + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + ASSERT_OK_AND_ASSIGN(auto poll_info, + client->PollFlightInfo(FlightDescriptor::Command("ping"))); + ASSERT_TRUE(poll_info->progress.has_value()); + ASSERT_DOUBLE_EQ(0.5, *poll_info->progress); + ASSERT_TRUE(poll_info->descriptor.has_value()); + ASSERT_EQ("poll-next", poll_info->descriptor->cmd); + + // Cleanup before looking at the recording: it was written on a gRPC thread. + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + ASSERT_EQ("ping", flight_server.last_descriptor().cmd); +} + +// A GetFlightInfo handler whose future never resolves: the deadline is the +// only thing that can end the RPC. +class HangingUnaryServer final : public AsyncGenericFlightServerBase { + public: + arrow::Future> GetFlightInfoAsync( + const ServerCallContext&, const FlightDescriptor&) override { + return arrow::Future>::Make(); + } +}; + +TEST(AsyncGrpcTest, GetFlightInfoDeadlineWithPendingHandlerFinishes) { + // The handler's future never resolves and the client's deadline ends the + // RPC: the reactor must be released even though the handler never completed, + // or Shutdown() below would block. + const int64_t before = detail::AsyncReactorBase::live_reactors(); + + HangingUnaryServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + FlightCallOptions impatient; + impatient.timeout = std::chrono::milliseconds(500); + auto status = + client->GetFlightInfo(impatient, FlightDescriptor::Command("hang")).status(); + EXPECT_FALSE(status.ok()) << "the deadline must expire while the handler hangs"; + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); + + EXPECT_TRUE(WaitFor([&] { + return detail::AsyncReactorBase::live_reactors() == before; + })) << "a reactor leaked: still alive after the server stopped (before=" + << before << ", now=" << detail::AsyncReactorBase::live_reactors() << ")"; +} + +TEST(AsyncGrpcTest, UnaryReactorPropagatesServerError) { + TestFlightServer flight_server; + FlightServerOptions options(LocalhostLocation()); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN(auto client, ConnectTo(flight_server.port())); + + // The server's KeyError must reach the client as a status failure carrying + // the server's message, not as a successful (empty) response. The message is + // matched by substring because this build defines ARROW_EXTRA_ERROR_CONTEXT, + // which appends the call-site trail to the server's message. + auto status = client->GetFlightInfo(FlightDescriptor::Command("fail")).status(); + ASSERT_FALSE(status.ok()); + ASSERT_EQ(arrow::StatusCode::KeyError, status.code()); + ASSERT_THAT(status.message(), ::testing::HasSubstr("no such flight")); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +TEST(AsyncGrpcTest, GetFlightInfoAsyncCompletesOnAnotherThread) { + // The handler completes its future from a helper thread after the dispatch + // has returned: the answer must arrive on the client anyway, which is the + // whole point of the async handler. + class DeferredInfoServer : public AsyncGenericFlightServerBase { + public: + arrow::Future> GetFlightInfoAsync( + const ServerCallContext&, const FlightDescriptor& request) override { + auto out = arrow::Future>::Make(); + std::thread([out, request]() mutable { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + auto info = + FlightInfo::Make(*arrow::schema({arrow::field("value", arrow::int64())}), + request, std::vector{}, -1, -1); + if (!info.ok()) { + out.MarkFinished(info.status()); + return; + } + out.MarkFinished(std::make_shared(std::move(*info))); + }).detach(); + return out; + } + }; + + ASSERT_OK_AND_ASSIGN(auto location, Location::Parse("grpc://localhost:0")); + DeferredInfoServer flight_server; + FlightServerOptions options(location); + ASSERT_OK(flight_server.Init(options)); + ASSERT_OK_AND_ASSIGN( + auto client_location, + Location::Parse("grpc://localhost:" + std::to_string(flight_server.port()))); + ASSERT_OK_AND_ASSIGN(auto client, FlightClient::Connect(client_location)); + + ASSERT_OK_AND_ASSIGN(auto info, + client->GetFlightInfo(FlightDescriptor::Command("deferred"))); + ipc::DictionaryMemo dict_memo; + ASSERT_OK_AND_ASSIGN(auto schema, info->GetSchema(&dict_memo)); + ASSERT_EQ(1, schema->num_fields()); + ASSERT_EQ("value", schema->field(0)->name()); + + ASSERT_OK(client->Close()); + ASSERT_OK(flight_server.Shutdown()); + ASSERT_OK(flight_server.Wait()); +} + +} // namespace arrow::flight::transport::grpc diff --git a/cpp/src/arrow/flight/transport/grpc/customize_grpc.h b/cpp/src/arrow/flight/transport/grpc/customize_grpc.h index 4a64e4e07214..1974111709b4 100644 --- a/cpp/src/arrow/flight/transport/grpc/customize_grpc.h +++ b/cpp/src/arrow/flight/transport/grpc/customize_grpc.h @@ -105,8 +105,8 @@ namespace grpc { // Those two functions are defined in serialization_internal.cc // Write FlightData to a grpc::ByteBuffer without extra copying -::grpc::Status FlightDataSerialize(const arrow::flight::FlightPayload& msg, - ::grpc::ByteBuffer* out, bool* own_buffer); +ARROW_FLIGHT_EXPORT ::grpc::Status FlightDataSerialize( + const arrow::flight::FlightPayload& msg, ::grpc::ByteBuffer* out, bool* own_buffer); // Read internal::FlightData from grpc::ByteBuffer containing FlightData // protobuf without copying diff --git a/cpp/src/arrow/flight/transport/grpc/grpc_server.cc b/cpp/src/arrow/flight/transport/grpc/grpc_server.cc index 2a100caeff76..507bef4693f7 100644 --- a/cpp/src/arrow/flight/transport/grpc/grpc_server.cc +++ b/cpp/src/arrow/flight/transport/grpc/grpc_server.cc @@ -25,6 +25,10 @@ #include #include +#include "arrow/flight/transport/grpc/async_grpc_compat.h" +#if ARROW_FLIGHT_HAS_ASYNC_SERVER +# include "arrow/flight/transport/grpc/async_grpc_service.h" +#endif #include "arrow/flight/transport/grpc/customize_grpc.h" #include @@ -113,7 +117,7 @@ class GrpcServerAuthSender : public ServerAuthSender { ::grpc::ServerReaderWriter* stream_; }; -using GrpcServerCallContext = transport::grpc::GrpcServerCallContext; +using SyncServerCallContext = transport::grpc::GrpcServerCallContext; using GrpcContextHelper = transport::grpc::GrpcServerCallContextHelper; // A ServerDataStream for streaming data to the client. @@ -224,7 +228,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status Handshake( ServerContext* context, ::grpc::ServerReaderWriter* stream) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.MakeCallContext(FlightMethod::Handshake, context, &flight_context)); if (!helper_.auth_handler()) { @@ -242,7 +246,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status ListFlights(ServerContext* context, const pb::Criteria* request, ServerWriter* writer) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::ListFlights, context, &flight_context)); // Retrieve the listing from the implementation @@ -265,7 +269,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status GetFlightInfo(ServerContext* context, const pb::FlightDescriptor* request, pb::FlightInfo* response) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::GetFlightInfo, context, &flight_context)); @@ -292,7 +296,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status PollFlightInfo(ServerContext* context, const pb::FlightDescriptor* request, pb::PollInfo* response) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::PollFlightInfo, context, &flight_context)); @@ -317,7 +321,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status GetSchema(ServerContext* context, const pb::FlightDescriptor* request, pb::SchemaResult* response) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::GetSchema, context, &flight_context)); @@ -342,7 +346,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status DoGet(ServerContext* context, const pb::Ticket* request, ServerWriter* writer) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::DoGet, context, &flight_context)); CHECK_ARG_NOT_NULL(flight_context, request, "ticket cannot be null"); @@ -358,7 +362,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status DoPut( ServerContext* context, ::grpc::ServerReaderWriter* reader) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::DoPut, context, &flight_context)); PutDataStream stream(reader); @@ -368,7 +372,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status DoExchange( ServerContext* context, ::grpc::ServerReaderWriter* stream) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::DoExchange, context, &flight_context)); ExchangeDataStream data_stream(stream); @@ -378,7 +382,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status ListActions(ServerContext* context, const pb::Empty* request, ServerWriter* writer) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::ListActions, context, &flight_context)); // Retrieve the listing from the implementation @@ -390,7 +394,7 @@ class GrpcServiceHandler final : public FlightService::Service { ::grpc::Status DoAction(ServerContext* context, const pb::Action* request, ServerWriter* writer) { - GrpcServerCallContext flight_context(context); + SyncServerCallContext flight_context(context); GRPC_RETURN_NOT_GRPC_OK( helper_.CheckAuth(FlightMethod::DoAction, context, &flight_context)); CHECK_ARG_NOT_NULL(flight_context, request, "Action cannot be null"); @@ -430,7 +434,16 @@ class GrpcServiceHandler final : public FlightService::Service { // The ServerTransport implementation for gRPC. Manages the gRPC server itself. class GrpcServerTransport : public internal::ServerTransport { public: - using internal::ServerTransport::ServerTransport; + GrpcServerTransport(FlightServerBase* base, + std::shared_ptr memory_manager, + AsyncGenericFlightServerBase* async_base = nullptr) + : internal::ServerTransport(base, std::move(memory_manager)) { +#if ARROW_FLIGHT_HAS_ASYNC_SERVER + async_base_ = async_base; +#else + ARROW_UNUSED(async_base); +#endif + } static arrow::Result> Make( FlightServerBase* base, std::shared_ptr memory_manager) { @@ -438,9 +451,56 @@ class GrpcServerTransport : public internal::ServerTransport { new GrpcServerTransport(base, std::move(memory_manager))); } + /// \brief Like Make(), for a server that is an AsyncGenericFlightServerBase. + /// + /// This is the async mode: the transport serves the async server's handlers + /// through the generic callback service, with no FlightServerBase involved. + static arrow::Result> MakeAsync( + AsyncGenericFlightServerBase* async_base, + std::shared_ptr memory_manager) { +#if ARROW_FLIGHT_HAS_ASYNC_SERVER + return std::unique_ptr( + new GrpcServerTransport(/*base=*/nullptr, std::move(memory_manager), async_base)); +#else + ARROW_UNUSED(async_base); + ARROW_UNUSED(memory_manager); + return Status::NotImplemented( + "The async Flight server requires gRPC's generic callback API " + "(arrow was built against gRPC < 1.65)"); +#endif + } + Status Init(const FlightServerOptions& options, const arrow::util::Uri& uri) override { - grpc_service_.reset( - new GrpcServiceHandler(options.auth_handler, options.middleware, this)); +#if ARROW_FLIGHT_HAS_ASYNC_SERVER + if (async_base_ != nullptr) { + if (options.auth_handler) { + return Status::NotImplemented( + "AsyncGenericFlightServerBase does not support a blocking auth handler; " + "override Handshake/ValidateToken instead"); + } + using AsyncHelper = GrpcServerCallContextHelper<::grpc::CallbackServerContext>; + AsyncHelper::HandshakeFn handshake = + [base = async_base_](const ServerCallContext& context, + const std::string& request, std::string* response) { + return base->Handshake(context, request, response); + }; + AsyncHelper::ValidateTokenFn validate_token = + [base = async_base_](const ServerCallContext& context, const std::string& token, + std::string* peer_identity) { + return base->ValidateToken(context, token, peer_identity); + }; + async_helper_ = std::make_shared( + /*auth_handler=*/nullptr, options.middleware, std::move(validate_token)); + async_service_ = std::make_unique( + async_base_, async_helper_, std::move(handshake)); + // The else and its body are split by #endif, so cpplint cannot pair the + // braces; the #if branch above opens the block it closes. + } else // NOLINT(readability/braces) +#endif + { + grpc_service_.reset( + new GrpcServiceHandler(options.auth_handler, options.middleware, this)); + } ::grpc::ServerBuilder builder; int port = 0; @@ -452,7 +512,14 @@ class GrpcServerTransport : public internal::ServerTransport { builder.AddListeningPort(endpoint.address, endpoint.credentials, &port); } - builder.RegisterService(grpc_service_.get()); +#if ARROW_FLIGHT_HAS_ASYNC_SERVER + if (async_service_) { + builder.RegisterCallbackGenericService(async_service_.get()); + } else // NOLINT(readability/braces) +#endif + { + builder.RegisterService(grpc_service_.get()); + } ConfigureServerBuilderOptions(options, &builder); grpc_server_ = builder.BuildAndStart(); @@ -484,6 +551,12 @@ class GrpcServerTransport : public internal::ServerTransport { private: std::unique_ptr grpc_service_; +#if ARROW_FLIGHT_HAS_ASYNC_SERVER + AsyncGenericFlightServerBase* async_base_ = nullptr; + std::shared_ptr> + async_helper_; + std::unique_ptr async_service_; +#endif std::unique_ptr<::grpc::Server> grpc_server_; Location location_; }; @@ -496,6 +569,11 @@ void InitializeFlightGrpcServer() { auto* registry = flight::internal::GetDefaultTransportRegistry(); for (const auto& transport : {"grpc", "grpc+tls", "grpc+tcp", "grpc+unix"}) { ARROW_CHECK_OK(registry->RegisterServer(transport, GrpcServerTransport::Make)); + // Registered even when gRPC is too old for the async server: MakeAsync then + // answers NotImplemented with the reason, instead of Init() failing with the + // generic "no async server transport implementation" error. + ARROW_CHECK_OK( + registry->RegisterAsyncServer(transport, GrpcServerTransport::MakeAsync)); } }); } diff --git a/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.h b/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.h index cdafb53a3e12..7964ce87df22 100644 --- a/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.h +++ b/cpp/src/arrow/flight/transport/grpc/grpc_server_internal.h @@ -17,6 +17,7 @@ #pragma once +#include #include #include #include @@ -26,7 +27,9 @@ #include #include "arrow/flight/server.h" +#include "arrow/flight/server_auth.h" #include "arrow/flight/server_middleware.h" +#include "arrow/flight/transport/grpc/async_grpc_compat.h" #include "arrow/flight/transport/grpc/util_internal.h" #include "arrow/util/uri.h" @@ -112,9 +115,24 @@ class GrpcServerCallContext : public ServerCallContext { template class GrpcServerCallContextHelper { public: + /// A token validator for servers that have no ServerAuthHandler: the async + /// server plugs AsyncGenericFlightServerBase::ValidateToken in here. + using ValidateTokenFn = + std::function; + + /// The async handshake hook: the transport reads one Handshake request, + /// calls this, then writes the response it filled in. + using HandshakeFn = + std::function; + GrpcServerCallContextHelper(std::shared_ptr auth_handler, - MiddlewareFactoryList middleware) - : auth_handler_(std::move(auth_handler)), middleware_(std::move(middleware)) {} + MiddlewareFactoryList middleware, + ValidateTokenFn validate_token = {}) + : auth_handler_(std::move(auth_handler)), + middleware_(std::move(middleware)), + validate_token_(std::move(validate_token)) {} // Authenticate the client (if applicable) and construct the call context ::grpc::Status MakeCallContext(FlightMethod method, GrpcContext* context, @@ -165,7 +183,11 @@ class GrpcServerCallContextHelper { } else { flight_context->peer_identity_ = ""; } - } else { + if (!validate_token_) { + return MakeCallContext(method, context, flight_context); + } + } + { const auto client_metadata = context->client_metadata(); const auto [auth_header, auth_header_end] = client_metadata.equal_range(kGrpcAuthHeader); @@ -174,7 +196,10 @@ class GrpcServerCallContextHelper { token = std::string(auth_header->second.data(), auth_header->second.length()); } auto auth_status = - auth_handler_->IsValid(*flight_context, token, &flight_context->peer_identity_); + validate_token_ + ? validate_token_(*flight_context, token, &flight_context->peer_identity_) + : auth_handler_->IsValid(*flight_context, token, + &flight_context->peer_identity_); if (!auth_status.ok()) { return flight_context->FinishRequest(auth_status); } @@ -187,6 +212,7 @@ class GrpcServerCallContextHelper { private: std::shared_ptr auth_handler_; MiddlewareFactoryList middleware_; + ValidateTokenFn validate_token_; }; struct GrpcServerEndpoint { diff --git a/cpp/src/arrow/flight/transport/grpc/serialization_internal.cc b/cpp/src/arrow/flight/transport/grpc/serialization_internal.cc index ee13a8e2022f..a2ac73c70dc0 100644 --- a/cpp/src/arrow/flight/transport/grpc/serialization_internal.cc +++ b/cpp/src/arrow/flight/transport/grpc/serialization_internal.cc @@ -17,6 +17,10 @@ #include "arrow/flight/transport/grpc/serialization_internal.h" +#include + +#include "arrow/flight/types.h" + #include #include #include @@ -125,7 +129,12 @@ class GrpcBuffer : public MutableBuffer { grpc_slice slice_; }; -// Destructor callback for grpc::Slice +} // namespace + +// Destructor callback for grpc::Slice. Declared here (rather than in the +// header) so -Wmissing-declarations is happy. +void ReleaseBuffer(void* buf_ptr); + void ReleaseBuffer(void* buf_ptr) { delete reinterpret_cast*>(buf_ptr); } @@ -149,7 +158,24 @@ arrow::Result<::grpc::Slice> SliceFromBuffer(const std::shared_ptr& buf) return slice; } -} // namespace +// Initialize arrow Buffer from gRPC ByteBuffer +Status WrapGrpcBuffer(::grpc::ByteBuffer* cpp_buf, std::shared_ptr* out) { + return GrpcBuffer::Wrap(cpp_buf, out); +} + +arrow::Result BytesFromBuffer(const ::grpc::ByteBuffer& buffer, + std::string_view error_message) { + std::vector<::grpc::Slice> slices; + if (!buffer.Dump(&slices).ok()) { + return MakeFlightError(FlightStatusCode::Internal, std::string(error_message)); + } + std::string bytes; + bytes.reserve(buffer.Length()); + for (const auto& slice : slices) { + bytes.append(reinterpret_cast(slice.begin()), slice.size()); + } + return bytes; +} ::grpc::Status FlightDataSerialize(const FlightPayload& msg, ByteBuffer* out, bool* own_buffer) { diff --git a/cpp/src/arrow/flight/transport/grpc/serialization_internal.h b/cpp/src/arrow/flight/transport/grpc/serialization_internal.h index 5c347fd4f819..36285e0bdec8 100644 --- a/cpp/src/arrow/flight/transport/grpc/serialization_internal.h +++ b/cpp/src/arrow/flight/transport/grpc/serialization_internal.h @@ -21,6 +21,8 @@ #pragma once #include +#include +#include #include "arrow/flight/protocol_internal.h" #include "arrow/flight/transport/grpc/protocol_grpc_internal.h" @@ -28,12 +30,40 @@ #include "arrow/result.h" namespace arrow { + +class Buffer; + namespace flight { namespace transport { namespace grpc { namespace pb = arrow::flight::protocol; +/// Zero-copy conversion of an Arrow Buffer to a gRPC Slice. +/// +/// The Buffer lifetime is tied to the Slice (a shared_ptr is attached to the +/// slice's destructor callback) +/// \param[in] buf The Arrow Buffer to convert. +/// \return A gRPC Slice that shares ownership of the Buffer. +arrow::Result<::grpc::Slice> SliceFromBuffer(const std::shared_ptr& buf); + +/// Zero-copy conversion of a gRPC ByteBuffer into an Arrow Buffer. +/// +/// The Arrow Buffer owns a reference to the gRPC slice, so it stays valid +/// after the ByteBuffer goes away. +/// \param[in] cpp_buf The gRPC ByteBuffer to wrap. +/// \param[out] out The resulting Arrow Buffer that shares ownership of the gRPC slices. +/// \return Status indicating success or failure. +Status WrapGrpcBuffer(::grpc::ByteBuffer* cpp_buf, std::shared_ptr* out); + +/// Copy a gRPC ByteBuffer's slices into one contiguous string. +/// \param[in] buffer The gRPC ByteBuffer to copy from. +/// \param[in] error_message The error message to use if the buffer cannot be dumped. +/// \return The buffer's bytes, or an IOError carrying FlightStatusCode::Internal +/// on failure. +ARROW_FLIGHT_EXPORT arrow::Result BytesFromBuffer( + const ::grpc::ByteBuffer& buffer, std::string_view error_message); + /// Write Flight message on gRPC stream with zero-copy optimizations. // Returns Invalid if the payload is ill-formed // Returns true if the payload was written, false if it was not diff --git a/cpp/src/arrow/flight/transport/grpc/util_internal.cc b/cpp/src/arrow/flight/transport/grpc/util_internal.cc index 95b205d67808..8dbdf3eb1952 100644 --- a/cpp/src/arrow/flight/transport/grpc/util_internal.cc +++ b/cpp/src/arrow/flight/transport/grpc/util_internal.cc @@ -310,9 +310,10 @@ static ::grpc::Status ToRawGrpcStatus(const Status& arrow_status) { return {grpc_code, std::move(transport_status.message)}; } -/// Convert an Arrow status to a gRPC status, and add extra headers to -/// the response to encode the original Arrow status. -::grpc::Status ToGrpcStatus(const Status& arrow_status, ::grpc::ServerContext* ctx) { +namespace { +// Shared body: convert and attach the Arrow status as trailing metadata. +template +::grpc::Status ToGrpcStatusWithContext(const Status& arrow_status, Context* ctx) { ::grpc::Status status = ToRawGrpcStatus(arrow_status); if (!status.ok() && ctx) { const std::string code = ToChars(static_cast(arrow_status.code())); @@ -327,9 +328,23 @@ ::grpc::Status ToGrpcStatus(const Status& arrow_status, ::grpc::ServerContext* c ctx->AddTrailingMetadata(kBinaryErrorDetailsKey, fsd->extra_info()); } } - return status; } +} // namespace + +/// Convert an Arrow status to a gRPC status, and add extra headers to +/// the response to encode the original Arrow status. +::grpc::Status ToGrpcStatus(const Status& arrow_status, ::grpc::ServerContext* ctx) { + return ToGrpcStatusWithContext(arrow_status, ctx); +} + +#if GRPC_CPP_VERSION_CHECK(1, 65, 0) +/// Same, for the callback API's server context. +::grpc::Status ToGrpcStatus(const Status& arrow_status, + ::grpc::CallbackServerContext* ctx) { + return ToGrpcStatusWithContext(arrow_status, ctx); +} +#endif #if GRPC_CPP_VERSION_CHECK(1, 80, 0) Status FromAbslStatus(const ::absl::Status& absl_status) { diff --git a/cpp/src/arrow/flight/transport/grpc/util_internal.h b/cpp/src/arrow/flight/transport/grpc/util_internal.h index 7e306d574a58..5d215620221c 100644 --- a/cpp/src/arrow/flight/transport/grpc/util_internal.h +++ b/cpp/src/arrow/flight/transport/grpc/util_internal.h @@ -42,6 +42,7 @@ namespace grpc { class Status; +class CallbackServerContext; } // namespace grpc @@ -107,6 +108,12 @@ ARROW_FLIGHT_EXPORT ::grpc::Status ToGrpcStatus(const Status& arrow_status, ::grpc::ServerContext* ctx = nullptr); +#if GRPC_CPP_VERSION_CHECK(1, 65, 0) +ARROW_FLIGHT_EXPORT +::grpc::Status ToGrpcStatus(const Status& arrow_status, + ::grpc::CallbackServerContext* ctx); +#endif + // gRPC 1.80.0 or later use absl::Status. #if GRPC_CPP_VERSION_CHECK(1, 80, 0) /// Convert an Abseil status to an Arrow status. diff --git a/cpp/src/arrow/flight/transport_server_internal.cc b/cpp/src/arrow/flight/transport_server_internal.cc index be159afa5d9e..9a5afa19f46b 100644 --- a/cpp/src/arrow/flight/transport_server_internal.cc +++ b/cpp/src/arrow/flight/transport_server_internal.cc @@ -298,6 +298,62 @@ void ServerSignalState::WaitForSignals( } } +ServerLifecycle::ServerLifecycle(const char* server_name) : server_name_(server_name) {} + +ServerLifecycle::~ServerLifecycle() = default; + +Status ServerLifecycle::Init(const FlightServerOptions& options, + std::unique_ptr transport) { + ARROW_ASSIGN_OR_RAISE(auto uri, options.location.uri()); + transport_ = std::move(transport); + return transport_->Init(options, uri); +} + +int ServerLifecycle::port() const { return location().port(); } + +Location ServerLifecycle::location() const { return transport_->location(); } + +Status ServerLifecycle::SetShutdownOnSignals(const std::vector& signals) { + return signal_state_.SetShutdownOnSignals(signals); +} + +Status ServerLifecycle::Serve() { + if (!transport_) { + return Status::UnknownError("Server did not start properly"); + } + return signal_state_.Serve( + [this]() -> Status { return transport_->Wait(); }, + [this](const std::chrono::system_clock::time_point* deadline) -> Status { + return Shutdown(deadline); + }, + "Error shutting down server"); +} + +int ServerLifecycle::GotSignal() const { return signal_state_.GotSignal(); } + +Status ServerLifecycle::Shutdown(const std::chrono::system_clock::time_point* deadline) { + if (!transport_) { + return Status::Invalid("Shutdown() on uninitialized ", server_name_); + } + return signal_state_.Shutdown( + [this](const std::chrono::system_clock::time_point* maybe_deadline) -> Status { + if (maybe_deadline) { + return transport_->Shutdown(*maybe_deadline); + } + return transport_->Shutdown(); + }, + deadline); +} + +Status ServerLifecycle::Wait() { + return signal_state_.Wait([this] { + if (!transport_) { + return Status::Invalid("Wait() on uninitialized ", server_name_); + } + return transport_->Wait(); + }); +} + Status WriteDataStream(std::unique_ptr data_stream, ServerDataStream* stream) { if (!data_stream) { diff --git a/cpp/src/arrow/flight/transport_server_internal.h b/cpp/src/arrow/flight/transport_server_internal.h index 2112d613c6b1..d1f28319787e 100644 --- a/cpp/src/arrow/flight/transport_server_internal.h +++ b/cpp/src/arrow/flight/transport_server_internal.h @@ -37,6 +37,7 @@ namespace arrow::flight::internal { class ServerDataStream; class PeekableFlightDataReader; +class ServerTransport; /// \brief Adapt TransportDataStream to the FlightMessageReader /// interface for DoPut. @@ -302,4 +303,39 @@ class ARROW_FLIGHT_EXPORT ServerSignalState { const char* shutdown_warning_ = nullptr; }; +/// \brief The server lifecycle shared by FlightServerBase and +/// AsyncGenericFlightServerBase. +/// +/// Owns the server's transport and signal state and implements Init(), the +/// signal-aware Serve()/Shutdown()/Wait() trio and the location accessors, so +/// both server classes have one implementation of them. +class ServerLifecycle { + public: + /// \param[in] server_name names the server class. Used in error logs. + explicit ServerLifecycle(const char* server_name); + ~ServerLifecycle(); + + /// \brief Take ownership of \p transport and initialize it. + /// + /// \param[in] options the server options the transport is initialized with. + /// \param[in] transport the transport serving the server, created by the + /// server class (the sync and async servers come from different registry + /// factories). + Status Init(const FlightServerOptions& options, + std::unique_ptr transport); + + int port() const; + Location location() const; + Status SetShutdownOnSignals(const std::vector& signals); + Status Serve(); + int GotSignal() const; + Status Shutdown(const std::chrono::system_clock::time_point* deadline); + Status Wait(); + + private: + const char* server_name_; + std::unique_ptr transport_; + ServerSignalState signal_state_; +}; + } // namespace arrow::flight::internal diff --git a/cpp/src/arrow/flight/types.cc b/cpp/src/arrow/flight/types.cc index c4b81184cc67..08e891557b42 100644 --- a/cpp/src/arrow/flight/types.cc +++ b/cpp/src/arrow/flight/types.cc @@ -1052,6 +1052,9 @@ Status ResultStream::Drain() { } FlightStreamChunk::FlightStreamChunk() noexcept = default; +FlightStreamChunk::FlightStreamChunk(std::shared_ptr data, + std::shared_ptr app_metadata) noexcept + : data(std::move(data)), app_metadata(std::move(app_metadata)) {} FlightStreamChunk::~FlightStreamChunk() = default; arrow::Result>> diff --git a/cpp/src/arrow/flight/types.h b/cpp/src/arrow/flight/types.h index dd18c0e0786a..d3049a5e8eed 100644 --- a/cpp/src/arrow/flight/types.h +++ b/cpp/src/arrow/flight/types.h @@ -1166,6 +1166,8 @@ class ARROW_FLIGHT_EXPORT ResultStream { struct ARROW_FLIGHT_EXPORT FlightStreamChunk { public: FlightStreamChunk() noexcept; + FlightStreamChunk(std::shared_ptr data, + std::shared_ptr app_metadata) noexcept; ~FlightStreamChunk(); std::shared_ptr data; diff --git a/cpp/src/arrow/util/windows_fixup.h b/cpp/src/arrow/util/windows_fixup.h index 42e74f4a7857..03dc238820cb 100644 --- a/cpp/src/arrow/util/windows_fixup.h +++ b/cpp/src/arrow/util/windows_fixup.h @@ -38,14 +38,6 @@ # undef DeleteFile # endif -// Other annoying Windows macro definitions... -# ifdef IN -# undef IN -# endif -# ifdef OUT -# undef OUT -# endif - // Note that we can't undefine OPTIONAL, because it can be used in other // Windows headers...