EIC code displayed by LXR master/​include/​arrow/​util/​async_util.h	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

File indexing completed on 2025-08-28 08:27:02

 // 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

This page was automatically generated by the 2.3.7 LXR engine. The LXR team