dispatcher_impl.h

//
// Copyright 2025 Pixar
//
// Licensed under the terms set forth in the LICENSE.txt file available at
// https://openusd.org/license.
//
#ifndef PXR_BASE_WORK_TBB_DISPATCHER_IMPL_H
#define PXR_BASE_WORK_TBB_DISPATCHER_IMPL_H
 
#include "pxr/pxr.h"
#include "pxr/base/work/api.h"
 
// Blocked range is not used in this file, but this header happens to pull in
// the TBB version header in a way that works in all TBB versions.
#include <tbb/blocked_range.h>
#include <tbb/concurrent_vector.h>
#if TBB_INTERFACE_VERSION_MAJOR >= 12
#include <tbb/task_group.h>
#else
#include <tbb/task.h>
#endif
 
#include <functional>
#include <type_traits>
#include <utility>
 
PXR_NAMESPACE_OPEN_SCOPE
 
class WorkImpl_Dispatcher
{
public:
    WORK_API WorkImpl_Dispatcher();
 
    WORK_API ~WorkImpl_Dispatcher() noexcept;
 
    WorkImpl_Dispatcher(WorkImpl_Dispatcher const &) = delete;
    WorkImpl_Dispatcher &operator=(WorkImpl_Dispatcher const &) = delete;
 
    template <class Callable>
    inline void Run(Callable &&c) {
#if TBB_INTERFACE_VERSION_MAJOR >= 12
        _taskGroup.run(std::forward<Callable>(c));
#else
        _rootTask->spawn(
            *new(_rootTask->allocate_additional_child_of(*_rootTask))
                _InvokerTaskWrapper<typename std::remove_reference<Callable>::type>(
                std::forward<Callable>(c)));
#endif
    }
    
    WORK_API void Reset();
 
    WORK_API void Wait();
 
    WORK_API void Cancel();
 
#if TBB_INTERFACE_VERSION_MAJOR < 12
    template <class Fn>
    struct _InvokerTaskWrapper : public tbb::task {
        explicit _InvokerTaskWrapper(Fn &&fn)
            : _fn(std::move(fn)) {}
 
        explicit _InvokerTaskWrapper(Fn const &fn)
            : _fn(fn) {}
 
        virtual tbb::task* execute() {
            // In anticipation of OneTBB, ensure that _fn meets OneTBB's
            // requirement that a task's call operator must be const.
            const_cast<_InvokerTaskWrapper const *>(this)->_fn();
            return NULL;
        }
    private:
        Fn _fn;
    };
#endif
    // Task group context to run tasks in.
    tbb::task_group_context _context;
#if TBB_INTERFACE_VERSION_MAJOR >= 12
    // Custom task group that lets us implement thread safe concurrent wait.
    class _TaskGroup : public tbb::task_group {
    public:
        _TaskGroup(tbb::task_group_context& ctx) : tbb::task_group(ctx) {}
         inline tbb::detail::d1::wait_context& _GetInternalWaitContext();
    };
 
    _TaskGroup _taskGroup;
#else
    // Root task that allows us to cancel tasks invoked directly by this
    // dispatcher.
    tbb::empty_task* _rootTask;
#endif
 
};
 
PXR_NAMESPACE_CLOSE_SCOPE
 
 
#endif // PXR_BASE_WORK_TBB_DISPATCHER_IMPL_H