fix: close issue #249 — Emscripten OOM in CKKS unit tests

BAder82t · pascoec · commit dcf4952fef7c · 2026-07-08T14:21:56.000-04:00
PR #1179 added ClearBootstrapPrecom and ClearSchemeSwitchPrecom helpers but did not wire them into ReleaseAllContexts. This change completes that work, then re-enables the EVAL_FAST_ROTATION and CONTEXT_WITH_SERTYPE tests that PR #985 had to guard under !defined(__EMSCRIPTEN__) for the original std::bad_alloc failure. Library changes (no hot-path touched): - ReleaseAllContexts now walks every live context and calls ClearAllCKKSCaches before clearing the static maps. - Adds malloc_zone_pressure_relief(Apple) and _heapmin(MSVC) to AllocTrim. - Add ClearAllCKKSCaches convenience on CryptoContextImpl. - ClearBootstrapPrecom and ClearSchemeSwitchPrecom bodies made noexcept-safe so the walk does not throw on contexts that have not Enable(FHE)'d or Enable(SCHEMESWITCH)'d. Test changes: - Mid-test ClearAllCKKSCaches between JSON and BINARY rounds in UnitTestCKKSrnsSerialize. - UTCKKSCacheClear extends its WITH_TCM skip to __EMSCRIPTEN__ (the heap-probe helper has no emmalloc backend). - Remove the !defined(__EMSCRIPTEN__) guards around the affected rows. CI: - Add .github/workflows/emscripten.yml: builds with emsdk 3.1.59 and runs pke_tests/core_tests under node in two slices. Validated on macOS + WASM: - Native: pke_tests 1892/1892 pass, core_tests 158/158 pass. - WASM: CKKS slice 939/941 pass (2 heap-probe tests skipped), BFV+BGV+binfhe slice 324/324 pass, peak 1.6 GB. Closes #249.
diff --git a/.github/workflows/emscripten.yml b/.github/workflows/emscripten.yml
@@ -0,0 +1,61 @@
+name: Emscripten
+
+# Builds with emcc and runs the WASM unit tests under node. Guards against
+# re-introducing the cache-leak class that issue #249 tracked.
+
+on:
+  pull_request:
+    branches: [dev]
+  push:
+    branches: [dev, main]
+  workflow_dispatch:
+
+jobs:
+  build-and-test:
+    runs-on: ubuntu-22.04
+    timeout-minutes: 60
+    steps:
+      - uses: actions/checkout@v5
+      - uses: actions/setup-node@v5
+        with: { node-version: '20' }
+
+      - name: Cache emsdk
+        id: cache-emsdk
+        uses: actions/cache@v4
+        with:
+          path: ~/emsdk
+          key: emsdk-3.1.59
+
+      - name: Install emsdk
+        if: steps.cache-emsdk.outputs.cache-hit != 'true'
+        run: |
+          git clone --depth 1 --branch 3.1.59 https://github.com/emscripten-core/emsdk.git ~/emsdk
+          ~/emsdk/emsdk install 3.1.59
+          ~/emsdk/emsdk activate 3.1.59
+
+      - name: Configure
+        run: |
+          source ~/emsdk/emsdk_env.sh
+          emcmake cmake -S . -B build \
+            -DCMAKE_BUILD_TYPE=Release -DBUILD_UNITTESTS=ON \
+            -DBUILD_BENCHMARKS=OFF -DBUILD_EXAMPLES=OFF -DBUILD_EXTRAS=OFF \
+            -DBUILD_STATIC=ON -DBUILD_SHARED=OFF -DWITH_OPENMP=OFF -DNATIVE_SIZE=64 \
+            -DCMAKE_EXE_LINKER_FLAGS="-s INITIAL_MEMORY=536870912 -s ALLOW_MEMORY_GROWTH=1 -s MAXIMUM_MEMORY=4294967296 -s WASM=1 -s NODERAWFS=1 -s STACK_SIZE=5MB -s ASSERTIONS=0"
+
+      - name: Build
+        run: |
+          source ~/emsdk/emsdk_env.sh
+          cmake --build build --target core_tests pke_tests -j2
+
+      - name: core_tests
+        working-directory: build/unittest
+        run: node --max-old-space-size=8192 ./core_tests.js --gtest_color=no
+
+      # Split pke_tests across two processes — each stays well below the 4 GB WASM cap.
+      - name: pke_tests (CKKS + general)
+        working-directory: build/unittest
+        run: node --max-old-space-size=8192 ./pke_tests.js --gtest_filter='UTCKKSRNS*:UTCKKSCacheClear*:UTGENERAL*' --gtest_color=no
+
+      - name: pke_tests (BFV + BGV + binfhe)
+        working-directory: build/unittest
+        run: node --max-old-space-size=8192 ./pke_tests.js --gtest_filter='UTBFV*:UTBGV*:UTBINFHE*' --gtest_color=no
diff --git a/src/core/lib/utils/memory.cpp b/src/core/lib/utils/memory.cpp
@@ -30,7 +30,12 @@
 //==================================================================================
 #include "utils/memory.h"
 
-#if defined(__GLIBC__)
+#if defined(__APPLE__)
+    #include <mach/mach.h>
+    #include <malloc/malloc.h>
+#elif defined(__GLIBC__)
+    #include <malloc.h>
+#elif defined(_MSC_VER)
     #include <malloc.h>
 #endif
 
@@ -67,9 +72,23 @@ namespace {
 // Not async-signal-safe and acquires allocator locks; call only from normal execution context,
 // never from a signal handler. Returns true if a trim was attempted, false on unsupported builds.
 bool AllocTrim() {
-#if defined(__GLIBC__)
+#if defined(__APPLE__)
+    vm_address_t* zones = nullptr;
+    unsigned count      = 0;
+    if (malloc_get_all_zones(mach_task_self(), nullptr, &zones, &count) == KERN_SUCCESS && zones) {
+        for (unsigned i = 0; i < count; ++i) {
+            auto* z = reinterpret_cast<malloc_zone_t*>(zones[i]);
+            if (z)
+                malloc_zone_pressure_relief(z, 0);
+        }
+    }
+    return true;
+#elif defined(__GLIBC__)
     malloc_trim(0);
     return true;
+#elif defined(_MSC_VER)
+    _heapmin();
+    return true;
 #else
     return false;
 #endif
diff --git a/src/pke/include/cryptocontext.h b/src/pke/include/cryptocontext.h
@@ -626,16 +626,24 @@ class CryptoContextImpl : public Serializable {
     * @brief Clear CKKS bootstrap precomputations cached by this context's scheme.
     */
     void ClearBootstrapPrecom() noexcept {
-        VerifyCKKSScheme(__func__);
-        m_scheme->ClearBootstrapPrecom();
+        if (m_scheme)
+            m_scheme->ClearBootstrapPrecom();
     }
 
     /**
     * @brief Clear CKKS/FHEW scheme-switch precomputations cached by this context's scheme.
     */
     void ClearSchemeSwitchPrecom() noexcept {
-        VerifyCKKSScheme(__func__);
-        m_scheme->ClearSchemeSwitchPrecom();
+        if (m_scheme)
+            m_scheme->ClearSchemeSwitchPrecom();
+    }
+
+    /**
+    * @brief Release all scheme-level caches (bootstrap + scheme-switch).
+    */
+    void ClearAllCKKSCaches() noexcept {
+        ClearBootstrapPrecom();
+        ClearSchemeSwitchPrecom();
     }
 
     /**
diff --git a/src/pke/include/cryptocontextfactory.h b/src/pke/include/cryptocontextfactory.h
@@ -64,6 +64,10 @@ class CryptoContextFactory {
 
 public:
     static void ReleaseAllContexts() {
+        for (auto& cc : AllContexts) {
+            if (cc)
+                cc->ClearAllCKKSCaches();
+        }
         if (!AllContexts.empty())
             AllContexts[0]->ClearStaticMapsAndVectors();
         AllContexts.clear();
diff --git a/src/pke/include/schemebase/base-scheme.h b/src/pke/include/schemebase/base-scheme.h
@@ -1147,13 +1147,13 @@ class SchemeBase {
     }
 
     void ClearBootstrapPrecom() noexcept {
-        VerifyFHEEnabled(__func__);
-        m_FHE->ClearBootstrapPrecom();
+        if (m_FHE)
+            m_FHE->ClearBootstrapPrecom();
     }
 
     void ClearSchemeSwitchPrecom() noexcept {
-        VerifySchemeSwitchEnabled(__func__);
-        m_SchemeSwitch->ClearSchemeSwitchPrecom();
+        if (m_SchemeSwitch)
+            m_SchemeSwitch->ClearSchemeSwitchPrecom();
     }
 
     std::shared_ptr<std::map<uint32_t, EvalKey<Element>>> EvalBootstrapKeyGen(const PrivateKey<Element> privateKey,
diff --git a/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp b/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp
@@ -30,14 +30,13 @@
 //==================================================================================
 
 #include "binfhecontext.h"
+#include "gtest/gtest.h"
 #include "openfhe.h"
 #include "scheme/ckksrns/ckksrns-fhe.h"
 #include "scheme/ckksrns/ckksrns-schemeswitching.h"
 #include "scheme/scheme-swch-params.h"
 #include "utils/memory.h"
 
-#include "gtest/gtest.h"
-
 #include <vector>
 
 #if defined(__GLIBC__)
@@ -116,9 +115,13 @@ CryptoContext<DCRTPoly> MakeSchemeSwitchCC() {
 
 class UTCKKSCacheClear : public ::testing::Test {
 protected:
-#if defined(WITH_TCM)
+#if defined(WITH_TCM) || defined(__EMSCRIPTEN__)
     void SetUp() override {
+#if defined(WITH_TCM)
         GTEST_SKIP() << "Heap usage checks are not stable with tcmalloc enabled";
+#else
+        GTEST_SKIP() << "Heap probe unavailable under Emscripten";
+#endif
     }
 #endif
 
diff --git a/src/pke/unittest/utckksrns/UnitTestCKKSrns.cpp b/src/pke/unittest/utckksrns/UnitTestCKKSrns.cpp
@@ -398,7 +398,6 @@ static std::vector<TEST_CASE_UTCKKSRNS> testCases = {
     { EVAL_FAST_ROTATION, "39", {CKKSRNS_SCHEME, RING_DIM, 7,     DFLT,     DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, BV,     FLEXIBLEAUTOEXT, DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,     DFLT,    DFLT},   BATCH},
     { EVAL_FAST_ROTATION, "40", {CKKSRNS_SCHEME, RING_DIM, 7,     DFLT,     DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, HYBRID, FLEXIBLEAUTOEXT, DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,     DFLT,    DFLT},   BATCH},
 #endif
-#if !defined(__EMSCRIPTEN__)
     { EVAL_FAST_ROTATION, "41", {CKKSRNS_SCHEME, RING_DIM, 7,     DFLT,     DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, BV,     FIXEDMANUAL,     DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,     DFLT,    DFLT},   RING_DIM_HALF},
     { EVAL_FAST_ROTATION, "42", {CKKSRNS_SCHEME, RING_DIM, 7,     DFLT,     DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, BV,     FIXEDAUTO,       DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,     DFLT,    DFLT},   RING_DIM_HALF},
     { EVAL_FAST_ROTATION, "43", {CKKSRNS_SCHEME, RING_DIM, 7,     DFLT,     DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, HYBRID, FIXEDMANUAL,     DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,     DFLT,    DFLT},   RING_DIM_HALF},
@@ -527,7 +526,6 @@ static std::vector<TEST_CASE_UTCKKSRNS> testCases = {
     // TestType,              Descr, Scheme,        RDim,   MultDepth, SModSize, DSize, BatchSz, SecKeyDist, MaxRelinSkDeg, FModSize, SecLvl,       KSTech, ScalTech,    LDigits, PtMod, StdDev, EvalAddCt, KSCt, MultTech, EncTech, PREMode
     { SMALL_SCALING_MOD_SIZE, "01", {CKKSRNS_SCHEME, 32768, 19,        22,       DFLT,  DFLT,    DFLT,       DFLT,          23,       DFLT,         DFLT,   FIXEDMANUAL, DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,     DFLT,    DFLT}, },
     { SMALL_SCALING_MOD_SIZE, "02", {CKKSRNS_SCHEME, 32768, 16,        50,       DFLT,  DFLT,    DFLT,       DFLT,          50,       HEStd_NotSet, DFLT,   DFLT,        DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,     DFLT,    DFLT}, },
-#endif
     // ==========================================
     // TestType,   Descr, Scheme,        RDim, MultDepth, SModSize, DSize, BatchSz, SecKeyDist, MaxRelinSkDeg, FModSize, SecLvl,       KSTech, ScalTech,        LDigits, PtMod, StdDev, EvalAddCt, KSCt, MultTech, EncTech, PREMode, MultipartyMode, decryptionNoiseMode, ExecutionMode, NoiseEstimate, RegisterWordSize, compositeDegree, CKKSDataType
     { EVALCOMPLEX, "01", {CKKSRNS_SCHEME, RING_DIM, 7,     DFLT,     DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, BV,     FIXEDMANUAL,     DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,     DFLT,    DFLT, DFLT,  DFLT,   DFLT,      DFLT, DFLT, DFLT, COMPLEX}, },
diff --git a/src/pke/unittest/utckksrns/UnitTestCKKSrnsSerialize.cpp b/src/pke/unittest/utckksrns/UnitTestCKKSrnsSerialize.cpp
@@ -119,7 +119,6 @@ constexpr uint32_t BATCH      = 16;
 // clang-format off
 static std::vector<TEST_CASE_UTCKKSRNS_SER> testCases = {
     // TestType,            Descr, Scheme,         RDim,     MultDepth,  SModSize, DSize, BatchSz, SecKeyDist, MaxRelinSkDeg, FModSize, SecLvl,       KSTech, ScalTech,        LDigits, PtMod, StdDev, EvalAddCt, KSCt, MultTech,  EncTech, PREMode
-#if !defined(__EMSCRIPTEN__)
     { CONTEXT_WITH_SERTYPE, "01", {CKKSRNS_SCHEME, RING_DIM, MULT_DEPTH, SMODSIZE, DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, BV,     FIXEDMANUAL,     DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,      DFLT,    DFLT}, },
     { CONTEXT_WITH_SERTYPE, "02", {CKKSRNS_SCHEME, RING_DIM, MULT_DEPTH, SMODSIZE, DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, BV,     FIXEDAUTO,       DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,      DFLT,    DFLT}, },
     { CONTEXT_WITH_SERTYPE, "03", {CKKSRNS_SCHEME, RING_DIM, MULT_DEPTH, SMODSIZE, DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, HYBRID, FIXEDMANUAL,     DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,      DFLT,    DFLT}, },
@@ -129,7 +128,6 @@ static std::vector<TEST_CASE_UTCKKSRNS_SER> testCases = {
     { CONTEXT_WITH_SERTYPE, "06", {CKKSRNS_SCHEME, RING_DIM, MULT_DEPTH, SMODSIZE, DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, HYBRID, FLEXIBLEAUTO,    DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,      DFLT,    DFLT}, },
     { CONTEXT_WITH_SERTYPE, "07", {CKKSRNS_SCHEME, RING_DIM, MULT_DEPTH, SMODSIZE, DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, BV,     FLEXIBLEAUTOEXT, DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,      DFLT,    DFLT}, },
     { CONTEXT_WITH_SERTYPE, "08", {CKKSRNS_SCHEME, RING_DIM, MULT_DEPTH, SMODSIZE, DSIZE, BATCH,   DFLT,       DFLT,          DFLT,     HEStd_NotSet, HYBRID, FLEXIBLEAUTOEXT, DFLT,    DFLT,  DFLT,   DFLT,      DFLT, DFLT,      DFLT,    DFLT}, },
-#endif
 #endif
     // ==========================================
     // TestType,            Descr, Scheme,         RDim,     MultDepth,  SModSize, DSize, BatchSz, SecKeyDist, MaxRelinSkDeg, FModSize, SecLvl,       KSTech, ScalTech,        LDigits, PtMod, StdDev, EvalAddCt, KSCt, MultTech,  EncTech, PREMode
@@ -189,6 +187,7 @@ class UTCKKSRNS_SER : public ::testing::TestWithParam<TEST_CASE_UTCKKSRNS_SER> {
         CryptoContext<Element> cc(UnitTestGenerateContext(testData.params));
 
         UnitTestContextWithSertype(cc, SerType::JSON, "json");
+        cc->ClearAllCKKSCaches();
         UnitTestContextWithSertype(cc, SerType::BINARY, "binary");
     }
 

Original file line number	Diff line number	Diff line change
`@@ -1147,13 +1147,13 @@ class SchemeBase {`
`1147`	`1147`	`}`
`1148`	`1148`
`1149`	`1149`	`void ClearBootstrapPrecom() noexcept {`
`1150`		`- VerifyFHEEnabled(__func__);`
`1151`		`- m_FHE->ClearBootstrapPrecom();`
	`1150`	`+ if (m_FHE)`
	`1151`	`+ m_FHE->ClearBootstrapPrecom();`
`1152`	`1152`	`}`
`1153`	`1153`
`1154`	`1154`	`void ClearSchemeSwitchPrecom() noexcept {`
`1155`		`- VerifySchemeSwitchEnabled(__func__);`
`1156`		`- m_SchemeSwitch->ClearSchemeSwitchPrecom();`
	`1155`	`+ if (m_SchemeSwitch)`
	`1156`	`+ m_SchemeSwitch->ClearSchemeSwitchPrecom();`
`1157`	`1157`	`}`
`1158`	`1158`
`1159`	`1159`	`std::shared_ptr<std::map<uint32_t, EvalKey<Element>>> EvalBootstrapKeyGen(const PrivateKey<Element> privateKey,`