codeCache_ut.cpp

/*
 * Copyright 2026, Datadog, Inc.
 * SPDX-License-Identifier: Apache-2.0
 */

#include <gtest/gtest.h>
#include "codeCache.h"

#include <memory>
#include <stdio.h>

// memoryUsage() must reflect the actual per-symbol name length, not a fixed
// per-symbol constant: adding one symbol grows the reported usage by exactly
// the allocation size of its name string.
TEST(CodeCacheTest, MemoryUsageCountsSymbolNameLength) {
    CodeCache cc("testlib");
    long long base = cc.memoryUsage();

    const char *name =
        "a_long_symbol_name_well_beyond_sizeof_NativeFunc_padding_xxxxxxxxxxxx";
    cc.add(reinterpret_cast<const void *>(0x1000), 16, name);

    EXPECT_EQ(NativeFunc::allocSize(name), (size_t)(cc.memoryUsage() - base));
}

// A longer name costs more than a shorter one — the reported size scales with
// the actual symbol-name length.
TEST(CodeCacheTest, MemoryUsageGrowsWithNameLength) {
    CodeCache shortc("lib");
    CodeCache longc("lib");
    long long short_base = shortc.memoryUsage();
    long long long_base = longc.memoryUsage();

    shortc.add(reinterpret_cast<const void *>(0x1000), 8, "f");
    longc.add(reinterpret_cast<const void *>(0x1000), 8,
              "an_extremely_long_symbol_name_padded_out_further_and_further_1234");

    EXPECT_GT(longc.memoryUsage() - long_base, shortc.memoryUsage() - short_base);
}

// A freshly-constructed, empty cache reports exactly the blob array (counted at
// the full CodeBlob size, not a pointer's worth) plus its own name — nothing
// more. An exact check so over-counting or an inflated constant would fail.
TEST(CodeCacheTest, MemoryUsageCountsFullBlobArray) {
    CodeCache cc("lib");
    EXPECT_EQ((long long)INITIAL_CODE_CACHE_CAPACITY * sizeof(CodeBlob)
                  + (long long)NativeFunc::allocSize("lib"),
              cc.memoryUsage());
}

// The DWARF unwind table contributes length * sizeof(FrameDesc). Populate it so
// a regression (e.g. * flipped to /) in that term is caught; the table is
// zero-length in the other tests.
TEST(CodeCacheTest, MemoryUsageCountsDwarfTable) {
    CodeCache cc("lib");
    long long base = cc.memoryUsage();

    const int entries = 7;
    // setDwarfTable() takes ownership; ~CodeCache frees it. Contents are
    // irrelevant here — memoryUsage() reads only the length.
    FrameDesc *table = (FrameDesc *)malloc(entries * sizeof(FrameDesc));
    cc.setDwarfTable(table, entries);

    EXPECT_EQ((long long)entries * sizeof(FrameDesc), cc.memoryUsage() - base);
}

// Copying a CodeCache must deep-copy each blob's name string. With the previous
// shallow memcpy the copy shared the originals' _name pointers, so destroying
// both double-freed them. Here both copies are populated, sized equally, and
// destroyed — a double-free would abort (or trip ASan).
TEST(CodeCacheTest, CopyIsDeepAndDoesNotDoubleFree) {
    CodeCache src("lib");
    src.add(reinterpret_cast<const void *>(0x1000), 16, "sym_one");
    src.add(reinterpret_cast<const void *>(0x2000), 16, "sym_two_longer_name");
    long long expected = src.memoryUsage();

    {
        CodeCache copy(src);
        EXPECT_EQ(expected, copy.memoryUsage());
    }  // copy destroyed — must not free src's name strings

    EXPECT_EQ(expected, src.memoryUsage());  // src's names still intact
}  // src destroyed — no double free

class CodeCacheArrayTest : public ::testing::Test {
protected:
    CodeCacheArray arr;

    void TearDown() override {
        arr.freeForTest();
    }
};

// CodeCacheArray::memoryUsage() is the aggregation production actually consumes
// (Profiler::dump()). Verify it sums the per-library values — guards the loop
// bound and the accumulation against regressions the per-instance tests miss.
TEST_F(CodeCacheArrayTest, ArrayMemoryUsageSumsLibraries) {
    auto a = std::make_unique<CodeCache>("liba");
    auto b = std::make_unique<CodeCache>("libb");
    // Populate before registering: add() must run pre-publication.
    a->add(reinterpret_cast<const void *>(0x1000), 16, "alpha_symbol");
    b->add(reinterpret_cast<const void *>(0x2000), 16, "beta_symbol_longer_name");

    long long expected = (long long)a->memoryUsage() + (long long)b->memoryUsage();
    ASSERT_TRUE(arr.add(a.get()));
    a.release();
    ASSERT_TRUE(arr.add(b.get()));
    b.release();
    EXPECT_EQ(expected, (long long)arr.memoryUsage());
}

// freeForTest() must restore the NM gauges to their pre-add baseline and
// leave the array reusable for a later cycle.
TEST_F(CodeCacheArrayTest, FreeForTestRestoresAccountingAndAllowsReuse) {
    const long long live_before = NativeMem::live(NM_NATIVE_SYMBOLS);
    const long long overhead_before = NativeMem::overhead(NM_NATIVE_SYMBOLS);
    arr.freeForTest();
    EXPECT_EQ(live_before, NativeMem::live(NM_NATIVE_SYMBOLS));
    EXPECT_EQ(overhead_before, NativeMem::overhead(NM_NATIVE_SYMBOLS));

    for (int cycle = 0; cycle < 2; cycle++) {
        auto a = std::make_unique<CodeCache>("cleanup_liba");
        auto b = std::make_unique<CodeCache>("cleanup_libb_with_longer_name");
        a->add(reinterpret_cast<const void *>(0x1000), 16, "alpha_symbol");
        b->add(reinterpret_cast<const void *>(0x2000), 32, "beta_symbol_longer_name");
        b->add(reinterpret_cast<const void *>(0x3000), 8, "gamma");
        CodeCache *first = a.get();
        CodeCache *second = b.get();
        const long long live = a->memoryUsage() + b->memoryUsage();
        const long long overhead = a->nameOverhead() + b->nameOverhead();
        ASSERT_GT(live, 0);

        ASSERT_TRUE(arr.add(a.get()));
        a.release();
        ASSERT_TRUE(arr.add(b.get()));
        b.release();
        ASSERT_EQ(2, arr.count());
        EXPECT_EQ(first, arr[0]);
        EXPECT_EQ(second, arr[1]);
        EXPECT_EQ(live, static_cast<long long>(arr.memoryUsage()));
        EXPECT_EQ(live_before + live, NativeMem::live(NM_NATIVE_SYMBOLS));
        EXPECT_EQ(overhead_before + overhead, NativeMem::overhead(NM_NATIVE_SYMBOLS));

        for (int cleanup = 0; cleanup < 2; cleanup++) {
            arr.freeForTest();
            EXPECT_EQ(live_before, NativeMem::live(NM_NATIVE_SYMBOLS));
            EXPECT_EQ(overhead_before, NativeMem::overhead(NM_NATIVE_SYMBOLS));
            EXPECT_EQ(0, arr.count());
            EXPECT_EQ(0u, arr.memoryUsage());
            for (int slot = 0; slot < MAX_NATIVE_LIBS; slot++) {
                EXPECT_EQ(nullptr, arr.at(slot)) << slot;
            }
        }
    }
}

// The two-phase add used under JitCodeCache's stubs lock: prepareAdd() hands
// over a grown array only when the cache is full, commitAdd() returns the
// replaced array exactly then, and the result matches add()'s documented
// accounting (blob array at the grown capacity plus each name allocation).
TEST(CodeCacheTest, TwoPhaseAddGrowsOnlyWhenFull) {
    CodeCache cc("lib");
    long long names = 0;
    char name[32];
    const int ADDS = INITIAL_CODE_CACHE_CAPACITY + 1;
    for (int i = 0; i < ADDS; i++) {
        snprintf(name, sizeof(name), "sym_%d", i);
        CodeCache::PreparedAdd prepared;
        cc.prepareAdd(name, prepared);
        bool full = i == INITIAL_CODE_CACHE_CAPACITY;
        EXPECT_EQ(full, prepared.grown != nullptr) << i;
        CodeBlob *retired = cc.commitAdd(reinterpret_cast<const void *>(0x1000 + i * 16),
                                         16, prepared);
        EXPECT_EQ(full, retired != nullptr) << i;
        delete[] retired;
        names += (long long)NativeFunc::allocSize(name);
    }
    EXPECT_EQ((long long)INITIAL_CODE_CACHE_CAPACITY * 2 * sizeof(CodeBlob)
                  + (long long)NativeFunc::allocSize("lib") + names,
              cc.memoryUsage());
    for (int i = 0; i < ADDS; i++) {
        snprintf(name, sizeof(name), "sym_%d", i);
        EXPECT_EQ(reinterpret_cast<const void *>(0x1000 + i * 16), cc.findSymbol(name)) << i;
    }
}

// prepareAdd() keeps add()'s replacement of non-printable name characters.
TEST(CodeCacheTest, TwoPhaseAddReplacesNonPrintableName) {
    CodeCache cc("lib");
    CodeCache::PreparedAdd prepared;
    cc.prepareAdd("a\x01" "b", prepared);
    delete[] cc.commitAdd(reinterpret_cast<const void *>(0x1000), 16, prepared);
    EXPECT_EQ(reinterpret_cast<const void *>(0x1000), cc.findSymbol("a?b"));
}

// A PreparedAdd that is never committed frees its allocations and leaves the
// cache unchanged, including when it holds a grown blob array. LeakSanitizer
// (ASan builds) reports the allocations if the destructor does not free them.
TEST(CodeCacheTest, UncommittedPrepareFreesAndLeavesCacheUnchanged) {
    CodeCache cc("lib");
    char name[32];
    for (int i = 0; i < INITIAL_CODE_CACHE_CAPACITY; i++) {
        snprintf(name, sizeof(name), "sym_%d", i);
        cc.add(reinterpret_cast<const void *>(0x1000 + i * 16), 16, name);
    }
    long long usage = cc.memoryUsage();
    {
        CodeCache::PreparedAdd prepared;
        cc.prepareAdd("never_committed", prepared);
        EXPECT_NE(nullptr, prepared.name);
        EXPECT_NE(nullptr, prepared.grown);
    }
    EXPECT_EQ(usage, cc.memoryUsage());
    EXPECT_EQ(nullptr, cc.findSymbol("never_committed"));
}

// commitAdd() takes ownership: the PreparedAdd no longer refers to the name or
// the grown array, so its destructor cannot free storage the cache uses.
TEST(CodeCacheTest, CommitAddTakesOwnership) {
    CodeCache cc("lib");
    CodeCache::PreparedAdd prepared;
    cc.prepareAdd("owned", prepared);
    delete[] cc.commitAdd(reinterpret_cast<const void *>(0x1000), 16, prepared);
    EXPECT_EQ(nullptr, prepared.name);
    EXPECT_EQ(nullptr, prepared.grown);
    EXPECT_EQ(reinterpret_cast<const void *>(0x1000), cc.findSymbol("owned"));
}