resourceschedule_ffrt/include/tm/task_factory.h

/*
 * Copyright (c) 2023 Huawei Device Co., Ltd.
 * Licensed 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.
 */

#ifndef TASK_FACTORY_HPP
#define TASK_FACTORY_HPP

#include "tm/task_base.h"
#include "util/cb_func.h"
#include "util/slab.h"

namespace ffrt {
template <typename T>
class TaskFactory {
public:
    static TaskFactory<T>& Instance();

    static T* Alloc()
    {
        return Instance().alloc_();
    }

    static void Free(T* task)
    {
        Instance().free_(task);
    }

    static void Free_(T* task)
    {
        if (Instance().free__ != nullptr) {
            Instance().free__(task);
        }
    }

    static std::vector<void*> GetUnfreedMem()
    {
        if (Instance().getUnfreedMem_ != nullptr) {
            return Instance().getUnfreedMem_();
        }
        return {};
    }

    static std::size_t GetUnfreedMemSize()
    {
        if (Instance().getUnfreedMemSize_ != nullptr) {
            return Instance().getUnfreedMemSize_();
        }
        return 0;
    }

    static std::vector<void*> GetUnfreedTasksFiltered()
    {
        LockMem();
        std::vector<void*> unfreed = GetUnfreedMem();
        // Filter out tasks where the reference count increment failed.
        unfreed.erase(
            std::remove_if(unfreed.begin(), unfreed.end(),
                [](void* task) {
                    return !IncDeleteRefIfPositive(reinterpret_cast<TaskBase*>(task));
                }),
            unfreed.end()
        );
        UnlockMem();
        return unfreed;
    }

    static bool HasBeenFreed(T* task)
    {
        if (Instance().hasBeenFreed_ != nullptr) {
            return Instance().hasBeenFreed_(task);
        }
        return true;
    }

    static void LockMem()
    {
        if (Instance().lockMem_ != nullptr) {
            Instance().lockMem_();
        }
    }

    static void UnlockMem()
    {
        if (Instance().unlockMem_ != nullptr) {
            Instance().unlockMem_();
        }
    }

    static void RegistCb(
        typename TaskAllocCB<T>::Alloc alloc,
        typename TaskAllocCB<T>::Free free,
        typename TaskAllocCB<T>::Free_ free_ = nullptr,
        typename TaskAllocCB<T>::GetUnfreedMem getUnfreedMem = nullptr,
        typename TaskAllocCB<T>::GetUnfreedMemSize getUnfreedMemSize = nullptr,
        typename TaskAllocCB<T>::HasBeenFreed hasBeenFreed = nullptr,
        typename TaskAllocCB<T>::LockMem lockMem = nullptr,
        typename TaskAllocCB<T>::UnlockMem unlockMem = nullptr)
    {
        Instance().alloc_ = alloc;
        Instance().free_ = free;
        Instance().free__ = free_;
        Instance().getUnfreedMem_ = getUnfreedMem;
        Instance().getUnfreedMemSize_ = getUnfreedMemSize;
        Instance().hasBeenFreed_ = hasBeenFreed;
        Instance().lockMem_ = lockMem;
        Instance().unlockMem_ = unlockMem;
    }

private:
    typename TaskAllocCB<T>::Alloc alloc_ = nullptr;
    typename TaskAllocCB<T>::Free free_ = nullptr;
    typename TaskAllocCB<T>::Free_ free__ = nullptr;
    typename TaskAllocCB<T>::GetUnfreedMem getUnfreedMem_ = nullptr;
    typename TaskAllocCB<T>::GetUnfreedMemSize getUnfreedMemSize_ = nullptr;
    typename TaskAllocCB<T>::HasBeenFreed hasBeenFreed_ = nullptr;
    typename TaskAllocCB<T>::LockMem lockMem_ = nullptr;
    typename TaskAllocCB<T>::UnlockMem unlockMem_ = nullptr;
};

template <typename T>
class TaskMemScopedLock {
public:
    TaskMemScopedLock()
    {
        TaskFactory<T>::LockMem();
    }

    ~TaskMemScopedLock()
    {
        TaskFactory<T>::UnlockMem();
    }
};

template <typename FactoryTaskType, typename AllocatorTaskType = FactoryTaskType>
void RegisterTaskFactoryCallbacks()
{
    ffrt::TaskFactory<FactoryTaskType>::RegistCb(
        [] () -> FactoryTaskType* {
            return ffrt::SimpleAllocator<AllocatorTaskType>::AllocMem();
        },
        [] (FactoryTaskType* task) {
            ffrt::SimpleAllocator<AllocatorTaskType>::FreeMem(static_cast<AllocatorTaskType*>(task));
        },
        [] (FactoryTaskType* task) {
            ffrt::SimpleAllocator<AllocatorTaskType>::FreeMem_(static_cast<AllocatorTaskType*>(task));
        },
        ffrt::SimpleAllocator<AllocatorTaskType>::getUnfreedMem,
        ffrt::SimpleAllocator<AllocatorTaskType>::getUnfreedMemSize,
        [] (FactoryTaskType* task) {
            return ffrt::SimpleAllocator<AllocatorTaskType>::HasBeenFreed(static_cast<AllocatorTaskType*>(task));
        },
        ffrt::SimpleAllocator<AllocatorTaskType>::LockMem,
        ffrt::SimpleAllocator<AllocatorTaskType>::UnlockMem);
}
} // namespace ffrt

#endif