perf(storage): avoid full-size copies of FHE keys on the keygen store… - #631
Merged
Conversation
Consolidated Tests Results 2026-06-04 - 14:03:37Test Results
Details
test-reporter: Run #2562
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 🔄 This comment has been updated |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR reduces peak memory usage when persisting large TFHE/FHE key material by avoiding unnecessary deep clones/copies during storage and key-material wiring.
Changes:
- Avoids cloning the full serialized blob when uploading versioned data to S3.
- Streams
safe_serializeoutput directly to disk via a temp file + atomic rename to avoid buffering multi‑GiB blobs in memory. - Shares a single
Arc<CompressedXofKeySet>allocation across private/public key material and avoids deep-cloning before decompression.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| core/service/src/vault/storage/s3.rs | Removes a full Vec<u8> clone before S3 upload to reduce peak RSS. |
| core/service/src/util/file_handling.rs | Switches versioned writes to stream to a temp file and atomically rename into place (no in-memory Vec<u8>). |
| core/service/src/engine/threshold/service/kms_impl.rs | Adds an Arc-accepting constructor and removes a deep clone before decompression. |
| core/service/src/engine/threshold/service/key_generator.rs | Wraps the compressed keyset in a single Arc and reuses it for both stored public data and private material. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
maksymsur
force-pushed
the
maksymsur/optim/2834/mem-consumption
branch
from
June 3, 2026 15:36
3d24d4f to
341f241
Compare
maksymsur
force-pushed
the
maksymsur/optim/2834/mem-consumption
branch
from
June 4, 2026 13:09
2a5bfd3 to
5eba3ce
Compare
maksymsur
force-pushed
the
maksymsur/optim/2834/mem-consumption
branch
from
June 5, 2026 08:47
f3d753c to
1653548
Compare
Consolidated Tests Results 2026-06-05 - 08:48:23Test Results
Details
test-reporter: Run #5134
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-05 - 08:49:19Test Results
Details
test-reporter: Run #5134
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-05 - 09:31:03Test Results
Details
test-reporter: Run #2586
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-10 - 22:57:38Test Results
Details
test-reporter: Run #2637
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
maksymsur
force-pushed
the
maksymsur/optim/2834/mem-consumption
branch
from
June 11, 2026 04:43
1111aa7 to
f3a709d
Compare
- expand_keys: `CompressedXofKeySet::decompress` takes `&self` (tfhe 1.6.1), so decompress through the `Arc` directly instead of first cloning the whole (multi-GiB) compressed keyset. Removes a per-key ~GiB clone on the first decryption after a key is loaded; drops the now-stale tfhe-rs #3469 TODO. - file_handling: name the streaming temp file dot-prefixed (`.<name>.partial`) so a crash-leftover is skipped by directory listings instead of breaking `FileStorage::all_data_ids`, which parses every non-hidden filename as a RequestId.
safe_write_element_versioned used a fixed `.<name>.partial` temp name and relied on the storage-layer mutex to avoid collisions. That mutex is per-instance, so it guards neither separate processes, nor separate storage instances over the same directory, nor the non-storage callers (kms-custodian, core-client). Two concurrent writers to the same path could then create, truncate and interleave the same temp file, yielding a corrupted final file or a spurious rename failure. Add a per-writer suffix (process id + a process-local atomic counter) to the temp name, keeping the dot-prefix so directory listings still skip leftovers. Add unit tests for round-trip, no leftover temp file after a successful write, concurrent writes to the same path, and destination preservation when the write fails before the rename.
destroy_mpc_epoch deleted only the on-disk key material, leaving the (tens-of-GiB) decompressed keys resident in ThresholdCryptoMaterialStorage.fhe_keys until process restart. Also drop the cache entries for the destroyed epoch. This is pure reclaim with no decryption-behaviour change: decrypts under a destroyed epoch already fail epoch validation, so the freed entries are unreachable. Add a fhe_key_cache_size gauge, recorded in the existing periodic metrics loops for both the threshold and centralized variants, so in-memory key-cache growth is observable in production.
safe_write_element_versioned writes to a uniquely-named dot-prefixed sibling temp file and renames it into place. On a failed write (create/serialize/flush/rename) the partial was left behind; the per-writer unique suffix means these accumulate across repeated failures and waste disk space. Add a TempFileGuard that removes the temp file on any early return and is disarmed only after a successful rename, so the rename-failure path is covered too.
- fsync the temp file (and best-effort its parent directory) before the atomic rename in safe_write_element_versioned, so the rename publishes durable bytes even across a power loss; add path context to the serialize and rename error paths - resharing: decompress the compressed keyset in place (&self) and move it into PublicKeyMaterial at its last use instead of deep-cloning the multi-GiB keyset twice - add missing sunshine tests for the new pub items: PublicKeyMaterial::from_arc, CoreMetrics::record_fhe_key_cache_size (including the i64 clamp) and the centralized cached_fhe_key_count - fix comment accuracy (tfhe version reference, serialized blob sizes are bounded by SAFE_SER_SIZE_LIMIT)
maksymsur
force-pushed
the
maksymsur/optim/2834/mem-consumption
branch
from
June 11, 2026 04:44
f3a709d to
21b29bd
Compare
Consolidated Tests Results 2026-06-11 - 05:30:51Test Results
Details
test-reporter: Run #2639
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Keep only the constraints the code cannot show (temp-name uniqueness, lock ordering, why serialization runs inline); drop change-narration and restore the missing rustdoc on record_fhe_key_cache_size.
Contributor
Author
… sharing - replace the hand-rolled temp-file naming and cleanup guard with tempfile::NamedTempFile (unique O_EXCL name, auto-delete on drop); keep the buffered writer, the pre-rename fsync and the best-effort parent-dir fsync - let PublicKeyMaterial own the keyset and hand out shared refs via compressed_keyset() instead of the from_arc constructor; the keygen store path builds the material first and reuses its allocation for the public-storage set - log the destroy-epoch cache purge unconditionally (kept at info: it is the operational signal that multi-GiB key material was reclaimed)
Consolidated Tests Results 2026-06-12 - 11:21:34Test Results
Details
test-reporter: Run #5226
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-12 - 11:22:32Test Results
Details
test-reporter: Run #5226
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-12 - 12:03:07Test Results
Details
test-reporter: Run #2678
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
dvdplm
reviewed
Jun 12, 2026
dvdplm
reviewed
Jun 12, 2026
dvdplm
reviewed
Jun 12, 2026
Consolidated Tests Results 2026-06-12 - 14:23:13Test Results
Details
test-reporter: Run #5240
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-12 - 14:23:18Test Results
Details
test-reporter: Run #5240
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-12 - 14:56:37Test Results
Details
test-reporter: Run #5246
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-12 - 14:57:52Test Results
Details
test-reporter: Run #5246
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Keep the file sync_all before the rename (prevents a truncated file at the destination); without the directory fsync the worst case on power loss is that the write looks like it never happened.
Consolidated Tests Results 2026-06-12 - 15:35:49Test Results
Details
test-reporter: Run #2698
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
The gauge keeps its previous value when the cache lock is contended instead of stalling the whole system-metrics loop behind long-held decryption read guards and queued writers.
Consolidated Tests Results 2026-06-12 - 15:48:17Test Results
Details
test-reporter: Run #5252
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-12 - 15:49:24Test Results
Details
test-reporter: Run #5252
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
Consolidated Tests Results 2026-06-12 - 16:24:57Test Results
Details
test-reporter: Run #2703
🎉 All tests passed!TestsView All Tests
🍂 No flaky tests in this run. Github Test Reporter by CTRF 💚 |
dvdplm
approved these changes
Jun 12, 2026
dvdplm
left a comment
dvdplm
left a comment
Contributor
There was a problem hiding this comment.
LGTM, thank you for your patience! :)
This was referenced Jul 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description of changes
Reduces peak memory usage when handling multi-GiB FHE key material:
Arc<CompressedXofKeySet>between private and public key material instead of deep-cloning the keyset&self) instead of cloning them firstfhe_key_cache_sizegaugeIssue ticket number and link
https://github.com/zama-ai/kms-internal/issues/2834
PR Checklist
I attest that all checked items are satisfied. Any deviation is clearly justified above.
chore: ...).TODO(#issue).unwrap/expect/paniconly in tests or for invariant bugs (documented if present).devopslabel + infra notified + infra-team reviewer assigned.!and affected teams notified.Zeroize+ZeroizeOnDropimplemented.unsafe; if unavoidable: minimal, justified, documented, and test/fuzz covered.Dependency Update Questionnaire (only if deps changed or added)
Answer in the
Cargo.tomlnext to the dependency (or here if updating):More details and explanations for the checklist and dependency updates can be found in CONTRIBUTING.md