461 lines
19 KiB
C
461 lines
19 KiB
C
![]() |
// 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 <functional>
|
||
|
#include <list>
|
||
|
#include <memory>
|
||
|
|
||
|
#include "arrow/result.h"
|
||
|
#include "arrow/status.h"
|
||
|
#include "arrow/util/cancel.h"
|
||
|
#include "arrow/util/functional.h"
|
||
|
#include "arrow/util/future.h"
|
||
|
#include "arrow/util/iterator.h"
|
||
|
#include "arrow/util/mutex.h"
|
||
|
#include "arrow/util/thread_pool.h"
|
||
|
#include "arrow/util/tracing.h"
|
||
|
|
||
|
namespace arrow {
|
||
|
|
||
|
using internal::FnOnce;
|
||
|
|
||
|
namespace util {
|
||
|
|
||
|
/// A utility which keeps tracks of, and schedules, asynchronous tasks
|
||
|
///
|
||
|
/// An asynchronous task has a synchronous component and an asynchronous component.
|
||
|
/// The synchronous component typically schedules some kind of work on an external
|
||
|
/// resource (e.g. the I/O thread pool or some kind of kernel-based asynchronous
|
||
|
/// resource like io_uring). The asynchronous part represents the work
|
||
|
/// done on that external resource. Executing the synchronous part will be referred
|
||
|
/// to as "submitting the task" since this usually includes submitting the asynchronous
|
||
|
/// portion to the external thread pool.
|
||
|
///
|
||
|
/// By default the scheduler will submit the task (execute the synchronous part) as
|
||
|
/// soon as it is added, assuming the underlying thread pool hasn't terminated or the
|
||
|
/// scheduler hasn't aborted. In this mode, the scheduler is simply acting as
|
||
|
/// a simple task group.
|
||
|
///
|
||
|
/// A task scheduler starts with an initial task. That task, and all subsequent tasks
|
||
|
/// are free to add subtasks. Once all submitted tasks finish the scheduler will
|
||
|
/// finish. Note, it is not an error to add additional tasks after a scheduler has
|
||
|
/// aborted. These tasks will be ignored and never submitted. The scheduler returns a
|
||
|
/// future which will complete when all submitted tasks have finished executing. Once all
|
||
|
/// tasks have been finished the scheduler is invalid and should no longer be used.
|
||
|
///
|
||
|
/// Task failure (either the synchronous portion or the asynchronous portion) will cause
|
||
|
/// the scheduler to enter an aborted state. The first such failure will be reported in
|
||
|
/// the final task future.
|
||
|
class ARROW_EXPORT AsyncTaskScheduler {
|
||
|
public:
|
||
|
/// Destructor for AsyncTaskScheduler
|
||
|
///
|
||
|
/// The lifetime of the task scheduled is managed automatically. The scheduler
|
||
|
/// will remain valid while any tasks are running (and can always be safely accessed)
|
||
|
/// within tasks) and will be destroyed as soon as all tasks have finished.
|
||
|
virtual ~AsyncTaskScheduler() = default;
|
||
|
/// An interface for a task
|
||
|
///
|
||
|
/// Users may want to override this, for example, to add priority
|
||
|
/// information for use by a queue.
|
||
|
class Task {
|
||
|
public:
|
||
|
virtual ~Task() = default;
|
||
|
/// Submit the task
|
||
|
///
|
||
|
/// This will be called by the scheduler at most once when there
|
||
|
/// is space to run the task. This is expected to be a fairly quick
|
||
|
/// function that simply submits the actual task work to an external
|
||
|
/// resource (e.g. I/O thread pool).
|
||
|
///
|
||
|
/// If this call fails then the scheduler will enter an aborted state.
|
||
|
virtual Result<Future<>> operator()() = 0;
|
||
|
/// The cost of the task
|
||
|
///
|
||
|
/// A ThrottledAsyncTaskScheduler can be used to limit the number of concurrent tasks.
|
||
|
/// A custom cost may be used, for example, if you would like to limit the number of
|
||
|
/// tasks based on the total expected RAM usage of the tasks (this is done in the
|
||
|
/// scanner)
|
||
|
virtual int cost() const { return 1; }
|
||
|
/// The name of the task
|
||
|
///
|
||
|
/// This is used for debugging and traceability. The returned view must remain
|
||
|
/// valid for the lifetime of the task.
|
||
|
virtual std::string_view name() const = 0;
|
||
|
|
||
|
/// a span tied to the lifetime of the task, for internal use only
|
||
|
tracing::Span span;
|
||
|
};
|
||
|
|
||
|
/// Add a task to the scheduler
|
||
|
///
|
||
|
/// If the scheduler is in an aborted state this call will return false and the task
|
||
|
/// will never be run. This is harmless and does not need to be guarded against.
|
||
|
///
|
||
|
/// The return value for this call can usually be ignored. There is little harm in
|
||
|
/// attempting to add tasks to an aborted scheduler. It is only included for callers
|
||
|
/// that want to avoid future task generation to save effort.
|
||
|
///
|
||
|
/// \param task the task to submit
|
||
|
///
|
||
|
/// A task's name must remain valid for the duration of the task. It is used for
|
||
|
/// debugging (e.g. when debugging a deadlock to see which tasks still remain) and for
|
||
|
/// traceability (the name will be used for spans assigned to the task)
|
||
|
///
|
||
|
/// \return true if the task was submitted or queued, false if the task was ignored
|
||
|
virtual bool AddTask(std::unique_ptr<Task> task) = 0;
|
||
|
|
||
|
/// Adds an async generator to the scheduler
|
||
|
///
|
||
|
/// The async generator will be visited, one item at a time. Submitting a task
|
||
|
/// will consist of polling the generator for the next future. The generator's future
|
||
|
/// will then represent the task itself.
|
||
|
///
|
||
|
/// This visits the task serially without readahead. If readahead or parallelism
|
||
|
/// is desired then it should be added in the generator itself.
|
||
|
///
|
||
|
/// The generator itself will be kept alive until all tasks have been completed.
|
||
|
/// However, if the scheduler is aborted, the generator will be destroyed as soon as the
|
||
|
/// next item would be requested.
|
||
|
///
|
||
|
/// \param generator the generator to submit to the scheduler
|
||
|
/// \param visitor a function which visits each generator future as it completes
|
||
|
/// \param name a name which will be used for each submitted task
|
||
|
template <typename T>
|
||
|
bool AddAsyncGenerator(std::function<Future<T>()> generator,
|
||
|
std::function<Status(const T&)> visitor, std::string_view name);
|
||
|
|
||
|
template <typename Callable>
|
||
|
struct SimpleTask : public Task {
|
||
|
SimpleTask(Callable callable, std::string_view name)
|
||
|
: callable(std::move(callable)), name_(name) {}
|
||
|
SimpleTask(Callable callable, std::string name)
|
||
|
: callable(std::move(callable)), owned_name_(std::move(name)) {
|
||
|
name_ = *owned_name_;
|
||
|
}
|
||
|
Result<Future<>> operator()() override { return callable(); }
|
||
|
std::string_view name() const override { return name_; }
|
||
|
Callable callable;
|
||
|
std::string_view name_;
|
||
|
std::optional<std::string> owned_name_;
|
||
|
};
|
||
|
|
||
|
/// Add a task with cost 1 to the scheduler
|
||
|
///
|
||
|
/// \param callable a "submit" function that should return a future
|
||
|
/// \param name a name for the task
|
||
|
///
|
||
|
/// `name` must remain valid until the task has been submitted AND the returned
|
||
|
/// future completes. It is used for debugging and tracing.
|
||
|
///
|
||
|
/// \see AddTask for more details
|
||
|
template <typename Callable>
|
||
|
bool AddSimpleTask(Callable callable, std::string_view name) {
|
||
|
return AddTask(std::make_unique<SimpleTask<Callable>>(std::move(callable), name));
|
||
|
}
|
||
|
|
||
|
/// Add a task with cost 1 to the scheduler
|
||
|
///
|
||
|
/// This is an overload of \see AddSimpleTask that keeps `name` alive
|
||
|
/// in the task.
|
||
|
template <typename Callable>
|
||
|
bool AddSimpleTask(Callable callable, std::string name) {
|
||
|
return AddTask(
|
||
|
std::make_unique<SimpleTask<Callable>>(std::move(callable), std::move(name)));
|
||
|
}
|
||
|
|
||
|
/// Construct a scheduler
|
||
|
///
|
||
|
/// \param initial_task The initial task which is responsible for adding
|
||
|
/// the first subtasks to the scheduler.
|
||
|
/// \param abort_callback A callback that will be triggered immediately after a task
|
||
|
/// fails while other tasks may still be running. Nothing needs to be done here,
|
||
|
/// when a task fails the scheduler will stop accepting new tasks and eventually
|
||
|
/// return the error. However, this callback can be used to more quickly end
|
||
|
/// long running tasks that have already been submitted. Defaults to doing
|
||
|
/// nothing.
|
||
|
/// \param stop_token An optional stop token that will allow cancellation of the
|
||
|
/// scheduler. This will be checked before each task is submitted and, in the
|
||
|
/// event of a cancellation, the scheduler will enter an aborted state. This is
|
||
|
/// a graceful cancellation and submitted tasks will still complete.
|
||
|
/// \return A future that will be completed when the initial task and all subtasks have
|
||
|
/// finished.
|
||
|
static Future<> Make(
|
||
|
FnOnce<Status(AsyncTaskScheduler*)> initial_task,
|
||
|
FnOnce<void(const Status&)> abort_callback = [](const Status&) {},
|
||
|
StopToken stop_token = StopToken::Unstoppable());
|
||
|
|
||
|
/// A span tracking execution of the scheduler's tasks, for internal use only
|
||
|
virtual const tracing::Span& span() const = 0;
|
||
|
};
|
||
|
|
||
|
class ARROW_EXPORT ThrottledAsyncTaskScheduler : public AsyncTaskScheduler {
|
||
|
public:
|
||
|
/// An interface for a task queue
|
||
|
///
|
||
|
/// A queue's methods will not be called concurrently
|
||
|
class Queue {
|
||
|
public:
|
||
|
virtual ~Queue() = default;
|
||
|
/// Push a task to the queue
|
||
|
///
|
||
|
/// \param task the task to enqueue
|
||
|
virtual void Push(std::unique_ptr<Task> task) = 0;
|
||
|
/// Pop the next task from the queue
|
||
|
virtual std::unique_ptr<Task> Pop() = 0;
|
||
|
/// Peek the next task in the queue
|
||
|
virtual const Task& Peek() = 0;
|
||
|
/// Check if the queue is empty
|
||
|
virtual bool Empty() = 0;
|
||
|
/// Purge the queue of all items
|
||
|
virtual void Purge() = 0;
|
||
|
virtual std::size_t Size() const = 0;
|
||
|
};
|
||
|
|
||
|
class Throttle {
|
||
|
public:
|
||
|
virtual ~Throttle() = default;
|
||
|
/// Acquire amt permits
|
||
|
///
|
||
|
/// If nullopt is returned then the permits were immediately
|
||
|
/// acquired and the caller can proceed. If a future is returned then the caller
|
||
|
/// should wait for the future to complete first. When the returned future completes
|
||
|
/// the permits have NOT been acquired and the caller must call Acquire again
|
||
|
///
|
||
|
/// \param amt the number of permits to acquire
|
||
|
virtual std::optional<Future<>> TryAcquire(int amt) = 0;
|
||
|
/// Release amt permits
|
||
|
///
|
||
|
/// This will possibly complete waiting futures and should probably not be
|
||
|
/// called while holding locks.
|
||
|
///
|
||
|
/// \param amt the number of permits to release
|
||
|
virtual void Release(int amt) = 0;
|
||
|
|
||
|
/// The size of the largest task that can run
|
||
|
///
|
||
|
/// Incoming tasks will have their cost latched to this value to ensure
|
||
|
/// they can still run (although they will be the only thing allowed to
|
||
|
/// run at that time).
|
||
|
virtual int Capacity() = 0;
|
||
|
|
||
|
/// Pause the throttle
|
||
|
///
|
||
|
/// Any tasks that have been submitted already will continue. However, no new tasks
|
||
|
/// will be run until the throttle is resumed.
|
||
|
virtual void Pause() = 0;
|
||
|
/// Resume the throttle
|
||
|
///
|
||
|
/// Allows task to be submitted again. If there is a max_concurrent_cost limit then
|
||
|
/// it will still apply.
|
||
|
virtual void Resume() = 0;
|
||
|
};
|
||
|
|
||
|
/// Pause the throttle
|
||
|
///
|
||
|
/// Any tasks that have been submitted already will continue. However, no new tasks
|
||
|
/// will be run until the throttle is resumed.
|
||
|
virtual void Pause() = 0;
|
||
|
/// Resume the throttle
|
||
|
///
|
||
|
/// Allows task to be submitted again. If there is a max_concurrent_cost limit then
|
||
|
/// it will still apply.
|
||
|
virtual void Resume() = 0;
|
||
|
/// Return the number of tasks queued but not yet submitted
|
||
|
virtual std::size_t QueueSize() = 0;
|
||
|
|
||
|
/// Create a throttled view of a scheduler
|
||
|
///
|
||
|
/// Tasks added via this view will be subjected to the throttle and, if the tasks cannot
|
||
|
/// run immediately, will be placed into a queue.
|
||
|
///
|
||
|
/// Although a shared_ptr is returned it should generally be assumed that the caller
|
||
|
/// is being given exclusive ownership. The shared_ptr is used to share the view with
|
||
|
/// queued and submitted tasks and the lifetime of those is unpredictable. It is
|
||
|
/// important the caller keep the returned pointer alive for as long as they plan to add
|
||
|
/// tasks to the view.
|
||
|
///
|
||
|
/// \param scheduler a scheduler to submit tasks to after throttling
|
||
|
///
|
||
|
/// This can be the root scheduler, another throttled scheduler, or a task group. These
|
||
|
/// are all composable.
|
||
|
///
|
||
|
/// \param max_concurrent_cost the maximum amount of cost allowed to run at any one time
|
||
|
///
|
||
|
/// If a task is added that has a cost greater than max_concurrent_cost then its cost
|
||
|
/// will be reduced to max_concurrent_cost so that it is still possible for the task to
|
||
|
/// run.
|
||
|
///
|
||
|
/// \param queue the queue to use when tasks cannot be submitted
|
||
|
///
|
||
|
/// By default a FIFO queue will be used. However, a custom queue can be provided if
|
||
|
/// some tasks have higher priority than other tasks.
|
||
|
static std::shared_ptr<ThrottledAsyncTaskScheduler> Make(
|
||
|
AsyncTaskScheduler* scheduler, int max_concurrent_cost,
|
||
|
std::unique_ptr<Queue> queue = NULLPTR);
|
||
|
|
||
|
/// @brief Create a ThrottledAsyncTaskScheduler using a custom throttle
|
||
|
///
|
||
|
/// \see Make
|
||
|
static std::shared_ptr<ThrottledAsyncTaskScheduler> MakeWithCustomThrottle(
|
||
|
AsyncTaskScheduler* scheduler, std::unique_ptr<Throttle> throttle,
|
||
|
std::unique_ptr<Queue> queue = NULLPTR);
|
||
|
};
|
||
|
|
||
|
/// A utility to keep track of a collection of tasks
|
||
|
///
|
||
|
/// Often it is useful to keep track of some state that only needs to stay alive
|
||
|
/// for some small collection of tasks, or to perform some kind of final cleanup
|
||
|
/// when a collection of tasks is finished.
|
||
|
///
|
||
|
/// For example, when scanning, we need to keep the file reader alive while all scan
|
||
|
/// tasks run for a given file, and then we can gracefully close it when we finish the
|
||
|
/// file.
|
||
|
class ARROW_EXPORT AsyncTaskGroup : public AsyncTaskScheduler {
|
||
|
public:
|
||
|
/// Destructor for the task group
|
||
|
///
|
||
|
/// The destructor might trigger the finish callback. If the finish callback fails
|
||
|
/// then the error will be reported as a task on the scheduler.
|
||
|
///
|
||
|
/// Failure to destroy the async task group will not prevent the scheduler from
|
||
|
/// finishing. If the scheduler finishes before the async task group is done then
|
||
|
/// the finish callback will be run immediately when the async task group finishes.
|
||
|
///
|
||
|
/// If the scheduler has aborted then the finish callback will not run.
|
||
|
~AsyncTaskGroup() = default;
|
||
|
/// Create an async task group
|
||
|
///
|
||
|
/// The finish callback will not run until the task group is destroyed and all
|
||
|
/// tasks are finished so you will generally want to reset / destroy the returned
|
||
|
/// unique_ptr at some point.
|
||
|
///
|
||
|
/// \param scheduler The underlying scheduler to submit tasks to
|
||
|
/// \param finish_callback A callback that will be run only after the task group has
|
||
|
/// been destroyed and all tasks added by the group have
|
||
|
/// finished.
|
||
|
///
|
||
|
/// Note: in error scenarios the finish callback may not run. However, it will still,
|
||
|
/// of course, be destroyed.
|
||
|
static std::unique_ptr<AsyncTaskGroup> Make(AsyncTaskScheduler* scheduler,
|
||
|
FnOnce<Status()> finish_callback);
|
||
|
};
|
||
|
|
||
|
/// Create a task group that is also throttled
|
||
|
///
|
||
|
/// This is a utility factory that creates a throttled view of a scheduler and then
|
||
|
/// wraps that throttled view with a task group that destroys the throttle when finished.
|
||
|
///
|
||
|
/// \see ThrottledAsyncTaskScheduler
|
||
|
/// \see AsyncTaskGroup
|
||
|
/// \param target the underlying scheduler to submit tasks to
|
||
|
/// \param max_concurrent_cost the maximum amount of cost allowed to run at any one time
|
||
|
/// \param queue the queue to use when tasks cannot be submitted
|
||
|
/// \param finish_callback A callback that will be run only after the task group has
|
||
|
/// been destroyed and all tasks added by the group have finished
|
||
|
ARROW_EXPORT std::unique_ptr<ThrottledAsyncTaskScheduler> MakeThrottledAsyncTaskGroup(
|
||
|
AsyncTaskScheduler* target, int max_concurrent_cost,
|
||
|
std::unique_ptr<ThrottledAsyncTaskScheduler::Queue> queue,
|
||
|
FnOnce<Status()> finish_callback);
|
||
|
|
||
|
// Defined down here to avoid circular dependency between AsyncTaskScheduler and
|
||
|
// AsyncTaskGroup
|
||
|
template <typename T>
|
||
|
bool AsyncTaskScheduler::AddAsyncGenerator(std::function<Future<T>()> generator,
|
||
|
std::function<Status(const T&)> visitor,
|
||
|
std::string_view name) {
|
||
|
struct State {
|
||
|
State(std::function<Future<T>()> generator, std::function<Status(const T&)> visitor,
|
||
|
std::unique_ptr<AsyncTaskGroup> task_group, std::string_view name)
|
||
|
: generator(std::move(generator)),
|
||
|
visitor(std::move(visitor)),
|
||
|
task_group(std::move(task_group)),
|
||
|
name(name) {}
|
||
|
std::function<Future<T>()> generator;
|
||
|
std::function<Status(const T&)> visitor;
|
||
|
std::unique_ptr<AsyncTaskGroup> task_group;
|
||
|
std::string_view name;
|
||
|
};
|
||
|
struct SubmitTask : public Task {
|
||
|
explicit SubmitTask(std::unique_ptr<State> state_holder)
|
||
|
: state_holder(std::move(state_holder)) {}
|
||
|
|
||
|
struct SubmitTaskCallback {
|
||
|
SubmitTaskCallback(std::unique_ptr<State> state_holder, Future<> task_completion)
|
||
|
: state_holder(std::move(state_holder)),
|
||
|
task_completion(std::move(task_completion)) {}
|
||
|
void operator()(const Result<T>& maybe_item) {
|
||
|
if (!maybe_item.ok()) {
|
||
|
task_completion.MarkFinished(maybe_item.status());
|
||
|
return;
|
||
|
}
|
||
|
const auto& item = *maybe_item;
|
||
|
if (IsIterationEnd(item)) {
|
||
|
task_completion.MarkFinished();
|
||
|
return;
|
||
|
}
|
||
|
Status visit_st = state_holder->visitor(item);
|
||
|
if (!visit_st.ok()) {
|
||
|
task_completion.MarkFinished(std::move(visit_st));
|
||
|
return;
|
||
|
}
|
||
|
state_holder->task_group->AddTask(
|
||
|
std::make_unique<SubmitTask>(std::move(state_holder)));
|
||
|
task_completion.MarkFinished();
|
||
|
}
|
||
|
std::unique_ptr<State> state_holder;
|
||
|
Future<> task_completion;
|
||
|
};
|
||
|
|
||
|
Result<Future<>> operator()() {
|
||
|
Future<> task = Future<>::Make();
|
||
|
// Consume as many items as we can (those that are already finished)
|
||
|
// synchronously to avoid recursion / stack overflow.
|
||
|
while (true) {
|
||
|
Future<T> next = state_holder->generator();
|
||
|
if (next.TryAddCallback(
|
||
|
[&] { return SubmitTaskCallback(std::move(state_holder), task); })) {
|
||
|
return task;
|
||
|
}
|
||
|
ARROW_ASSIGN_OR_RAISE(T item, next.result());
|
||
|
if (IsIterationEnd(item)) {
|
||
|
task.MarkFinished();
|
||
|
return task;
|
||
|
}
|
||
|
ARROW_RETURN_NOT_OK(state_holder->visitor(item));
|
||
|
}
|
||
|
}
|
||
|
|
||
|
std::string_view name() const { return state_holder->name; }
|
||
|
|
||
|
std::unique_ptr<State> state_holder;
|
||
|
};
|
||
|
std::unique_ptr<AsyncTaskGroup> task_group =
|
||
|
AsyncTaskGroup::Make(this, [] { return Status::OK(); });
|
||
|
AsyncTaskGroup* task_group_view = task_group.get();
|
||
|
std::unique_ptr<State> state_holder = std::make_unique<State>(
|
||
|
std::move(generator), std::move(visitor), std::move(task_group), name);
|
||
|
task_group_view->AddTask(std::make_unique<SubmitTask>(std::move(state_holder)));
|
||
|
return true;
|
||
|
}
|
||
|
|
||
|
} // namespace util
|
||
|
} // namespace arrow
|