src/wasm/module_tier_ladder.cpp

#ifdef WITH_RIVE_SCRIPTING_WASM

#include "rive/wasm/module_tier_ladder.hpp"

#ifdef RIVE_WASM_TIER_LADDER

#include <algorithm>
#include <atomic>
#include <cerrno>
#include <cstdio>
#include <filesystem>
#include <fstream>
#include <iterator>
#include <fcntl.h>
#include <cstdlib>
#include <cstring>
#include <signal.h>
#include <spawn.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <unistd.h>

extern char** environ;

namespace rive
{

ModuleTierLadder& ModuleTierLadder::instance()
{
    static ModuleTierLadder* ladder = new ModuleTierLadder();
    return *ladder;
}

ModuleTierLadder::~ModuleTierLadder()
{
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        m_terminating = true;
        for (Job* job : m_running)
        {
            job->cancelled = true;
#ifndef RIVE_NX
            if (job->pid > 0)
            {
                kill(job->pid, SIGKILL);
            }
#endif
        }
    }
    m_workAvailable.notify_all();
    for (auto& worker : m_workers)
    {
        worker.join();
    }
}

void ModuleTierLadder::configure(const std::string& wamrcPath,
                                 const std::string& cacheDir)
{
    std::unique_lock<std::mutex> lock(m_mutex);
    m_wamrcPath = wamrcPath;
    m_versionProbed = false;
    if (m_wamrcPath.empty())
    {
        if (const char* env = getenv("RIVE_WAMRC"))
        {
            m_wamrcPath = env;
        }
    }
    m_cacheDir = cacheDir;
}

bool ModuleTierLadder::enabled()
{
    std::unique_lock<std::mutex> lock(m_mutex);
    if (m_wamrcPath.empty() && getenv("RIVE_WAMRC") != nullptr)
    {
        m_wamrcPath = getenv("RIVE_WAMRC");
    }
    // Self-configure from env so the first module load, which happens
    // before any host configure call, can stage its pristine bytes.
    if (m_cacheDir.empty() && !m_wamrcPath.empty())
    {
        const char* dirEnv = getenv("RIVE_AOT_CACHE_DIR");
#ifdef RIVE_NX
        // No temp_directory_path in this libc++, and the platform cannot map
        // executable pages anyway, so the ladder only runs when pointed at a
        // directory explicitly.
        m_cacheDir = dirEnv != nullptr ? dirEnv : "";
#else
        m_cacheDir =
            dirEnv != nullptr
                ? dirEnv
                : (std::filesystem::temp_directory_path() / "rive_aot_cache")
                      .string();
#endif
    }
    if (m_wamrcPath.empty() || m_cacheDir.empty())
    {
        return false;
    }
    wamrcVersion();
    return m_wamrcUsable;
}

void ModuleTierLadder::onArrival(ArrivalCallback callback)
{
    std::unique_lock<std::mutex> lock(m_mutex);
    m_arrival = std::move(callback);
}

const std::string& ModuleTierLadder::wamrcVersion()
{
    // Callers hold m_mutex. The version keys the cache so a compiler swap
    // reads as misses, never as stale native code.
    if (!m_versionProbed)
    {
        m_versionProbed = true;
        m_wamrcVersion = "unknown";
#ifndef RIVE_NX
        std::string cmd = m_wamrcPath + " --version 2>/dev/null";
        if (FILE* pipe = popen(cmd.c_str(), "r"))
        {
            char line[128] = {0};
            if (fgets(line, sizeof(line), pipe) != nullptr)
            {
                std::string v(line);
                while (!v.empty() && (v.back() == '\n' || v.back() == ' '))
                {
                    v.pop_back();
                }
                for (char& c : v)
                {
                    if (c == ' ' || c == '/')
                    {
                        c = '-';
                    }
                }
                if (!v.empty())
                {
                    m_wamrcVersion = v;
                }
            }
            pclose(pipe);
        }
        // Every compile passes --rive-interrupt, which an unpatched wamrc
        // rejects; say so once rather than fail each compile quietly.
        std::string probe =
            m_wamrcPath + " --rive-interrupt --version >/dev/null 2>&1";
        FILE* probePipe = popen(probe.c_str(), "r");
        m_wamrcUsable = probePipe != nullptr && pclose(probePipe) == 0;
        if (!m_wamrcUsable)
        {
            fprintf(stderr,
                    "wasm aot: %s lacks --rive-interrupt; build wamrc from "
                    "our patched WAMR (runtime/scripting/wamr_patches). AOT "
                    "is off.\n",
                    m_wamrcPath.c_str());
        }
#endif
    }
    return m_wamrcVersion;
}

// The module exports __riveHostBudget. A byte search suffices: a false
// hit only adds interrupt checks.
static bool asksHostBudget(const std::string& wasmPath)
{
    std::ifstream in(wasmPath, std::ios::binary);
    std::string bytes((std::istreambuf_iterator<char>(in)),
                      std::istreambuf_iterator<char>());
    return bytes.find("__riveHostBudget") != std::string::npos;
}

std::string ModuleTierLadder::keyedCacheDir()
{
    // The revision folds our wamrc flag choices into the cache key; bump
    // it whenever species flags change or stale artifacts get served on a
    // cache hit.
    std::string dir = m_cacheDir + "/" + wamrcVersion() + "-r8";
    mkdir(m_cacheDir.c_str(), 0755);
    mkdir(dir.c_str(), 0755);
    return dir;
}

std::string ModuleTierLadder::artifactName(uint64_t moduleKey,
                                           TierSpecies species)
{
    const char* pattern = "%016llx.aot";
    if (species == TierSpecies::o0)
    {
        pattern = "%016llx.o0.aot";
    }
    else if (species == TierSpecies::hw)
    {
        pattern = "%016llx.hw.aot";
    }
    char name[64];
    snprintf(name, sizeof(name), pattern, (unsigned long long)moduleKey);
    return name;
}

std::string ModuleTierLadder::artifactPath(uint64_t moduleKey,
                                           TierSpecies species)
{
    std::unique_lock<std::mutex> lock(m_mutex);
    if (m_wamrcPath.empty() || m_cacheDir.empty())
    {
        return std::string();
    }
    std::string path = keyedCacheDir() + "/" + artifactName(moduleKey, species);
    struct stat st;
    if (stat(path.c_str(), &st) == 0)
    {
        return path;
    }
    return std::string();
}

// The cache is shared across processes and a name must never be visible
// half-written: write a process-unique temp, then atomically rename.
static bool writeFileAtomic(const std::string& path, Span<const uint8_t> bytes)
{
    std::string tmpPath = path + "." + std::to_string(getpid()) + ".tmp";
    FILE* f = fopen(tmpPath.c_str(), "wb");
    if (f == nullptr)
    {
        return false;
    }
    size_t written = fwrite(bytes.data(), 1, bytes.size(), f);
    if (fclose(f) != 0 || written != bytes.size())
    {
        unlink(tmpPath.c_str());
        return false;
    }
    if (rename(tmpPath.c_str(), path.c_str()) != 0)
    {
        unlink(tmpPath.c_str());
        return false;
    }
    return true;
}

void ModuleTierLadder::stagePristine(uint64_t moduleKey,
                                     Span<const uint8_t> moduleBytes)
{
    if (!enabled())
    {
        return;
    }
    std::unique_lock<std::mutex> lock(m_mutex);
    char wasmName[64];
    snprintf(wasmName,
             sizeof(wasmName),
             "%016llx.wasm",
             (unsigned long long)moduleKey);
    std::string wasmPath = keyedCacheDir() + "/" + wasmName;
    struct stat st;
    if (stat(wasmPath.c_str(), &st) == 0)
    {
        return;
    }
    writeFileAtomic(wasmPath, moduleBytes);
}

std::string ModuleTierLadder::compileSync(uint64_t moduleKey,
                                          Span<const uint8_t> moduleBytes,
                                          TierSpecies species)
{
    if (!enabled())
    {
        return std::string();
    }
    std::string existing = artifactPath(moduleKey, species);
    if (!existing.empty())
    {
        return existing;
    }
    stagePristine(moduleKey, moduleBytes);
    std::string wasmPath;
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        char wasmName[64];
        snprintf(wasmName,
                 sizeof(wasmName),
                 "%016llx.wasm",
                 (unsigned long long)moduleKey);
        wasmPath = keyedCacheDir() + "/" + wasmName;
    }
    // Runs on the caller's thread and never joins m_running, so lane
    // supersession cannot kill it out from under the boot.
    Job job{std::string(), moduleKey, species, wasmPath, 0};
    if (!runWamrc(job))
    {
        return std::string();
    }
    return artifactPath(moduleKey, species);
}

void ModuleTierLadder::schedule(const std::string& laneId,
                                uint64_t moduleKey,
                                Span<const uint8_t> moduleBytes)
{
    if (!enabled())
    {
        return;
    }
    std::unique_lock<std::mutex> lock(m_mutex);
    uint64_t generation = m_nextGeneration++;
    m_laneGeneration[laneId] = generation;
    // Supersede the lane: queued stale jobs get skipped by the generation
    // check; running ones die now so the pool frees up for the new hash.
    // Same-hash jobs are not stale — a rebuild of identical bytes must not
    // kill its own compile — so they ride forward on the new generation
    // and their species skip re-queueing below.
    std::vector<TierSpecies> inFlight;
    for (Job* job : m_running)
    {
        if (job->laneId != laneId || job->generation >= generation)
        {
            continue;
        }
        // A job already cancelled by an earlier supersede will not produce
        // an artifact, so it cannot stand in for this request even when the
        // bytes match: let its species re-queue below.
        if (job->moduleKey == moduleKey && !job->cancelled)
        {
            job->generation = generation;
            inFlight.push_back(job->species);
        }
        else if (job->moduleKey != moduleKey && job->pid > 0 && !job->cancelled)
        {
            job->cancelled = true;
#ifndef RIVE_NX
            kill(job->pid, SIGKILL);
#endif
        }
    }
    for (Job& job : m_queue)
    {
        if (job.laneId == laneId && job.moduleKey == moduleKey &&
            !job.cancelled)
        {
            job.generation = generation;
            inFlight.push_back(job.species);
        }
    }

    std::string dir = keyedCacheDir();
    char wasmName[64];
    snprintf(wasmName,
             sizeof(wasmName),
             "%016llx.wasm",
             (unsigned long long)moduleKey);
    std::string wasmPath = dir + "/" + wasmName;

    // On guard-page builds the top rung is the hw species: no inline
    // bounds checks and growth never moves the memory base.
#ifdef RIVE_WASM_HW_BOUNDS
    constexpr TierSpecies kTopSpecies = TierSpecies::hw;
#else
    constexpr TierSpecies kTopSpecies = TierSpecies::o3;
#endif
    std::vector<TierSpecies> wanted;
    if (moduleBytes.size() <= kStraightToO3Bytes)
    {
        wanted = {kTopSpecies};
    }
    else
    {
        wanted = {TierSpecies::o0, kTopSpecies};
    }

    bool queued = false;
    for (TierSpecies species : wanted)
    {
        if (std::find(inFlight.begin(), inFlight.end(), species) !=
            inFlight.end())
        {
            continue;
        }
        std::string artifact = dir + "/" + artifactName(moduleKey, species);
        struct stat st;
        if (stat(artifact.c_str(), &st) == 0)
        {
            if (m_arrival)
            {
                Artifact ready{moduleKey, species, artifact};
                lock.unlock();
                m_arrival(ready);
                lock.lock();
            }
            continue;
        }
        if (!queued)
        {
            // A pristine stage from before the in-place load wins; the bytes
            // passed here may already be loader-rewritten.
            struct stat wasmStat;
            if (stat(wasmPath.c_str(), &wasmStat) != 0 &&
                !writeFileAtomic(wasmPath, moduleBytes))
            {
                return;
            }
            queued = true;
        }
        m_queue.push_back(
            Job{laneId, moduleKey, species, wasmPath, generation});
    }
    if (queued)
    {
        ensureWorkers();
        m_workAvailable.notify_all();
    }
}

void ModuleTierLadder::ensureWorkers()
{
    if (!m_workers.empty())
    {
        return;
    }
    for (unsigned i = 0; i < kWorkerCount; i++)
    {
        m_workers.emplace_back([this] { workerLoop(); });
    }
}

void ModuleTierLadder::workerLoop()
{
    std::unique_lock<std::mutex> lock(m_mutex);
    for (;;)
    {
        m_workAvailable.wait(lock, [this] {
            return m_terminating || !m_queue.empty();
        });
        if (m_terminating)
        {
            return;
        }
        Job job = m_queue.front();
        m_queue.pop_front();
        auto lane = m_laneGeneration.find(job.laneId);
        if (lane != m_laneGeneration.end() && job.generation < lane->second)
        {
            continue;
        }
        m_running.push_back(&job);
        lock.unlock();
        bool ok = runWamrc(job);
        lock.lock();
        m_running.erase(std::find(m_running.begin(), m_running.end(), &job));
        if (ok && !job.cancelled && m_arrival)
        {
            Artifact ready{job.moduleKey,
                           job.species,
                           keyedCacheDir() + "/" +
                               artifactName(job.moduleKey, job.species)};
            lock.unlock();
            m_arrival(ready);
            lock.lock();
        }
        if (m_queue.empty() && m_running.empty())
        {
            m_idle.notify_all();
        }
    }
}

bool ModuleTierLadder::runWamrc(Job& job)
{
    // Callers dropped m_mutex; only job fields and immutable config are
    // touched off-lock, except pid which the scheduler reads under lock to
    // kill superseded compiles.
    std::string finalPath;
    std::string wamrc;
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        finalPath =
            keyedCacheDir() + "/" + artifactName(job.moduleKey, job.species);
        wamrc = m_wamrcPath;
    }
    // Unique temp per compile: the same artifact can be raced by a worker
    // and a sync boot, or by two processes sharing the cache.
    static std::atomic<uint64_t> tmpSerial{0};
    std::string tmpPath = finalPath + "." + std::to_string(getpid()) + "-" +
                          std::to_string(++tmpSerial) + ".tmp";

    std::vector<std::string> args = {wamrc};
    if (job.species == TierSpecies::hw)
    {
        // Guard-page bounds; native stack checks stay sw (the runtime
        // builds with WASM_DISABLE_STACK_HW_BOUND_CHECK).
        args.push_back("--bounds-checks=0");
        args.push_back("--stack-bounds-checks=1");
    }
    else
    {
        // sw bounds run on every build; wamrc's default hw-bounds output
        // segfaults on runtimes without the guard-page trap handler.
        args.push_back("--bounds-checks=1");
    }
    if (job.species == TierSpecies::o0 || job.species == TierSpecies::hw)
    {
        args.push_back("--opt-level=0");
    }
    else
    {
        args.push_back("--opt-level=3");
    }
#if (defined(__APPLE__) || defined(_WIN32)) &&                                 \
    (defined(__aarch64__) || defined(_M_ARM64))
    // The platform clobbers x18; reserve it here so safety does not depend on
    // the wamrc we were handed. Features need an explicit cpu.
    args.push_back("--cpu=generic");
    args.push_back("--cpu-features=+reserve-x18");
#endif
    // Loop back-edges and calls check the watchdog only in modules that ask
    // for it; the module bytes, and so its cache key, decide.
    if (asksHostBudget(job.wasmPath))
    {
        args.push_back("--rive-interrupt");
    }
    args.push_back("-o");
    args.push_back(tmpPath);
    args.push_back(job.wasmPath);
    std::vector<char*> argv;
    for (auto& a : args)
    {
        argv.push_back(const_cast<char*>(a.c_str()));
    }
    argv.push_back(nullptr);

#if defined(RIVE_ANDROID) || defined(RIVE_NX)
    // No wamrc on device, and neither platform has posix_spawn; the ladder
    // never schedules compiles here.
    return false;
#else
    pid_t pid = -1;
    posix_spawn_file_actions_t actions;
    posix_spawn_file_actions_init(&actions);
    // A compiler's stdout is noise at editor runtime; stderr stays for
    // diagnosing failed compiles.
    posix_spawn_file_actions_addopen(&actions,
                                     STDOUT_FILENO,
                                     "/dev/null",
                                     O_WRONLY,
                                     0);
    int rc = posix_spawn(&pid,
                         wamrc.c_str(),
                         &actions,
                         nullptr,
                         argv.data(),
                         environ);
    posix_spawn_file_actions_destroy(&actions);
    if (rc != 0)
    {
        return false;
    }
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        job.pid = pid;
    }
    // The whole child runs below interactive priority; the editor's frame
    // loop must never contend with LLVM.
    setpriority(PRIO_PROCESS, pid, 10);

    int status = 0;
    while (waitpid(pid, &status, 0) < 0 && errno == EINTR)
    {
    }
    {
        std::unique_lock<std::mutex> lock(m_mutex);
        job.pid = -1;
    }
    if (job.cancelled || !WIFEXITED(status) || WEXITSTATUS(status) != 0)
    {
        if (!job.cancelled)
        {
            // wamrc prints usage errors to stdout, which is discarded; an
            // unlogged exit here is otherwise a silently dead lane.
            fprintf(stderr,
                    "[wasm] wamrc failed (exit %d) on %s\n",
                    WIFEXITED(status) ? WEXITSTATUS(status) : -1,
                    job.wasmPath.c_str());
        }
        unlink(tmpPath.c_str());
        return false;
    }
    // Atomic arrival: a partial artifact can never carry the final name.
    return rename(tmpPath.c_str(), finalPath.c_str()) == 0;
#endif
}

void ModuleTierLadder::drain()
{
    std::unique_lock<std::mutex> lock(m_mutex);
    m_idle.wait(lock, [this] { return m_queue.empty() && m_running.empty(); });
}

} // namespace rive

#endif // RIVE_WASM_TIER_LADDER
#endif // WITH_RIVE_SCRIPTING_WASM