// Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#include <OgnTestConcurrencyDatabase.h>
#include <omni/graph/exec/unstable/IExecutionContext.h>
#include <omni/graph/exec/unstable/IExecutionCurrentThread.h>
#include <omni/graph/exec/unstable/Stamp.h>
#include "../include/omni/graph/test/ConcurrencyState.h"

#include <atomic>
#include <chrono>
#include <cstdlib>
#include <mutex>

// This is the implementation of the OGN node defined in OgnTestConcurrency.ogn

namespace omni {
namespace graph {
namespace test {

namespace
{
    std::mutex g_resetMutex;
    exec::unstable::SyncStamp g_currentSyncStamp;
    std::atomic<uint64_t> g_currentConcurrency{0};
}

class OgnTestConcurrency
{
public:
    static bool compute(OgnTestConcurrencyDatabase& db)
    {
        auto& result = db.outputs.result();
        result = false;

        // Try to access the shared mutex via a shared lock. If the node is unable to do
        // that (i.e. because TestIsolate currently holds exclusive ownership), we'll
        // return early with result set to false.
        if (!getConcurrencyState().sharedMutex.try_lock_shared())
        {
            return true;
        }

        if (auto task = exec::unstable::getCurrentTask())
        {
            auto currentStamp = task->getContext()->getExecutionStamp();
            if(!g_currentSyncStamp.inSync(currentStamp))
            {
                std::unique_lock<std::mutex> lock(g_resetMutex);
                if(g_currentSyncStamp.makeSync(currentStamp))
                {
                    g_currentConcurrency = 0;
                }
            }

            g_currentConcurrency++;

            uint64_t expectedConcurrency = db.inputs.expected();
            std::chrono::milliseconds timeOut(db.inputs.timeOut());

            auto start = std::chrono::high_resolution_clock::now();

            std::chrono::milliseconds currentDuration;
            do {
                currentDuration = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now()-start);
                if(g_currentConcurrency.load() == expectedConcurrency)
                {
                    result = true;
                    break;
                }

                std::this_thread::yield();
            } while(currentDuration < timeOut);
        }

        // Release the shared mutex.
        getConcurrencyState().sharedMutex.unlock_shared();

        return true;
    }
};

REGISTER_OGN_NODE()

} // test
} // namespace graph
} // namespace omni
