NET_ThreadedIO.h

Go to the documentation of this file.

/*
 * PROPRIETARY INFORMATION.  This software is proprietary to
 * Side Effects Software Inc., and is not to be reproduced,
 * transmitted, or disclosed in any way without written permission.
 *
 * NAME:        NET_ThreadedIO.h
 *
 * COMMENTS:
 *
 */

#ifndef __NET_THREADEDIO_H__
#define __NET_THREADEDIO_H__

#include "NET_API.h"

#include <UT/UT_BoostAsio.h>
#include <UT/UT_Condition.h>
#include <UT/UT_Function.h>
#include <UT/UT_Lock.h>
#include <UT/UT_NonCopyable.h>
#include <UT/UT_ThreadedIO.h>
#include <UT/UT_UniquePtr.h>

/// NOTE: This is ONLY a proof of concept as a replacement for UT_ThreadedIO
/// If accepted this would likely move to $UT
class NET_API NET_ThreadedIO
{
public:
    /// Base Class to perform IO thread tasks
    class NET_API Task
    {
    public:
        Task() = default;
        virtual ~Task() = default;
        UT_NON_COPYABLE(Task)

        virtual void invoke() = 0;
        virtual exint memoryEstimate() const = 0;
    };

private:
    class NET_API CallableTask : public Task
    {
    public:
        CallableTask(UT_Function<void()> callable, exint mem_usage = 0) :
            myClb(std::move(callable)),
            myMemUsage(mem_usage)
        {
        }

        void invoke() override
        {
            myClb();
        }

        exint memoryEstimate() const override
        {
            return myMemUsage;
        }
    private:
        UT_Function<void()> myClb;
        exint myMemUsage;
    };

public:
    NET_ThreadedIO(ASIO_IOContext& ctx, exint max_cost = 0) :
        myStrand(hboost::asio::make_strand(ctx.get_executor())),
        myMaxCost(max_cost),
        myCurrentCost(0)
    {}

    /// Adds the task to the IO queue.
    /// This will block until there is enough free space on the queue.
    void postTask(UT_UniquePtr<Task> task)
    {
        exint cost = task->memoryEstimate();

        UT_AutoLock l(myLock);

        // Check if we'll exceed the cost.
        while (true)
        {
            if (cost && myMaxCost >= 0 && (myCurrentCost > 0)
                && (myCurrentCost + cost > myMaxCost))
            {
                // Too expensive to add right now, so await a trigger.
                myCond.waitForTrigger(myLock);
            }
            else
            {
                break;
            }
        }

        myCurrentCost += cost;
        // Queue the task for execution when its the tasks turn to run
        hboost::asio::post(myStrand, [this, work_task = std::move(task), cost]() {
            {
                UT_AutoLock l(myLock);
                myCurrentCost -= cost;
                myCond.triggerAll();
            }

            work_task->invoke();
        });

        myCond.triggerAll();
    }

    /// Adds the task to the IO queue.
    /// This will block until there is enough free space on the queue.
    void postTask(UT_Function<void()> func, exint cost = 0)
    {
        postTask(UTmakeUnique<CallableTask>(std::move(func), cost));
    }

private:
    exint myCurrentCost;
    exint myMaxCost;

    mutable UT_Lock myLock;
    UT_Condition myCond;

    // Used as the queue for posting work items.
    hboost::asio::strand<hboost::asio::io_context::executor_type> myStrand;
};

NET_API NET_ThreadedIO* NETgetThreadedIO();


#endif // __NET_THREADEDIO_H__