Copilot commented on code in PR #55419:
URL: https://github.com/apache/arrow/pull/55419#discussion_r4511764175
##########
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
Review Comment:
Both newly registered examples unconditionally include `<unistd.h>` and call
`getpid()`, but these targets are enabled on Windows as well
(`ADD_ARROW_EXAMPLE` has no platform filter). Enabling Flight examples with
MSVC now fails to compile. Use the Windows process-ID API behind a platform
guard in both sources, or restrict these targets to supported platforms.
##########
cpp/src/arrow/flight/transport/grpc/async_grpc_exchange.cc:
##########
@@ -0,0 +1,128 @@
+// 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<ExchangeReader>()),
+ writer_(std::make_shared<ExchangeWriter>(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();
Review Comment:
Cancellation can run concurrently with `OnReadDone()` or `OnWriteDone()`
after those callbacks have passed their `finished()` checks. These teardown
calls then mutate the same reader/writer state without synchronization: for
example, a read completion and `OnCancelled()` can both observe `pending_` as
unfinished and both call `MarkFinished()`, which Arrow futures do not permit.
Serialize cancellation with normal reader/writer state transitions, and move
futures out under that synchronization before completing them outside it to
allow reentrant continuations.
##########
cpp/src/arrow/flight/transport/grpc/async_grpc_do_get.cc:
##########
@@ -0,0 +1,166 @@
+// 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 <atomic>
+#include <memory>
+#include <optional>
+#include <utility>
+
+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<std::shared_ptr<AsyncFlightDataStream>> future =
+ base_->DoGetAsync(flight_context(), *ticket);
+ future.AddCallback(
+ [this, token = hold()](
+ const arrow::Result<std::shared_ptr<AsyncFlightDataStream>>&
result) {
+ if (!finished()) {
+ if (!result.ok()) {
+ FinishOnce(result.status());
+ } else {
+ async_data_stream_ = result.ValueOrDie();
Review Comment:
Publishing this shared pointer races with `OnCancel()` reading it in
`CloseStreamOnce()`. There is also a lost-close interleaving: the future
callback passes `finished()`, cancellation sets `closed_` while the stream is
still null, and then this assignment installs the stream; `WriteNextPayload()`
returns because the RPC is finished, so that stream never receives `Close()`.
Synchronize stream publication with cancellation so the result is either
rejected before adoption or adopted and closed exactly once.
##########
cpp/src/arrow/flight/flight_data_decoder.cc:
##########
@@ -0,0 +1,320 @@
+// 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 <algorithm>
+#include <deque>
+
+#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<ipc::Message> message, std::shared_ptr<Buffer>
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<Buffer> 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<std::unique_ptr<ipc::Message>> 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<Buffer> 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<ipc::Message> message,
+ std::shared_ptr<Buffer> app_metadata)
+ : message(std::move(message)), app_metadata(std::move(app_metadata)) {}
+ std::unique_ptr<ipc::Message> message;
+ std::shared_ptr<Buffer> app_metadata;
+ };
+
+ std::deque<QueuedMessage> queue_;
+ std::shared_ptr<Buffer> 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<std::unique_ptr<ipc::Message>> ReadNextMessage() override {
+ return reader_->ReadNextMessage();
+ }
+
+ private:
+ FlightDataMessageReader* reader_;
+};
+
+} // namespace
+
+class AsyncFlightMessageDecoder::AsyncFlightMessageDecoderImpl {
+ public:
+ AsyncFlightMessageDecoderImpl(std::shared_ptr<AsyncFlightDataListener>
listener,
+ ipc::IpcReadOptions options)
+ : listener_(std::move(listener)),
+ options_(std::move(options)),
+ message_reader_(std::make_unique<FlightDataMessageReader>()) {}
+
+ /// \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 {
+ // The descriptor of an upload arrives with the first message of the
+ // upload, which also carries the schema: report the descriptor first,
and
+ // decode the rest of that message once the listener's future resolves.
+ return listener_->OnDescriptor(*data.descriptor)
+ .Then([this, data = std::move(data)]() mutable {
+ return ConsumeDataAfterDescriptor(std::move(data));
+ });
+ }
+ }
+
+ std::shared_ptr<Schema> 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<FlightDataMessageReaderRef>(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<RecordBatch> 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<AsyncFlightDataListener> listener_;
+ ipc::IpcReadOptions options_;
+ // Kept here, not transferred to the RecordBatchStreamReader: the decoder
+ // reads the app metadata back out of it, and a failed Open() must leave it
+ // usable.
+ std::unique_ptr<FlightDataMessageReader> message_reader_;
+ std::shared_ptr<ipc::RecordBatchStreamReader> batch_reader_;
+};
+
+AsyncFlightMessageDecoder::AsyncFlightMessageDecoder(
+ std::shared_ptr<AsyncFlightDataListener> listener, ipc::IpcReadOptions
options)
+ :
impl_(std::make_unique<AsyncFlightMessageDecoderImpl>(std::move(listener),
+
std::move(options))) {}
+
+AsyncFlightMessageDecoder::~AsyncFlightMessageDecoder() = default;
+
+Future<> AsyncFlightMessageDecoder::Consume(std::shared_ptr<Buffer> 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<Schema> 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) {
+ 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<AsyncFlightDataListener>& listener,
+ FlightDataListenerTransport* transport) {
+ if (listener == nullptr) {
+ return;
+ }
+ std::lock_guard<std::mutex> lock(listener->transport_mutex_);
+ listener->transport_ = transport;
+ if (listener->writer_ != nullptr) {
+ std::lock_guard<std::mutex> writer_lock(listener->writer_->writer_mutex_);
+ listener->writer_->transport_ = transport;
+ }
+}
+
+void FlightDataListenerTransport::Clear(
+ const std::shared_ptr<AsyncFlightDataListener>& listener) {
+ if (listener == nullptr) {
+ return;
+ }
+ std::lock_guard<std::mutex> lock(listener->transport_mutex_);
+ listener->transport_ = nullptr;
+ if (listener->writer_ != nullptr) {
+ std::lock_guard<std::mutex> writer_lock(listener->writer_->writer_mutex_);
+ listener->writer_->transport_ = nullptr;
+ }
+}
+
+Status FlightDataListenerTransport::ReportFinish(
+ const std::shared_ptr<AsyncFlightDataListener>& listener, Status status) {
+ if (listener == nullptr) {
+ return Status::OK();
+ }
+ // Cancel() holds the listener's transport lock across its call into the
transport, and
+ // the ending it triggers reports from inside that call.
+ bool expected = false;
+ if (!listener->finished_.compare_exchange_strong(expected, true)) {
+ return Status::OK();
+ }
+ // OnFinish is application code, it may call back into the listener (as
+ // Cancel() does), which is why nothing is held here.
+ return listener->OnFinish(std::move(status)).status();
+}
+
+} // namespace internal
+
+Future<> AsyncFlightDataListener::Cancel(Status status) {
+ std::lock_guard<std::mutex> lock(transport_mutex_);
+ 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_->CancelUpload(std::move(status));
Review Comment:
`Cancel()` still holds `transport_mutex_` when this invokes
`CancelUpload()`, which synchronously calls `FinishUpload()` → `ReportFinish()`
→ the application's `OnFinish()`. If `OnFinish()` calls `Cancel()` (allowed by
the any-thread/any-time contract), it attempts to lock the same non-recursive
mutex and deadlocks. Pin the transport's lifetime under the mutex, then release
the mutex before invoking transport/application callbacks.
##########
cpp/src/arrow/flight/flight_data_decoder.cc:
##########
@@ -0,0 +1,320 @@
+// 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 <algorithm>
+#include <deque>
+
+#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<ipc::Message> message, std::shared_ptr<Buffer>
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<Buffer> 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<std::unique_ptr<ipc::Message>> 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<Buffer> 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<ipc::Message> message,
+ std::shared_ptr<Buffer> app_metadata)
+ : message(std::move(message)), app_metadata(std::move(app_metadata)) {}
+ std::unique_ptr<ipc::Message> message;
+ std::shared_ptr<Buffer> app_metadata;
+ };
+
+ std::deque<QueuedMessage> queue_;
+ std::shared_ptr<Buffer> 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<std::unique_ptr<ipc::Message>> ReadNextMessage() override {
+ return reader_->ReadNextMessage();
+ }
+
+ private:
+ FlightDataMessageReader* reader_;
+};
+
+} // namespace
+
+class AsyncFlightMessageDecoder::AsyncFlightMessageDecoderImpl {
+ public:
+ AsyncFlightMessageDecoderImpl(std::shared_ptr<AsyncFlightDataListener>
listener,
+ ipc::IpcReadOptions options)
+ : listener_(std::move(listener)),
+ options_(std::move(options)),
+ message_reader_(std::make_unique<FlightDataMessageReader>()) {}
+
+ /// \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 {
+ // The descriptor of an upload arrives with the first message of the
+ // upload, which also carries the schema: report the descriptor first,
and
+ // decode the rest of that message once the listener's future resolves.
+ return listener_->OnDescriptor(*data.descriptor)
+ .Then([this, data = std::move(data)]() mutable {
+ return ConsumeDataAfterDescriptor(std::move(data));
+ });
+ }
+ }
+
+ std::shared_ptr<Schema> 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<FlightDataMessageReaderRef>(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<RecordBatch> 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<AsyncFlightDataListener> listener_;
+ ipc::IpcReadOptions options_;
+ // Kept here, not transferred to the RecordBatchStreamReader: the decoder
+ // reads the app metadata back out of it, and a failed Open() must leave it
+ // usable.
+ std::unique_ptr<FlightDataMessageReader> message_reader_;
+ std::shared_ptr<ipc::RecordBatchStreamReader> batch_reader_;
+};
+
+AsyncFlightMessageDecoder::AsyncFlightMessageDecoder(
+ std::shared_ptr<AsyncFlightDataListener> listener, ipc::IpcReadOptions
options)
+ :
impl_(std::make_unique<AsyncFlightMessageDecoderImpl>(std::move(listener),
+
std::move(options))) {}
+
+AsyncFlightMessageDecoder::~AsyncFlightMessageDecoder() = default;
+
+Future<> AsyncFlightMessageDecoder::Consume(std::shared_ptr<Buffer> 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<Schema> 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) {
+ 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<AsyncFlightDataListener>& listener,
+ FlightDataListenerTransport* transport) {
+ if (listener == nullptr) {
+ return;
+ }
+ std::lock_guard<std::mutex> lock(listener->transport_mutex_);
+ listener->transport_ = transport;
+ if (listener->writer_ != nullptr) {
+ std::lock_guard<std::mutex> writer_lock(listener->writer_->writer_mutex_);
+ listener->writer_->transport_ = transport;
+ }
+}
+
+void FlightDataListenerTransport::Clear(
+ const std::shared_ptr<AsyncFlightDataListener>& listener) {
+ if (listener == nullptr) {
+ return;
+ }
+ std::lock_guard<std::mutex> lock(listener->transport_mutex_);
+ listener->transport_ = nullptr;
+ if (listener->writer_ != nullptr) {
+ std::lock_guard<std::mutex> writer_lock(listener->writer_->writer_mutex_);
+ listener->writer_->transport_ = nullptr;
+ }
+}
+
+Status FlightDataListenerTransport::ReportFinish(
+ const std::shared_ptr<AsyncFlightDataListener>& listener, Status status) {
+ if (listener == nullptr) {
+ return Status::OK();
+ }
+ // Cancel() holds the listener's transport lock across its call into the
transport, and
+ // the ending it triggers reports from inside that call.
+ bool expected = false;
+ if (!listener->finished_.compare_exchange_strong(expected, true)) {
+ return Status::OK();
+ }
+ // OnFinish is application code, it may call back into the listener (as
+ // Cancel() does), which is why nothing is held here.
+ return listener->OnFinish(std::move(status)).status();
Review Comment:
`Future::status()` waits for completion (`util/future.h:340–356`).
Consequently, an `OnFinish()` that returns a pending future blocks the gRPC
`OnWriteDone()`/`OnCancel()` callback until application work finishes; an
unresolved future even prevents cancellation teardown. Return/chain the finish
future asynchronously and retain the necessary listener/reactor state instead
of synchronously waiting here.
##########
cpp/src/arrow/flight/transport/grpc/async_grpc_service_internal.h:
##########
@@ -0,0 +1,205 @@
+// 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 <atomic>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <utility>
+
+#include <grpcpp/support/byte_buffer.h>
+#include <grpcpp/support/slice.h>
+
+#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 <typename ProtoT>
+::grpc::ByteBuffer MakeWriteBuffer(const ProtoT& message) {
+ const std::string bytes = message.SerializeAsString();
+ ::grpc::Slice slice(bytes);
+ return {&slice, 1};
+}
+
+template <typename PbT>
+arrow::Result<PbT> 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<int>(bytes.size()))) {
+ return arrow::Status::Invalid("Failed to parse ", what);
+ }
+ return pb;
+}
+
+template <typename PbT, typename T>
+arrow::Result<T> ParseProtoRequest(const ::grpc::ByteBuffer& buf,
std::string_view what) {
+ ARROW_ASSIGN_OR_RAISE(auto pb, ParseProtoRequest<PbT>(buf, what));
+ T out;
+ ARROW_RETURN_NOT_OK(arrow::flight::internal::FromProto(pb, &out));
+ return out;
+}
+
+template <typename T, typename PbT>
+arrow::Status SerializeOrNotFound(const std::shared_ptr<T>& 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)) {
+ Finish(flight_context_.FinishRequest(status));
Review Comment:
A future continuation can pass `!finished()`, then cancellation can call
`Finish()` and gRPC can run `OnDone()` before the continuation calls
`StartWrite()`. `HoldToken` keeps the reactor object alive, but not the
underlying gRPC call, so the write can access torn-down state. Serialize
termination with the finished check and read/write initiation across the unary,
streaming, DoGet, DoPut, and exchange reactors.
##########
cpp/src/arrow/flight/transport/grpc/grpc_server_internal.h:
##########
@@ -154,7 +172,7 @@ class GrpcServerCallContextHelper {
// Authenticate the client (if applicable) and construct the call context
::grpc::Status CheckAuth(FlightMethod method, GrpcContext* context,
GrpcServerCallContext<GrpcContext>* flight_context)
const {
- if (!auth_handler_) {
+ if (!auth_handler_ && !validate_token_) {
Review Comment:
Every async server installs `validate_token_`, so this condition skips
extracting the authenticated TLS peer identity even when the default no-op
validator is used. Handlers and middleware therefore receive an empty
`peer_identity()` for mutual-TLS clients. Populate the TLS identity when there
is no blocking auth handler, then invoke the token validator so the default
hook preserves that identity.
##########
cpp/src/arrow/flight/transport/grpc/async_grpc_do_put.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.
+
+#include "arrow/flight/transport/grpc/async_grpc_service_internal.h"
+
+#include "arrow/flight/transport/grpc/customize_grpc.h"
+
+#include <memory>
+#include <mutex>
+#include <utility>
+
+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<AsyncFlightDataListener> listener)
+ : AsyncReactorBase(std::move(flight_context)),
+ listener_(std::move(listener)),
+ decoder_(listener_) {
+ arrow::flight::internal::FlightDataListenerTransport::Install(listener_,
this);
+ StartRead(&request_buf_);
+ }
+
+ void CancelUpload(Status status) override {
+ if (status.ok()) {
+ status = Status::Invalid("Cancel() needs a non-OK status");
+ }
+ Hold();
+ FinishUpload(std::move(status));
+ ReleaseHold();
+ }
+
+ Status WriteMetadata(const Buffer& app_metadata) override {
+ std::lock_guard<std::mutex> 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) {
+ std::lock_guard<std::mutex> lock(write_mutex_);
Review Comment:
A client that half-closes DoPut without sending any messages reaches this
path without a descriptor and receives a successful acknowledgment and
`OnFinish(OK)`. The synchronous reader rejects this case with `Stream finished
before first message sent`. Reject an empty upload before entering the
clean-ending path so listeners cannot treat a descriptor-less request as valid.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
