Skip to content

Commit 43a4c82

Browse files
author
GitHub Actions
committed
Add in-memory heap storage implementation
- Implemented InMemoryHeapStorage using HashMap with Arc<RwLock> - Added --in-memory CLI flag to use in-memory storage - Provides thread-safe, zero-config storage for development and testing - Added comprehensive unit tests for in-memory storage - Added IN_MEMORY_STORAGE.md documentation - Updated main.rs to support in-memory storage option - InMemory variant added to AnyHeapStorage enum Benefits: - Fastest performance (no I/O overhead) - No external dependencies - Useful for testing, development, and ephemeral workloads - Rich API with utility methods (len, keys, clear, contains_key, remove) Limitations: - Data lost on process termination - Not suitable for production with persistent data needs - Memory bounded by available RAM
1 parent 1423567 commit 43a4c82

File tree

3 files changed

+333
-7
lines changed

3 files changed

+333
-7
lines changed

IN_MEMORY_STORAGE.md

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
# In-Memory Heap Storage
2+
3+
## Overview
4+
5+
This repository now includes an in-memory heap storage implementation that stores V8 heap snapshots in RAM rather than persisting them to disk or cloud storage. This is useful for:
6+
7+
- **Testing and Development**: Fast iteration without filesystem I/O overhead
8+
- **Temporary Sessions**: Sessions that don't require persistence beyond process lifetime
9+
- **Performance**: Faster read/write operations compared to disk or S3
10+
- **Simplicity**: No external dependencies or configuration required
11+
12+
## Implementation Details
13+
14+
### Architecture
15+
16+
The `InMemoryHeapStorage` struct uses:
17+
- `HashMap<String, Vec<u8>>` to store heap snapshots keyed by heap name
18+
- `Arc<RwLock<...>>` for thread-safe concurrent access
19+
- Tokio's async RwLock for non-blocking operations
20+
21+
### Features
22+
23+
- **Thread-safe**: Multiple concurrent readers with exclusive writer access
24+
- **Clone-able**: Shares the same underlying storage across clones
25+
- **Zero-configuration**: Works out of the box with no setup
26+
- **Rich API**: Additional utility methods beyond the basic `put`/`get` interface
27+
28+
### API
29+
30+
Basic HeapStorage trait methods:
31+
- `put(name: &str, data: &[u8]) -> Result<(), String>` - Store heap snapshot
32+
- `get(name: &str) -> Result<Vec<u8>, String>` - Retrieve heap snapshot
33+
34+
Additional utility methods:
35+
- `len() -> usize` - Get number of stored items
36+
- `is_empty() -> bool` - Check if storage is empty
37+
- `clear()` - Remove all stored data
38+
- `contains_key(name: &str) -> bool` - Check if key exists
39+
- `remove(name: &str) -> Option<Vec<u8>>` - Remove and return specific item
40+
- `keys() -> Vec<String>` - Get all stored keys
41+
42+
## Usage
43+
44+
### Command Line
45+
46+
Run the server with in-memory storage:
47+
48+
```bash
49+
# With stdio transport (default)
50+
cargo run --bin server -- --in-memory
51+
52+
# With HTTP transport
53+
cargo run --bin server -- --in-memory --http-port 8080
54+
55+
# With SSE transport
56+
cargo run --bin server -- --in-memory --sse-port 8080
57+
```
58+
59+
### Programmatic Usage
60+
61+
```rust
62+
use mcp::heap_storage::{InMemoryHeapStorage, HeapStorage};
63+
64+
// Create new in-memory storage
65+
let storage = InMemoryHeapStorage::new();
66+
67+
// Store data
68+
storage.put("heap1", b"snapshot data").await.unwrap();
69+
70+
// Retrieve data
71+
let data = storage.get("heap1").await.unwrap();
72+
73+
// Check existence
74+
if storage.contains_key("heap1").await {
75+
println!("Found heap1");
76+
}
77+
78+
// Get all keys
79+
let keys = storage.keys().await;
80+
println!("Stored heaps: {:?}", keys);
81+
82+
// Clear all data
83+
storage.clear().await;
84+
```
85+
86+
### Integration with AnyHeapStorage
87+
88+
```rust
89+
use mcp::heap_storage::{AnyHeapStorage, InMemoryHeapStorage};
90+
91+
let storage = AnyHeapStorage::InMemory(InMemoryHeapStorage::new());
92+
93+
// Use through the HeapStorage trait
94+
storage.put("key", b"value").await.unwrap();
95+
let value = storage.get("key").await.unwrap();
96+
```
97+
98+
## Comparison with Other Storage Backends
99+
100+
| Feature | In-Memory | File | S3 |
101+
|---------|-----------|------|-----|
102+
| Persistence | Process lifetime only | Permanent | Permanent |
103+
| Performance | Fastest | Fast | Slower (network) |
104+
| Setup Required | None | Directory path | AWS credentials + bucket |
105+
| Memory Usage | High (all in RAM) | Low | Low |
106+
| Distributed Access | No | Shared filesystem only | Yes |
107+
| Best For | Testing, dev, temp sessions | Single server production | Multi-server production |
108+
109+
## Configuration Options
110+
111+
The in-memory storage conflicts with other storage options:
112+
- Cannot be used with `--s3-bucket`
113+
- Cannot be used with `--directory-path`
114+
- Cannot be used with `--stateless`
115+
116+
## Limitations
117+
118+
1. **No Persistence**: All data is lost when the process terminates
119+
2. **Memory Bounded**: Limited by available RAM
120+
3. **Single Process**: Cannot be shared across multiple server instances
121+
4. **No Backup**: No automatic snapshot or recovery mechanism
122+
123+
## Testing
124+
125+
The implementation includes comprehensive unit tests:
126+
127+
```bash
128+
# Run all tests
129+
cargo test -p server
130+
131+
# Run only heap_storage tests
132+
cargo test -p server heap_storage
133+
```
134+
135+
Test coverage includes:
136+
- Basic put/get operations
137+
- Error handling for missing keys
138+
- Utility method functionality
139+
- Clone behavior (shared storage)
140+
- Integration with AnyHeapStorage enum
141+
142+
## When to Use
143+
144+
**Use in-memory storage when:**
145+
- Developing and testing locally
146+
- Running ephemeral workloads
147+
- Performance is critical and persistence is not
148+
- You want zero configuration overhead
149+
150+
**Don't use in-memory storage when:**
151+
- You need data to survive process restarts
152+
- Running in production with important data
153+
- Memory is constrained
154+
- Multiple servers need to share state
155+
156+
## Future Enhancements
157+
158+
Possible improvements:
159+
- Memory usage limits with LRU eviction
160+
- Periodic snapshots to disk
161+
- Memory metrics and monitoring
162+
- TTL (time-to-live) for entries

server/src/main.rs

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,23 +13,27 @@ use tokio_util::sync::CancellationToken;
1313

1414
mod mcp;
1515
use mcp::{StatelessService, StatefulService, initialize_v8};
16-
use mcp::heap_storage::{AnyHeapStorage, S3HeapStorage, FileHeapStorage};
16+
use mcp::heap_storage::{AnyHeapStorage, S3HeapStorage, FileHeapStorage, InMemoryHeapStorage};
1717

1818
/// Command line arguments for configuring heap storage
1919
#[derive(Parser, Debug)]
2020
#[command(author, version, about, long_about = None)]
2121
struct Cli {
2222

2323
/// S3 bucket name (required if --use-s3)
24-
#[arg(long, conflicts_with_all = ["directory_path", "stateless"])]
24+
#[arg(long, conflicts_with_all = ["directory_path", "stateless", "in_memory"])]
2525
s3_bucket: Option<String>,
2626

2727
/// Directory path for filesystem storage (required if --use-filesystem)
28-
#[arg(long, conflicts_with_all = ["s3_bucket", "stateless"])]
28+
#[arg(long, conflicts_with_all = ["s3_bucket", "stateless", "in_memory"])]
2929
directory_path: Option<String>,
3030

31+
/// Use in-memory storage (heap state persists only during process lifetime)
32+
#[arg(long, conflicts_with_all = ["s3_bucket", "directory_path", "stateless"])]
33+
in_memory: bool,
34+
3135
/// Run in stateless mode - no heap snapshots are saved or loaded
32-
#[arg(long, conflicts_with_all = ["s3_bucket", "directory_path"])]
36+
#[arg(long, conflicts_with_all = ["s3_bucket", "directory_path", "in_memory"])]
3337
stateless: bool,
3438

3539
/// HTTP port to listen on (if not specified, uses stdio transport)
@@ -77,12 +81,18 @@ async fn main() -> Result<()> {
7781
}
7882
} else {
7983
// Stateful mode - with heap persistence
80-
let heap_storage = if let Some(bucket) = cli.s3_bucket {
84+
let heap_storage = if cli.in_memory {
85+
tracing::info!("Using in-memory heap storage");
86+
AnyHeapStorage::InMemory(InMemoryHeapStorage::new())
87+
} else if let Some(bucket) = cli.s3_bucket {
88+
tracing::info!("Using S3 heap storage with bucket: {}", bucket);
8189
AnyHeapStorage::S3(S3HeapStorage::new(bucket).await)
8290
} else if let Some(dir) = cli.directory_path {
91+
tracing::info!("Using file heap storage with directory: {}", dir);
8392
AnyHeapStorage::File(FileHeapStorage::new(dir))
8493
} else {
8594
// default to file /tmp/mcp-v8-heaps
95+
tracing::info!("Using default file heap storage: /tmp/mcp-v8-heaps");
8696
AnyHeapStorage::File(FileHeapStorage::new("/tmp/mcp-v8-heaps"))
8797
};
8898

0 commit comments

Comments
 (0)