westonpace commented on a change in pull request #9528:
URL: https://github.com/apache/arrow/pull/9528#discussion_r588547359
##########
File path: cpp/src/arrow/csv/reader.cc
##########
@@ -934,22 +946,34 @@ class AsyncThreadedTableReader
Result<std::shared_ptr<TableReader>> MakeTableReader(
MemoryPool* pool, io::IOContext io_context,
std::shared_ptr<io::InputStream> input,
const ReadOptions& read_options, const ParseOptions& parse_options,
- const ConvertOptions& convert_options) {
+ const ConvertOptions& convert_options, StopToken stop_token) {
Review comment:
Why is the stop_token passed in explicitly here and it is grabbed from
the `io_context` in `MakeStreamingReader`? Ah, I see it is grabbed out of the
`io_context` earlier but it could be pulled out here and save an extra
parameter.
##########
File path: cpp/src/arrow/testing/gtest_util.cc
##########
@@ -576,6 +594,13 @@ void SleepFor(double seconds) {
std::chrono::nanoseconds(static_cast<int64_t>(seconds * 1e9)));
}
+void BusyWait(double seconds, std::function<bool()> predicate) {
Review comment:
Nice. I was just thinking of writing something like this.
##########
File path: cpp/src/arrow/csv/reader.cc
##########
@@ -848,6 +859,7 @@ class AsyncThreadedTableReader
}
Status Init() override {
+ // TODO pass stop token to input stream iterator?
Review comment:
For an iterator, no. In the future, if this becomes a generator,
possibly. The only way it could be useful is if the I/O could be cancelled
somehow. Generally that's not possible. With some non-blocking I/O schemes
you can at least give up on the user-land side of the I/O. For networked I/O
it may be possible but I seem to recall you saying S3 had no such mechanism.
You could maybe make use of the stop token in the readahead to stop the
readahead but that doesn't seem too urgent. It will fill up the readahead
queue, and then, there will be no active references, and it will all be cleaned
up. Although it may be nice to add a unit test for that scenario. I wonder if
I could get a consumer side reference count of some kind and abort it as soon
as all consumer references are lost :thinking: . I'll add a JIRA for it.
##########
File path: cpp/src/arrow/util/cancel.h
##########
@@ -0,0 +1,99 @@
+// 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 <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "arrow/status.h"
+#include "arrow/type_fwd.h"
+#include "arrow/util/macros.h"
+#include "arrow/util/visibility.h"
+
+namespace arrow {
+
+class StopToken;
+
+struct StopSourceImpl;
+
+/// EXPERIMENTAL
+class ARROW_EXPORT StopSource {
+ public:
+ StopSource();
+ ~StopSource();
+
+ // Consumer API (the side that stops)
+ void RequestStop();
+ void RequestStop(Status error);
+ void RequestStopFromSignal(int signum);
Review comment:
Seems unfortunate this has to be public. Is it possible to make it
private and friend it?
Nevermind, I see that this could be useful if a user wants fine-grained
control over stopping in which case they might need to register their own
signal handlers (since they can't use our signal stop source if they are using
their own stop source).
##########
File path: cpp/src/arrow/util/cancel.cc
##########
@@ -0,0 +1,167 @@
+// 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/util/cancel.h"
+
+#include <atomic>
+#include <mutex>
+#include <sstream>
+#include <utility>
+
+#include "arrow/result.h"
+#include "arrow/util/atomic_shared_ptr.h"
+#include "arrow/util/io_util.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/visibility.h"
+
+namespace arrow {
+
+#if ATOMIC_INT_LOCK_FREE != 2
+#error Lock-free atomic int required for signal safety
+#endif
+
+using internal::ReinstateSignalHandler;
+using internal::SetSignalHandler;
+using internal::SignalHandler;
+
+// NOTE: We care mainly about the making the common case (not cancelled) fast.
+
+struct StopSourceImpl {
+ std::atomic<int> requested_{0}; // will be -1 or signal number if requested
+ std::mutex mutex_;
+ Status cancel_error_;
+};
+
+StopSource::StopSource() : impl_(new StopSourceImpl) {}
+
+StopSource::~StopSource() = default;
+
+void StopSource::RequestStop() { RequestStop(Status::Cancelled("Operation
cancelled")); }
+
+void StopSource::RequestStop(Status st) {
+ std::lock_guard<std::mutex> lock(impl_->mutex_);
+ DCHECK(!st.ok());
+ if (!impl_->requested_) {
+ impl_->requested_ = -1;
+ impl_->cancel_error_ = std::move(st);
+ }
+}
+
+void StopSource::RequestStopFromSignal(int signum) {
+ // Only async-signal-safe code allowed here
+ impl_->requested_.store(signum);
Review comment:
I think you can get rid of this mutex. If you changed this to something
like...
```
if (!impl_->requested_.fetch_or(-1)) {
impl_->signum = signum;
}
```
`signum` would be a new (non-atomic) int member that replaces the mutex.
Then `RequestStop` becomes...
```
if(!impl->requested_.fetch_or(-1)) {
impl->cancel_error_ = std::move(st);
}
```
The main advantage is there would be no way to lose your thread when calling
`Poll`.
##########
File path: python/pyarrow/_csv.pyx
##########
@@ -34,8 +34,8 @@ from pyarrow.lib cimport (check_status, Field, MemoryPool,
Schema,
pyarrow_unwrap_batch, pyarrow_unwrap_table,
pyarrow_wrap_schema, pyarrow_wrap_table,
pyarrow_wrap_data_type, pyarrow_unwrap_data_type,
- Table, RecordBatch)
-from pyarrow.lib import frombytes, tobytes
+ Table, RecordBatch, StopToken)
Review comment:
Is there a JIRA to expose `StopSource` to python at some point? There
are reasons other than signals someone may want to cancel an operation. For
example, a GUI-based application may have a cancel button. A web server may
want to cancel if the TCP connection for some analysis request is lost.
##########
File path: cpp/src/arrow/util/cancel.cc
##########
@@ -0,0 +1,167 @@
+// 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/util/cancel.h"
+
+#include <atomic>
+#include <mutex>
+#include <sstream>
+#include <utility>
+
+#include "arrow/result.h"
+#include "arrow/util/atomic_shared_ptr.h"
+#include "arrow/util/io_util.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/visibility.h"
+
+namespace arrow {
+
+#if ATOMIC_INT_LOCK_FREE != 2
+#error Lock-free atomic int required for signal safety
+#endif
+
+using internal::ReinstateSignalHandler;
+using internal::SetSignalHandler;
+using internal::SignalHandler;
+
+// NOTE: We care mainly about the making the common case (not cancelled) fast.
+
+struct StopSourceImpl {
+ std::atomic<int> requested_{0}; // will be -1 or signal number if requested
+ std::mutex mutex_;
+ Status cancel_error_;
+};
+
+StopSource::StopSource() : impl_(new StopSourceImpl) {}
+
+StopSource::~StopSource() = default;
+
+void StopSource::RequestStop() { RequestStop(Status::Cancelled("Operation
cancelled")); }
+
+void StopSource::RequestStop(Status st) {
+ std::lock_guard<std::mutex> lock(impl_->mutex_);
+ DCHECK(!st.ok());
+ if (!impl_->requested_) {
+ impl_->requested_ = -1;
+ impl_->cancel_error_ = std::move(st);
+ }
+}
+
+void StopSource::RequestStopFromSignal(int signum) {
+ // Only async-signal-safe code allowed here
+ impl_->requested_.store(signum);
+}
+
+StopToken StopSource::token() { return StopToken(impl_); }
+
+bool StopToken::IsStopRequested() {
+ if (!impl_) {
+ return false;
+ }
+ return impl_->requested_.load() != 0;
+}
+
+Status StopToken::Poll() {
+ if (!impl_) {
+ return Status::OK();
+ }
+ if (!impl_->requested_.load()) {
+ return Status::OK();
+ }
+
+ std::lock_guard<std::mutex> lock(impl_->mutex_);
+ if (impl_->cancel_error_.ok()) {
+ auto signum = impl_->requested_.load();
+ DCHECK_GT(signum, 0);
+ impl_->cancel_error_ = internal::CancelledFromSignal(signum, "Operation
cancelled");
+ }
+ return impl_->cancel_error_;
+}
+
+namespace {
+
+void HandleSignal(int signum);
+
+struct SignalStopState {
+ struct SavedSignalHandler {
+ int signum;
+ SignalHandler handler;
+ };
+
+ Status RegisterHandlers(const std::vector<int>& signals) {
+ if (!saved_handlers.empty()) {
+ return Status::Invalid("Signal handlers already registered");
+ }
+ for (int signum : signals) {
+ ARROW_ASSIGN_OR_RAISE(auto handler,
+ SetSignalHandler(signum,
SignalHandler{&HandleSignal}));
+ saved_handlers.push_back({signum, handler});
+ }
+ return Status::OK();
+ }
+
+ void UnregisterHandlers() {
+ auto handlers = std::move(saved_handlers);
+ for (const auto& h : handlers) {
+ ARROW_CHECK_OK(SetSignalHandler(h.signum, h.handler).status());
+ }
+ }
+
+ ~SignalStopState() { UnregisterHandlers(); }
+
+ StopSource stop_source;
+ std::vector<SavedSignalHandler> saved_handlers;
+};
+
+std::shared_ptr<SignalStopState> g_signal_stop_state;
+
+void HandleSignal(int signum) {
+ ReinstateSignalHandler(signum, &HandleSignal);
+ std::shared_ptr<SignalStopState> state =
internal::atomic_load(&g_signal_stop_state);
Review comment:
If someone calls `ResetSignalStopSource` while you are handling the
signal wouldn't this become the sole reference and, on exit, attempt to destroy
it (which would not be async-signal-safe)?
Admittedly, probably not a common occurrence.
##########
File path: cpp/src/arrow/util/cancel_test.cc
##########
@@ -0,0 +1,302 @@
+// 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 <atomic>
+#include <cmath>
+#include <sstream>
+#include <string>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include <signal.h>
+#ifndef _WIN32
+#include <sys/time.h> // for setitimer()
+#endif
+
+#include "arrow/testing/gtest_util.h"
+#include "arrow/util/cancel.h"
+#include "arrow/util/future.h"
+#include "arrow/util/io_util.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/optional.h"
+
+namespace arrow {
+
+class CancelTest : public ::testing::Test {};
+
+TEST_F(CancelTest, StopBasics) {
+ {
+ StopSource source;
+ StopToken token = source.token();
+ ASSERT_FALSE(token.IsStopRequested());
+ ASSERT_OK(token.Poll());
+
+ source.RequestStop();
+ ASSERT_TRUE(token.IsStopRequested());
+ ASSERT_RAISES(Cancelled, token.Poll());
+ }
+ {
+ StopSource source;
+ StopToken token = source.token();
+ source.RequestStop(Status::IOError("Operation cancelled"));
+ ASSERT_TRUE(token.IsStopRequested());
+ ASSERT_RAISES(IOError, token.Poll());
+ }
+}
+
+TEST_F(CancelTest, StopTokenCopy) {
+ StopSource source;
+ StopToken token = source.token();
+ ASSERT_FALSE(token.IsStopRequested());
+ ASSERT_OK(token.Poll());
+
+ source.RequestStop();
+ ASSERT_TRUE(token.IsStopRequested());
+ ASSERT_RAISES(Cancelled, token.Poll());
+
+ StopToken new_token = token;
Review comment:
Minor nit, but maybe also test the case where you copy the token before
stopping.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
