Copilot commented on code in PR #55419: URL: https://github.com/apache/arrow/pull/55419#discussion_r4513775275
########## cpp/src/arrow/flight/flight_data_decoder.cc: ########## @@ -0,0 +1,336 @@ +// 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) { + // 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<std::mutex> 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<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(); + } + // Report exactly once: the first ending wins. + 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. Its future must not be + // waited on: this runs on a gRPC callback thread (OnWriteDone/OnCancel), and + // blocking it would stall every RPC multiplexed onto that thread. Chain the + // completion instead so a failing OnFinish is still reported. + listener->OnFinish(std::move(status)).AddCallback([](const Status& finish_status) { + ARROW_WARN_NOT_OK(finish_status, + "Reporting the end of an upload to the listener failed"); + }); Review Comment: This callback does not retain the listener, so RPC completion can destroy a per-RPC listener while its `OnFinish()` future is still pending. Asynchronous finish work that accesses listener members can then use freed memory. Capture the listener in the completion callback to keep it alive until that work finishes. ########## cpp/src/arrow/flight/flight_data_decoder.cc: ########## @@ -0,0 +1,336 @@ +// 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) { + // 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<std::mutex> 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<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(); + } + // Report exactly once: the first ending wins. + 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. Its future must not be + // waited on: this runs on a gRPC callback thread (OnWriteDone/OnCancel), and + // blocking it would stall every RPC multiplexed onto that thread. Chain the + // completion instead so a failing OnFinish is still reported. + listener->OnFinish(std::move(status)).AddCallback([](const Status& finish_status) { + ARROW_WARN_NOT_OK(finish_status, + "Reporting the end of an upload to the listener failed"); + }); + return Status::OK(); +} + +} // namespace internal + +Future<> AsyncFlightDataListener::Cancel(Status status) { + internal::FlightDataListenerTransport* transport; + { + std::lock_guard<std::mutex> lock(transport_mutex_); + transport = transport_; + } Review Comment: After `Cancel()` releases `transport_mutex_`, concurrent `DoPutReactor::OnDone()` can clear the pointer and delete the reactor before `CancelUpload()` is invoked, causing a use-after-free. The `Hold()` inside `CancelUpload()` is too late to protect entry. Acquire a lifetime reference under the listener mutex, then invoke cancellation unlocked so `OnFinish()` can safely re-enter listener methods. ########## 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(); +} + +void ExchangeReactor::OnDone() { + // The handler may still hold the reader/writer: cut their back-pointers + // before this reactor is deleted so a late call fails cleanly instead of + // dereferencing freed memory. + reader_->DetachReactor(); + writer_->DetachReactor(); + AsyncReactorBase::OnDone(); +} + +void ExchangeReactor::StartReadNext() { + if (!finished()) { + StartRead(&read_buf_); + } +} + +arrow::Status ExchangeReactor::WritePayload(FlightPayload payload) { + if (finished()) { + return arrow::Status::Cancelled("the exchange is over"); + } + 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()); + } + StartWrite(&write_buf_); Review Comment: Cancellation can run `FinishOnce()` after the `finished()` check but before `StartWrite()`, allowing an application thread to start I/O on a torn-down gRPC call even if the reactor object remains alive. Serialize the live check and I/O initiation against `FinishOnce()` using the shared operation mutex. Apply the same guard to `StartReadNext()` here and the decode continuation in `async_grpc_do_put.cc:136`. ########## cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_reader.cc: ########## @@ -0,0 +1,221 @@ +// 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 { + if (reader_->pending_.is_valid() && !reader_->pending_.is_finished()) { + reader_->ResolvePending(std::move(chunk)); + } else { + reader_->buffered_ = std::move(chunk); + } Review Comment: `OnNext()` reads `pending_` without `mutex_`, while cancellation can concurrently move and reset that same future, causing a data race and potentially a null dereference. Claim the pending demand or update `buffered_` under the reader mutex, then complete the claimed future outside the lock to allow reentrant continuations. ########## cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_reader.cc: ########## @@ -0,0 +1,221 @@ +// 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 { + if (reader_->pending_.is_valid() && !reader_->pending_.is_finished()) { + reader_->ResolvePending(std::move(chunk)); + } else { + reader_->buffered_ = std::move(chunk); + } + return arrow::Future<>::MakeFinished(); + } + + arrow::Future<> OnDescriptor(const FlightDescriptor&) override { + return arrow::Future<>::MakeFinished(); + } + + arrow::Status OnSchemaDecoded(std::shared_ptr<Schema>) override { + return arrow::Status::OK(); + } + + private: + ExchangeReader* reader_; +}; + +ExchangeReader::ExchangeReader() + : listener_(std::make_shared<ChunkListener>(this)), decoder_(listener_) {} + +const FlightDescriptor& ExchangeReader::descriptor() const { return descriptor_; } + +arrow::Future<FlightStreamChunk> ExchangeReader::NextAsync() { + using FSCFuture = arrow::Future<FlightStreamChunk>; + + { + std::lock_guard<std::mutex> 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(); + } + // 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<std::mutex> lock(mutex_); + data = std::move(*first_read_); + first_read_.reset(); + } + Consume(std::move(data)); + } else { + RequestRead(); + } + { + std::lock_guard<std::mutex> lock(mutex_); + return pending_; Review Comment: `Consume()` can resolve the demand inline, and a fast read can also complete before `NextAsync()` returns; `ResolvePending()` then resets `pending_`, so callers receive an invalid future and can crash when attaching `Then()`. Copy the newly created future under the initial lock and return that local handle after starting decoding or I/O. ########## cpp/src/arrow/flight/transport/grpc/async_grpc_exchange_writer.cc: ########## @@ -0,0 +1,222 @@ +// 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() { reactor_.store(nullptr); } + +arrow::Future<> ExchangeWriter::BeginAsync(std::shared_ptr<Schema> 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<PayloadSink>(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<Buffer> app_metadata) { + pending_app_metadata_ = std::move(app_metadata); + return WriteBatch(batch); +} + +arrow::Future<> ExchangeWriter::WriteMetadataAsync(std::shared_ptr<Buffer> 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<std::mutex> lock(mutex_); + write_in_flight_ = false; + } + FailFlush(MakeFlightError(FlightStatusCode::Internal, "Failed to write response")); + return; + } + { + std::lock_guard<std::mutex> lock(mutex_); + write_in_flight_ = false; + } + Pump(); +} + +void ExchangeWriter::OnCancelled() { + bool had_write; + { + std::lock_guard<std::mutex> 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<std::mutex> 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<std::mutex> lock(mutex_); + if (write_in_flight_) { + return; + } + reactor = reactor_.load(); + 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_. + switch (action) { + case Action::kDrained: + if (drained.is_valid() && !drained.is_finished()) { + drained.MarkFinished(); + } + return; + case Action::kNoReactor: + 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)); Review Comment: After `Pump()` loads `reactor_`, cancellation can complete the handler future and let `OnDone()` delete the reactor before this dereference; the atomic pointer does not preserve its lifetime. Acquire a reactor lifetime reference in synchronization with detachment and retain it through the call. Apply the same fix to `ExchangeReader::RequestRead()` in `async_grpc_exchange_reader.cc:206–218`. ########## cpp/examples/arrow/flight_async_generic_doget_example.cc: ########## @@ -0,0 +1,375 @@ +// 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 <windows.h> +#else +# include <unistd.h> +#endif +#include <csignal> +#if defined(__linux__) +# include <sys/prctl.h> +#endif + +#include <chrono> +#include <condition_variable> +#include <cstdint> +#include <cstdlib> +#include <iostream> +#include <memory> +#include <mutex> +#include <optional> +#include <string> +#include <thread> +#include <utility> + +#include <arrow/api.h> +#include <arrow/flight/api.h> +#include <gflags/gflags.h> + +// 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<int64_t>(::GetCurrentProcessId()); +#else + return static_cast<int64_t>(::getpid()); +#endif +} + +/// \brief A reader with `num_batches` batches of three int64 values. +class BatchReader : public arrow::RecordBatchReader { + public: + BatchReader(std::shared_ptr<arrow::Schema> schema, int32_t num_batches) + : schema_(std::move(schema)), num_batches_(num_batches) {} + + std::shared_ptr<arrow::Schema> schema() const override { return schema_; } + + arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch>* 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<arrow::Schema> 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<flight::FlightDataStream> stream; + std::shared_ptr<arrow::Schema> 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<std::optional<flight::FlightPayload>> 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<WorkerState> state) { +#if defined(__linux__) + // Name the thread, so /proc/<pid>/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<std::mutex> 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<std::optional<flight::FlightPayload>> 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(); + } + 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<flight::FlightDataStream> stream, int32_t delay_ms) + : state_(std::make_shared<WorkerState>()) { + state_->stream = std::move(stream); + state_->delay_ms = delay_ms; + std::shared_ptr<WorkerState> 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<arrow::Schema> schema() override { return state_->schema; } + + arrow::Future<flight::FlightPayload> 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<flight::FlightPayload>::MakeFinished( + state_->stream->GetSchemaPayload()); + } + + arrow::Future<std::optional<flight::FlightPayload>> NextAsync() override { + arrow::Future<std::optional<flight::FlightPayload>> future = + arrow::Future<std::optional<flight::FlightPayload>>::Make(); + { + std::lock_guard<std::mutex> 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 { + Stop(); + { + std::lock_guard<std::mutex> lock(state_->mutex); + if (state_->stream_closed) { + return arrow::Status::OK(); + } + state_->stream_closed = true; + } + return state_->stream->Close(); Review Comment: On cancellation, this runs on a gRPC thread while the worker may already have passed its `closed` check and be executing `state->stream->Next()` outside the mutex. `RecordBatchStream::Next()` and `Close()` both mutate the IPC writer, so calling them concurrently introduces a data race and can corrupt the stream. Have the worker perform the underlying stream's teardown after it stops producing, rather than closing it directly from the cancellation thread; keep that handoff non-blocking. -- 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]
