Skip to content

[Client Encryption] Add ArrayPool-backed pooled streams for reduced allocations in streaming encrypt/decrypt operations #5479

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
adamnova wants to merge 11 commits into
base: master
Choose a base branch
from
Open

[Client Encryption] Add ArrayPool-backed pooled streams for reduced allocations in streaming encrypt/decrypt operations #5479

adamnova wants to merge 11 commits into from

Conversation

@adamnova
Copy link
Contributor

@adamnova adamnova commented Nov 7, 2025 •
edited
Loading

This PR introduces memory-efficient pooled stream implementations for the Client Encryption Custom library to significantly reduce GC pressure during streaming encryption/decryption operations.

Key Changes:

  1. New PooledMemoryStream: A MemoryStream replacement that rents buffers from ArrayPool<byte>.Shared and properly returns them on disposal. Includes:

    • Automatic buffer clearing for security (sensitive encryption data never remains in pooled memory)
    • Support for both sync and async operations
    • Span-based read/write APIs on .NET 8+
    • Proper overflow protection for buffer capacity growth

  2. New PooledJsonSerializer (.NET 8+): Optimized System.Text.Json serialization utilities using pooled buffers:

    • SerializeToPooledStream<T>() - Returns a PooledMemoryStream that caller must dispose
    • DeserializeFromStream/Span with MaxDepth=64 for DoS protection
    • Thread-safe with independent stream instances per invocation

  3. New PooledStreamConfiguration (.NET 8+): Configurable initial capacities for pooled streams/buffers

  4. StreamProcessor improvements:

    • Added try-finally blocks to ensure encryptionPayloadWriter and bufferWriter disposal on exception paths
    • Fixed potential memory leaks when decryption context is null (lines 260-263, 144-148)
    • Replaced MemoryStream with PooledMemoryStream in encrypt/decrypt hot paths

  5. SystemTextJsonStreamAdapter fixes: Added disposal of PooledMemoryStream when decryption returns null context

Benchmark Results

Allocation reduction vs Newtonsoft baseline:

Document Size Encrypt Reduction Decrypt Reduction
~1 KB ~82% ~82%
~4 KB ~76% ~72%
~8 KB ~83% ~85%

Type of change

  • New feature (non-breaking change which adds functionality)

Testing

  • Added comprehensive unit tests for PooledMemoryStream (read/write, seek, disposal, capacity growth)
  • Added PooledJsonSerializer tests (serialization round-trips, max depth limits, null handling)
  • Added security-focused tests (buffer clearing, disposal patterns, sensitive data handling)
  • Added memory leak stress tests (1000+ iterations verifying no ArrayPool exhaustion)
  • Added exception-path resource disposal tests

#4678

@adamnova
Add ArrayPool-based stream pooling for System.Text.Json encryption/de…
ef6be1f 
…cryption

- Add PooledMemoryStream: ArrayPool-backed MemoryStream replacement
- Add PooledJsonSerializer: Optimized JSON serialization utilities
- Add PooledStreamConfiguration: Centralized pooling configuration
- Update SystemTextJsonStreamAdapter to use pooled streams
- Update EncryptionProcessor to use pooled streams
- Update MdeEncryptionProcessor to use pooled streams

Benefits:
- ~50% reduction in memory allocations for encrypt/decrypt operations
- Reduced GC pressure through buffer reuse
- Full Stream API compatibility
- Secure buffer clearing on disposal (configurable)
- Multi-target support (net8.0 + netstandard2.0)
@adamnova
Update performance benchmark documentation
46ba6de 
Replaces raw BenchmarkDotNet output with a summarized analysis of memory allocation efficiency for stream-based encryption/decryption versus the Newtonsoft baseline. Adds tables, key observations, and instructions for running benchmarks, focusing on allocation metrics rather than execution time.
@adamnova adamnova closed this Nov 12, 2025
@adamnova adamnova reopened this Dec 17, 2025
@adamnova adamnova changed the title Adds pooled stream to StreamProcessor to reduce allocations [Client Encryption] Adds pooled stream to StreamProcessor to reduce allocations Dec 17, 2025
@adamnova
Add remarks on thread safety and disposal; improve stream handling
3957750 
Added detailed XML remarks to PooledJsonSerializer, PooledMemoryStream, and PooledStreamConfiguration regarding thread safety, disposal, and performance/security considerations. Improved stream disposal and buffer management in StreamProcessor.Encryptor and SystemTextJsonStreamAdapter to ensure proper resource cleanup and prevent memory leaks. Fixed integer overflow checks in PooledMemoryStream for large write operations.
@adamnova
Add comprehensive unit tests for pooled resource classes
af5a12e 
- PooledMemoryStreamTests: 586 lines, comprehensive stream operation tests
- PooledJsonSerializerTests: 388 lines, serialization/deserialization tests
- PooledJsonSerializerSecurityTests: Security-focused tests for depth limits and disposal

Tests cover core functionality, error paths, async operations, disposal, and security hardening.
Additional edge case and integration tests will follow in subsequent commits.
@adamnova
Add critical integration test for PooledMemoryStream disposal
e4ee4b3 
Tests SystemTextJsonStreamAdapter.cs:52 disposal path to prevent memory leaks
when decryption returns null context.

- Enhanced existing test with detailed comments and assertions
- Added stress test with 1000 iterations to verify no memory leak
- Both tests pass successfully
@adamnova
Fix critical memory leaks and add comprehensive integration tests
78b074b 
FIXES:
- EncryptionProcessor.cs:263: Dispose PooledMemoryStream when context is null
- MdeEncryptionProcessor.cs:146: Dispose PooledMemoryStream when context is null

These fixes prevent ArrayPool exhaustion when decrypting non-encrypted payloads.

TESTS ADDED:
EncryptionProcessorTests.cs:
- DecryptStreamAsync_UsesPooledMemoryStream_DisposesCorrectly
- DecryptStreamAsync_WithNoEncryptionProperties_DisposesPooledMemoryStream
- DecryptStreamAsync_UsesPooledJsonSerializer_DeserializesCorrectly
- DecryptStreamAsync_RepeatedCalls_NoMemoryLeak (100 iterations stress test)

MdeEncryptionProcessorTests.cs:
- DecryptStreamAsync_CreatesPooledMemoryStream_SuccessPath
- DecryptStreamAsync_WithEmptyEncryptedPaths_HandlesCorrectly
- DecryptStreamAsync_RepeatedCalls_NoMemoryLeak (100 iterations stress test)

StreamProcessorEncryptorTests.cs:
- EncryptStreamAsync_TryFinallyBlock_DisposesResourcesOnSuccess
- EncryptStreamAsync_ExceptionDuringParsing_DisposesResourcesInFinally
- EncryptStreamAsync_CancellationRequested_DisposesResourcesInFinally
- EncryptStreamAsync_RepeatedCallsWithExceptions_NoResourceLeak (100 iterations)

All 11 new integration tests pass, verifying no memory leaks under stress.
@adamnova
Add specialized tests for pooled resource management
76e8aba 
NEW TEST FILES:
- PooledResourceMemoryLeakTests.cs (6 tests)
  * Documents memory leak dangers and proper disposal patterns
  * Stress tests with 1000+ iterations to verify no ArrayPool exhaustion
  * Tests concurrent usage patterns
  * Tests rapid capacity growth scenarios

- PooledResourceDisposalTests.cs (10 tests)
  * Verifies proper buffer return to ArrayPool
  * Tests clearOnReturn security feature (buffer zeroing)
  * Tests disposal edge cases (double dispose, zero capacity, empty stream)
  * Tests async disposal path
  * Verifies try-finally disposal in exception paths

RESULTS:
- All 306 tests pass (290 original + 16 new)
- Comprehensive coverage of pooled resource lifecycle
- Memory leak prevention patterns documented
- Security features (buffer clearing) verified
@adamnova
Dispose bufferWriter after use in StreamProcessor.Encryptor
3147c28 
Added calls to bufferWriter?.Dispose() before setting bufferWriter to null to ensure proper resource cleanup in the encryption process.
@adamnova adamnova marked this pull request as ready for review December 18, 2025 12:01
@adamnova adamnova changed the title [Client Encryption] Adds pooled stream to StreamProcessor to reduce allocations [Client Encryption] Add ArrayPool-backed pooled streams for reduced allocations in streaming encrypt/decrypt operations Dec 18, 2025
adamnova and others added 2 commits December 18, 2025 13:58
@adamnova @claude
Fix critical memory leaks in stream disposal paths and security gaps
337a094 
This commit addresses several critical issues in the pooled stream management:

1. Added back Debug.Assert(memoryStream.TryGetBuffer(out _)) in AeAesEncryptionProcessor
   - The assert was incorrectly removed in e4ee4b3
   - TryGetBuffer() is internal and accessible within the same assembly
   - The assert validates the stream implementation before using ToArray()

2. Fixed PooledMemoryStream security gaps to prevent data leakage:
   - SetLength: Zero newly exposed regions when expanding length
   - Write: Zero gaps between current length and write position
   - Dispose: Clear entire capacity (not just length) to handle SetLength(0)
   - EnsureCapacity: Clear new buffer before use and old buffer before return

3. Fixed PooledJsonSerializer exception safety:
   - Added try-catch in SerializeToPooledArray to return rented buffer on exception
   - Enhanced documentation to clarify caller responsibility for buffer return

4. Added exception handling in all stream disposal paths:
   - SystemTextJsonStreamAdapter.EncryptAsync and DecryptAsync
   - MdeEncryptionProcessor.DecryptAsync
   - EncryptionProcessor.DecryptAsync

All changes ensure that PooledMemoryStream instances are properly disposed
even in error paths, preventing ArrayPool buffer leaks.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
@adamnova @claude
Make PooledStreamConfiguration properties thread-safe with immutable …
ce5b8c1 
…configuration pattern

Refactored PooledStreamConfiguration to use an immutable configuration object pattern
for thread-safe access to configuration values. This eliminates race conditions from
mutable static properties.

Changes:
- Convert PooledStreamConfiguration from static class to sealed class with immutable properties
- Add Current property for thread-safe configuration reads
- Add SetConfiguration() method for atomic configuration swaps
- Replace mutable StreamInitialCapacity/BufferWriterInitialCapacity with init-only properties
- Consolidate StreamProcessor.InitialBufferSize into PooledStreamConfiguration.StreamProcessorBufferSize
- Update all usage sites to read from PooledStreamConfiguration.Current
- Update test files to use SetConfiguration() pattern for test-specific configurations

Thread-safety benefits:
- Atomic reference swaps ensure readers always see consistent configuration
- Immutable configuration objects prevent partial updates
- No torn reads or memory visibility issues across threads
- Tests can safely swap and restore configurations

All 306 tests pass (net8.0) and 166 tests pass (net6.0).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kirankumarkolli kirankumarkolli Awaiting requested review from kirankumarkolli kirankumarkolli is a code owner

@FabianMeiswinkel FabianMeiswinkel Awaiting requested review from FabianMeiswinkel FabianMeiswinkel is a code owner

@sboshra sboshra Awaiting requested review from sboshra sboshra is a code owner

@Pilchie Pilchie Awaiting requested review from Pilchie Pilchie is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@adamnova