Copilot commented on code in PR #50407: URL: https://github.com/apache/arrow/pull/50407#discussion_r4280507389
########## cpp/src/arrow/flight/transport_server_internal.cc: ########## @@ -0,0 +1,332 @@ +// 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_server_internal.h" + +#include "arrow/buffer.h" +#include "arrow/flight/serialization_internal.h" +#include "arrow/flight/transport_server.h" +#include "arrow/ipc/reader.h" + +namespace arrow::flight::internal { +namespace { + +class TransportIpcMessageReader : public ipc::MessageReader { + public: + TransportIpcMessageReader(std::shared_ptr<PeekableFlightDataReader> peekable_reader, + std::shared_ptr<MemoryManager> memory_manager, + std::shared_ptr<Buffer>* app_metadata) + : peekable_reader_(std::move(peekable_reader)), + memory_manager_(std::move(memory_manager)), + app_metadata_(app_metadata) {} + + arrow::Result<std::unique_ptr<ipc::Message>> ReadNextMessage() override { + if (stream_finished_) return nullptr; + FlightData* data; + peekable_reader_->Next(&data); + if (!data) { + stream_finished_ = true; + return nullptr; + } + if (data->body) { + ARROW_ASSIGN_OR_RAISE(data->body, Buffer::ViewOrCopy(data->body, memory_manager_)); + } + *app_metadata_ = std::move(data->app_metadata); + return data->OpenMessage(); + } + + private: + std::shared_ptr<PeekableFlightDataReader> peekable_reader_; + std::shared_ptr<MemoryManager> memory_manager_; + // A reference to TransportDataStream.app_metadata_. That class + // can't access the app metadata because when it Peek()s the stream, + // it may be looking at a dictionary batch, not the record + // batch. Updating it here ensures the reader is always updated with + // the last metadata message read. + std::shared_ptr<Buffer>* app_metadata_; + bool stream_finished_ = false; +}; + +} // namespace + +TransportMessageReader::TransportMessageReader( + ServerDataStream* stream, std::shared_ptr<MemoryManager> memory_manager) + : peekable_reader_(new PeekableFlightDataReader(stream)), + memory_manager_(std::move(memory_manager)) {} + +Status TransportMessageReader::Init() { + // Peek the first message to get the descriptor. + FlightData* data; + peekable_reader_->Peek(&data); + if (!data) { + return Status::IOError("Stream finished before first message sent"); + } + if (!data->descriptor) { + return Status::IOError("Descriptor missing on first message"); + } + descriptor_ = *data->descriptor; + // If there's a schema (=DoPut), also Open(). + if (data->metadata) { + return EnsureDataStarted(); + } + peekable_reader_->Next(&data); + return Status::OK(); +} + +const FlightDescriptor& TransportMessageReader::descriptor() const { return descriptor_; } + +arrow::Result<std::shared_ptr<Schema>> TransportMessageReader::GetSchema() { + RETURN_NOT_OK(EnsureDataStarted()); + return batch_reader_->schema(); +} + +arrow::Result<FlightStreamChunk> TransportMessageReader::Next() { + FlightStreamChunk out; + FlightData* data = nullptr; + peekable_reader_->Peek(&data); + if (!data) { + out.app_metadata = nullptr; + out.data = nullptr; + return out; + } + + if (!data->metadata) { + // Metadata-only (data->metadata is the IPC header) + out.app_metadata = data->app_metadata; + out.data = nullptr; + peekable_reader_->Next(&data); + return out; + } + + if (!batch_reader_) { + RETURN_NOT_OK(EnsureDataStarted()); + // re-peek here since EnsureDataStarted() advances the stream + return Next(); + } + RETURN_NOT_OK(batch_reader_->ReadNext(&out.data)); + out.app_metadata = std::move(app_metadata_); + return out; +} + +ipc::ReadStats TransportMessageReader::stats() const { + if (batch_reader_ == nullptr) { + return ipc::ReadStats{}; + } + return batch_reader_->stats(); +} + +Status TransportMessageReader::EnsureDataStarted() { + if (!batch_reader_) { + // peek() until we find the first data message; discard metadata + if (!peekable_reader_->SkipToData()) { + return Status::IOError("Client never sent a data message"); + } + auto message_reader = std::unique_ptr<ipc::MessageReader>( + new TransportIpcMessageReader(peekable_reader_, memory_manager_, &app_metadata_)); + ARROW_ASSIGN_OR_RAISE(batch_reader_, + ipc::RecordBatchStreamReader::Open(std::move(message_reader))); + } + return Status::OK(); +} + +namespace { + +/// \brief An IpcPayloadWriter for ServerDataStream. +/// +/// To support app_metadata and reuse the existing IPC infrastructure, +/// this takes a pointer to a buffer to be combined with the IPC +/// payload when writing a Flight payload. +class TransportMessagePayloadWriter : public ipc::internal::IpcPayloadWriter { + public: + TransportMessagePayloadWriter(ServerDataStream* stream, + std::shared_ptr<Buffer>* app_metadata) + : stream_(stream), app_metadata_(app_metadata) {} + + Status Start() override { return Status::OK(); } + + Status WritePayload(const ipc::IpcPayload& ipc_payload) override { + FlightPayload payload; + payload.ipc_message = ipc_payload; + if (ipc_payload.type == ipc::MessageType::RECORD_BATCH && *app_metadata_) { + payload.app_metadata = std::move(*app_metadata_); + } + ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload)); + if (!success) { + return MakeFlightError( + FlightStatusCode::Internal, + "Could not write record batch to stream (client disconnect?)"); + } + return Status::OK(); + } + + Status Close() override { + // Closing is handled one layer up in TransportMessageWriter::Close + return Status::OK(); + } + + private: + ServerDataStream* stream_; + std::shared_ptr<Buffer>* app_metadata_; +}; + +} // namespace + +TransportMessageWriter::TransportMessageWriter(ServerDataStream* stream) + : stream_(stream), + app_metadata_(nullptr), + ipc_options_(::arrow::ipc::IpcWriteOptions::Defaults()) {} + +Status TransportMessageWriter::Begin(const std::shared_ptr<Schema>& schema, + const ipc::IpcWriteOptions& options) { + if (batch_writer_) { + return Status::Invalid("This writer has already been started."); + } + ipc_options_ = options; + std::unique_ptr<ipc::internal::IpcPayloadWriter> payload_writer( + new TransportMessagePayloadWriter(stream_, &app_metadata_)); + ARROW_ASSIGN_OR_RAISE(batch_writer_, + ipc::internal::OpenRecordBatchWriter(std::move(payload_writer), + schema, ipc_options_)); + return Status::OK(); +} + +Status TransportMessageWriter::WriteRecordBatch(const RecordBatch& batch) { + return WriteWithMetadata(batch, nullptr); +} + +Status TransportMessageWriter::WriteMetadata(std::shared_ptr<Buffer> app_metadata) { + FlightPayload payload{}; + payload.app_metadata = std::move(app_metadata); + ARROW_ASSIGN_OR_RAISE(auto success, stream_->WriteData(payload)); + if (!success) { + ARROW_RETURN_NOT_OK(Close()); + return MakeFlightError(FlightStatusCode::Internal, + "Could not write metadata to stream (client disconnect?)"); + } + return Status::OK(); +} + +Status TransportMessageWriter::WriteWithMetadata(const RecordBatch& batch, + std::shared_ptr<Buffer> app_metadata) { + RETURN_NOT_OK(CheckStarted()); + app_metadata_ = std::move(app_metadata); + auto status = batch_writer_->WriteRecordBatch(batch); + if (!status.ok()) { + ARROW_RETURN_NOT_OK(Close()); + } + return status; +} + +Status TransportMessageWriter::Close() { + if (batch_writer_) { + RETURN_NOT_OK(batch_writer_->Close()); + } + return Status::OK(); +} + +ipc::WriteStats TransportMessageWriter::stats() const { + ARROW_CHECK_NE(batch_writer_, nullptr); + return batch_writer_->stats(); +} + +Status TransportMessageWriter::CheckStarted() { + if (!batch_writer_) { + return Status::Invalid("This writer is not started. Call Begin() with a schema"); + } + return Status::OK(); +} + +TransportMetadataWriter::TransportMetadataWriter(ServerDataStream* stream) + : stream_(stream) {} + +Status TransportMetadataWriter::WriteMetadata(const Buffer& buffer) { + return stream_->WritePutMetadata(buffer); +} + +std::atomic<ServerSignalState*> ServerSignalState::running_instance_; Review Comment: The old implementation explicitly required both atomics used by the signal handler (`running_instance_` and `got_signal_`) to be always lock-free. That guard was dropped during the move, so on a supported platform with a non-lock-free atomic this code can invoke a potentially non-signal-safe atomic operation from `HandleSignal`, making signal handling unsafe. Please restore the compile-time lock-free check (or otherwise avoid these atomics in the signal handler). -- 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]
