omni/extras/UniqueApp.h

File members: omni/extras/UniqueApp.h

// Copyright (c) 2021-2023, 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.
//
#pragma once

#include "../core/Omni.h"

#include <string>

#if OMNI_PLATFORM_WINDOWS
#    include "../../carb/CarbWindows.h"
#    include "../../carb/extras/WindowsPath.h"
#else
#    include <unistd.h>
#    include <sys/types.h>
#    include <sys/stat.h>
#    include <fcntl.h>
#endif

namespace omni
{
namespace extras
{

class UniqueApp
{
public:
    UniqueApp() = default;

    // Copy construction is disabled because this is a ref-counted object.
    UniqueApp(const UniqueApp&) = delete;

    UniqueApp(UniqueApp&& other)
    {
        *this = std::move(other);
    }

    UniqueApp& operator=(UniqueApp&& other)
    {
        m_guardPath = other.m_guardPath;
        m_guardName = other.m_guardName;
        m_launchGuard = other.m_launchGuard;
        m_exitGuard = other.m_exitGuard;

        other.m_guardPath = ".";
        other.m_guardName = kDefaultNamePrefix;
        other.m_launchGuard = kBadFileHandle;
        other.m_exitGuard = kBadFileHandle;

        return *this;
    }

    UniqueApp(const char* guardPath, const char* guardName)
    {
        setGuardPath(guardPath);
        setGuardName(guardName);
    }

    ~UniqueApp()
    {
        // Note: We are *intentionally* leaking any created guard objects here.  If either of
        //       them were to be closed on destruction of the object, undesirable effects would
        //       result:
        //          * If the launch guard was created, closing it would allow other instances
        //            of the unique app to successfully launch.
        //          * If this process 'connected' to the unique app process, closing the exit
        //            guard object would remove its reference and could allow the unique app to
        //            exit prematurely thinking all of its 'clients' had exited already.
        //
        //       An alternative to this would require forcing all callers to store the created
        //       object at a global level where it would live for the duration of the process.
        //       While this is certainly possible, enforcing that would be difficult at best.
    }

    void setGuardPath(const char* path)
    {
        if (path[0] == 0 || path == nullptr)
            path = ".";

        m_guardPath = path;
    }

    void setGuardName(const char* name)
    {
        if (name[0] == 0 || name == nullptr)
            name = kDefaultNamePrefix;

        m_guardName = name;
    }

    bool createLaunchGuard()
    {
        FileHandle handle;

        // this process has already created the launch guard object -> nothing to do => succeed.
        if (m_launchGuard != kBadFileHandle)
            return true;

#if OMNI_PLATFORM_WINDOWS
        std::string name = m_guardName + kLaunchLockExtension;

        handle = CreateEventA(nullptr, false, false, name.c_str());

        if (handle == nullptr)
            return false;

        if (GetLastError() == CARBWIN_ERROR_ALREADY_EXISTS)
        {
            CloseHandle(handle);
            return false;
        }
#else
        std::string path = _getGuardName(kLaunchLockExtension).c_str();
        handle = _openFile(path.c_str());

        if (handle == kBadFileHandle)
            return false;

        if (!_lockFile(handle, LockType::eExclusive))
        {
            close(handle);
            return false;
        }
#endif

        // save the exit guard event so we can destroy it later.
        m_launchGuard = handle;

        // intentionally 'leak' the guard handle here.  This will keep the handle open for the
        // entire remaining duration of the process.  This is important because the existence
        // of the guard object is what is used by other host apps to determine if the unique
        // app is already running.  Once the unique app process exits (or it makes a call to
        // @ref destroyLaunchGuard()), the OS will automatically close the guard object.
        return true;
    }

    void destroyLaunchGuard()
    {
        FileHandle fp = m_launchGuard;

        m_launchGuard = kBadFileHandle;
        _closeFile(fp);
    }

    bool checkLaunchGuard()
    {
#if OMNI_PLATFORM_WINDOWS
        HANDLE event;
        DWORD error;
        std::string name = m_guardName + kLaunchLockExtension;

        event = CreateEventA(nullptr, false, false, name.c_str());

        // failed to create the event handle (?!?) => fail.
        if (event == nullptr)
            return false;

        error = GetLastError();
        CloseHandle(event);

        return error == CARBWIN_ERROR_ALREADY_EXISTS;
#else
        FileHandle fp;
        bool success;

        fp = _openFile(_getGuardName(kLaunchLockExtension).c_str());

        if (fp == kBadFileHandle)
            return false;

        success = _lockFile(fp, LockType::eExclusive, LockAction::eTest);
        _closeFile(fp);
        return !success;
#endif
    }

    bool connectClientProcess()
    {
        bool success;
        FileHandle fp;

        // this object has already 'connected' to the unique app -> nothing to do => succeed.
        if (m_exitGuard != kBadFileHandle)
            return true;

        fp = _openFile(_getGuardName(kExitLockExtension).c_str());

        // failed to open the guard file (?!?) => fail.
        if (fp == kBadFileHandle)
            return false;

        // grab a shared lock to the file.  This will allow all clients to still also grab a
        // shared lock but will prevent the unique app from grabbing its exclusive lock
        // that it uses to determine whether any client apps are still 'connected'.
        success = _lockFile(fp, LockType::eShared);

        if (!success)
            _closeFile(fp);

        // save the exit guard handle in case we need to explicitly disconnect this app later.
        else
            m_exitGuard = fp;

        // intentionally 'leak' the file handle here.  Since the file lock is associated with
        // the file handle, we can't close it until the process exits otherwise the unique app
        // will think this client has 'disconnected'.  If we leak the handle, the OS will take
        // care of closing the handle when the process exits and that will automatically remove
        // this process's file lock.
        return success;
    }

    void disconnectClientProcess()
    {
        FileHandle fp = m_exitGuard;

        m_exitGuard = kBadFileHandle;
        _closeFile(fp);
    }

    bool haveAllClientsExited()
    {
        bool success;
        FileHandle fp;
        std::string path;

        path = _getGuardName(kExitLockExtension);
        fp = _openFile(path.c_str());

        // failed to open the guard file (?!?) => fail.
        if (fp == kBadFileHandle)
            return false;

        success = _lockFile(fp, LockType::eExclusive, LockAction::eTest);
        _closeFile(fp);

        // the file lock was successfully acquired -> no more clients are 'connected' => delete
        //   the lock file as a final cleanup step.
        if (success)
            _deleteFile(path.c_str());

        // all the clients have 'disconnected' when we're able to successfully grab an exclusive
        // lock on the file.  This means that all of the clients have exited and the OS has
        // released their shared locks on the file thus allowing us to grab an exclusive lock.
        return success;
    }

private:
    static constexpr const char* kLaunchLockExtension = ".lock";

    static constexpr const char* kExitLockExtension = ".exit";

    static constexpr const char* kDefaultNamePrefix = "nvidia-unique-app";

#if OMNI_PLATFORM_WINDOWS
    using FileHandle = HANDLE;
    static constexpr FileHandle kBadFileHandle = CARBWIN_INVALID_HANDLE_VALUE;
#else
    using FileHandle = int;

    static constexpr FileHandle kBadFileHandle = -1;
#endif

    enum class LockType
    {
        eShared,

        eExclusive,
    };

    enum class LockAction
    {
        eSet,
        eTest,
    };

    std::string _getGuardName(const char* extension) const
    {
        return m_guardPath + "/" + m_guardName + extension;
    }

    static FileHandle _openFile(const char* filename)
    {
        // make sure the path up to the file exists.  Note that this will likely just fail to
        // open the file below if creating the directories fail.
        _makeDirectories(filename);

#if OMNI_PLATFORM_WINDOWS
        std::wstring pathW = carb::extras::convertCarboniteToWindowsPath(filename);
        return CreateFileW(pathW.c_str(), CARBWIN_GENERIC_READ | CARBWIN_GENERIC_WRITE,
                           CARBWIN_FILE_SHARE_READ | CARBWIN_FILE_SHARE_WRITE, nullptr, CARBWIN_OPEN_ALWAYS, 0, nullptr);
#else
        return open(filename, O_CREAT | O_TRUNC | O_RDWR, S_IRUSR | S_IWUSR | S_IROTH | S_IRGRP);
#endif
    }

    static void _closeFile(FileHandle fp)
    {
#if OMNI_PLATFORM_WINDOWS
        CloseHandle(fp);
#else
        close(fp);
#endif
    }

    static void _deleteFile(const char* filename)
    {
#if OMNI_PLATFORM_WINDOWS
        std::wstring pathW = carb::extras::convertCarboniteToWindowsPath(filename);
        DeleteFileW(pathW.c_str());
#else
        unlink(filename);
#endif
    }

    static bool _lockFile(FileHandle fp, LockType type, LockAction action = LockAction::eSet)
    {
#if OMNI_PLATFORM_WINDOWS
        BOOL success;
        CARBWIN_OVERLAPPED ov = {};
        DWORD flags = CARBWIN_LOCKFILE_FAIL_IMMEDIATELY;

        if (type == LockType::eExclusive)
            flags |= CARBWIN_LOCKFILE_EXCLUSIVE_LOCK;

        success = LockFileEx(fp, flags, 0, 1, 0, reinterpret_cast<LPOVERLAPPED>(&ov));

        if (action == LockAction::eTest)
            UnlockFileEx(fp, 0, 1, 0, reinterpret_cast<LPOVERLAPPED>(&ov));

        return success;
#else
        int result;
        struct flock fl;

        fl.l_type = (type == LockType::eExclusive ? F_WRLCK : F_RDLCK);
        fl.l_whence = SEEK_SET;
        fl.l_start = 0;
        fl.l_len = 1;
        fl.l_pid = 0;
        result = fcntl(fp, (action == LockAction::eTest ? F_GETLK : F_SETLK), &fl);

        if (result != 0)
            return false;

        if (action == LockAction::eTest)
            return fl.l_type == F_UNLCK;

        return true;
#endif
    }

    static bool _makeDirectories(const std::string& path)
    {
        size_t start = 0;
#if CARB_PLATFORM_WINDOWS
        constexpr const char* separators = "\\/";

        if (path.size() > 3 && (path[1] == ':' || path[0] == '/' || path[0] == '\\'))
            start = 3;
#elif CARB_POSIX
        constexpr const char* separators = "/";

        if (path.size() > 1 && path[0] == '/')
            start = 1;
#else
        CARB_UNSUPPORTED_PLATFORM();
#endif

        for (;;)
        {
            size_t pos = path.find_first_of(separators, start);
            std::string next;

            if (pos == std::string::npos)
                break;

            next = path.substr(0, pos);
            start = pos + 1;

#if CARB_PLATFORM_WINDOWS
            std::wstring pathW = carb::extras::convertCarboniteToWindowsPath(next);
            DWORD attr = GetFileAttributesW(pathW.c_str());

            if (attr != CARBWIN_INVALID_FILE_ATTRIBUTES)
            {
                // already a directory -> nothing to do => skip it.
                if ((attr & CARBWIN_FILE_ATTRIBUTE_DIRECTORY) != 0)
                    continue;

                // exists but not a directory -> cannot continue => fail.
                else
                    return false;
            }

            // failed to create the directory -> cannot continue => fail.
            if (!CreateDirectoryW(pathW.c_str(), nullptr))
                return false;
#elif CARB_POSIX
            struct stat st;

            if (stat(next.c_str(), &st) == 0)
            {
                // already a directory -> nothing to do => skip it.
                if (S_ISDIR(st.st_mode))
                    continue;

                // exists but not a directory -> cannot continue => fail.
                else
                    return false;
            }

            // failed to create the directory -> cannot continue => fail.
            if (mkdir(next.c_str(), S_IRWXU | S_IRGRP | S_IROTH) != 0)
                return false;
#else
            CARB_UNSUPPORTED_PLATFORM();
#endif
        }

        return true;
    }

    std::string m_guardPath = ".";

    std::string m_guardName = kDefaultNamePrefix;

    FileHandle m_launchGuard = kBadFileHandle;

    FileHandle m_exitGuard = kBadFileHandle;
};

} // namespace extras
} // namespace omni