feoxdb/core/store/atomic.rs
1use bytes::Bytes;
2use std::sync::atomic::Ordering;
3use std::sync::Arc;
4
5use crate::constants::Operation;
6use crate::core::record::Record;
7use crate::error::{FeoxError, Result};
8
9use super::FeoxStore;
10
11impl FeoxStore {
12 /// Atomically increment a numeric counter.
13 ///
14 /// The value must be stored as an 8-byte little-endian i64. If the key doesn't exist,
15 /// it will be created with the given delta value. If it exists, the value will be
16 /// incremented atomically.
17 ///
18 /// # Value Format
19 ///
20 /// The value MUST be exactly 8 bytes representing a little-endian i64.
21 /// Use `i64::to_le_bytes()` to create the initial value:
22 /// ```rust,ignore
23 /// let zero: i64 = 0;
24 /// store.insert(b"counter", &zero.to_le_bytes())?;
25 /// ```
26 ///
27 /// # Arguments
28 ///
29 /// * `key` - The key of the counter
30 /// * `delta` - The amount to increment by (can be negative for decrement)
31 /// * `timestamp` - Optional timestamp for conflict resolution
32 ///
33 /// # Returns
34 ///
35 /// Returns the new value after incrementing.
36 ///
37 /// # Errors
38 ///
39 /// * `InvalidOperation` - Existing value is not exactly 8 bytes (not a valid i64)
40 /// * `OlderTimestamp` - Timestamp is not newer than existing record
41 ///
42 /// # Example
43 ///
44 /// ```rust
45 /// # use feoxdb::FeoxStore;
46 /// # fn main() -> feoxdb::Result<()> {
47 /// # let store = FeoxStore::new(None)?;
48 /// // Initialize counter with proper binary format
49 /// let initial: i64 = 0;
50 /// store.insert(b"visits", &initial.to_le_bytes())?;
51 ///
52 /// // Increment atomically
53 /// let val = store.atomic_increment(b"visits", 1)?;
54 /// assert_eq!(val, 1);
55 ///
56 /// // Increment by 5
57 /// let val = store.atomic_increment(b"visits", 5)?;
58 /// assert_eq!(val, 6);
59 ///
60 /// // Decrement by 2
61 /// let val = store.atomic_increment(b"visits", -2)?;
62 /// assert_eq!(val, 4);
63 ///
64 /// // Or create new counter directly (starts at delta value)
65 /// let downloads = store.atomic_increment(b"downloads", 100)?;
66 /// assert_eq!(downloads, 100);
67 /// # Ok(())
68 /// # }
69 /// ```
70 pub fn atomic_increment(&self, key: &[u8], delta: i64) -> Result<i64> {
71 self.atomic_increment_with_timestamp_and_ttl(key, delta, None, 0)
72 }
73
74 /// Atomically increment/decrement with explicit timestamp.
75 ///
76 /// This is the advanced version that allows manual timestamp control.
77 /// Most users should use `atomic_increment()` instead.
78 ///
79 /// # Arguments
80 ///
81 /// * `key` - The key to increment/decrement
82 /// * `delta` - Amount to add (negative to decrement)
83 /// * `timestamp` - Optional timestamp. If `None`, uses current time.
84 ///
85 /// # Errors
86 ///
87 /// * `OlderTimestamp` - Timestamp is not newer than existing record
88 pub fn atomic_increment_with_timestamp(
89 &self,
90 key: &[u8],
91 delta: i64,
92 timestamp: Option<u64>,
93 ) -> Result<i64> {
94 self.atomic_increment_with_timestamp_and_ttl(key, delta, timestamp, 0)
95 }
96
97 /// Atomically increment/decrement with TTL support.
98 ///
99 /// # Arguments
100 ///
101 /// * `key` - The key to increment/decrement
102 /// * `delta` - Amount to add (negative to decrement)
103 /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
104 ///
105 /// # Errors
106 ///
107 /// * `InvalidOperation` - Value is not a valid i64
108 pub fn atomic_increment_with_ttl(
109 &self,
110 key: &[u8],
111 delta: i64,
112 ttl_seconds: u64,
113 ) -> Result<i64> {
114 self.atomic_increment_with_timestamp_and_ttl(key, delta, None, ttl_seconds)
115 }
116
117 /// Atomically increment/decrement with explicit timestamp and TTL.
118 ///
119 /// # Arguments
120 ///
121 /// * `key` - The key to increment/decrement
122 /// * `delta` - Amount to add (negative to decrement)
123 /// * `timestamp` - Optional timestamp. If `None`, uses current time.
124 /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
125 ///
126 /// # Errors
127 ///
128 /// * `OlderTimestamp` - Timestamp is not newer than existing record
129 pub fn atomic_increment_with_timestamp_and_ttl(
130 &self,
131 key: &[u8],
132 delta: i64,
133 timestamp: Option<u64>,
134 ttl_seconds: u64,
135 ) -> Result<i64> {
136 self.validate_key(key)?;
137
138 let key_vec = key.to_vec();
139
140 let result = match self.hash_table.entry(key_vec.clone()) {
141 scc::hash_map::Entry::Occupied(mut entry) => {
142 let old_record = entry.get();
143
144 // Get timestamp inside the critical section to ensure it's always newer
145 let timestamp = match timestamp {
146 Some(0) | None => self.get_timestamp(),
147 Some(ts) => ts,
148 };
149
150 // Check if timestamp is valid
151 if timestamp < old_record.timestamp {
152 return Err(FeoxError::OlderTimestamp);
153 }
154
155 // Load value from memory or disk
156 let value = if let Some(val) = old_record.get_value() {
157 val.to_vec()
158 } else if let Some(value) = self
159 .cache
160 .as_ref()
161 .and_then(|cache| cache.get_for_record(key, old_record))
162 {
163 value.to_vec()
164 } else {
165 let value = self.load_value_from_disk(old_record)?;
166 if let Some(ref cache) = self.cache {
167 cache.insert_for_record(
168 key_vec.clone(),
169 Bytes::from(value.clone()),
170 Arc::clone(old_record),
171 );
172 }
173 value
174 };
175
176 let current_val = if value.len() == 8 {
177 let bytes = value
178 .get(..8)
179 .and_then(|slice| slice.try_into().ok())
180 .ok_or(FeoxError::InvalidNumericValue)?;
181 i64::from_le_bytes(bytes)
182 } else {
183 return Err(FeoxError::InvalidOperation);
184 };
185
186 let new_val = current_val.saturating_add(delta);
187 let new_value = new_val.to_le_bytes().to_vec();
188
189 // Create new record with TTL if specified
190 let new_record = if ttl_seconds > 0 {
191 let ttl_expiry = timestamp + (ttl_seconds * 1_000_000_000); // Convert to nanoseconds
192 Arc::new(Record::new_with_timestamp_ttl(
193 old_record.key.clone(),
194 new_value,
195 timestamp,
196 ttl_expiry,
197 ))
198 } else {
199 Arc::new(Record::new(old_record.key.clone(), new_value, timestamp))
200 };
201
202 let old_value_len = old_record.value_len;
203 let old_size = old_record.calculate_size();
204 let new_size = self.calculate_record_size(old_record.key.len(), 8);
205 let old_record_arc = Arc::clone(old_record);
206
207 // Atomically update the entry
208 old_record_arc.refcount.store(0, Ordering::Release);
209 entry.insert(Arc::clone(&new_record));
210
211 // Update skip list as well
212 self.tree.insert(key_vec.clone(), Arc::clone(&new_record));
213
214 // Update memory usage
215 if new_size > old_size {
216 self.stats
217 .memory_usage
218 .fetch_add(new_size - old_size, Ordering::AcqRel);
219 } else {
220 self.stats
221 .memory_usage
222 .fetch_sub(old_size - new_size, Ordering::AcqRel);
223 }
224
225 // Only do cache and persistence operations if not in memory-only mode
226 if !self.memory_only {
227 if self.enable_caching {
228 if let Some(ref cache) = self.cache {
229 cache.remove(&key_vec);
230 }
231 }
232
233 if let Some(ref wb) = self.write_buffer {
234 wb.add_write(Operation::Update, Arc::clone(&new_record), old_value_len)?;
235 wb.add_write(Operation::Delete, old_record_arc, old_value_len)?;
236 }
237 }
238
239 Ok(new_val)
240 }
241 scc::hash_map::Entry::Vacant(entry) => {
242 // Key doesn't exist, create it with initial value
243 // Get timestamp inside the critical section
244 let timestamp = match timestamp {
245 Some(0) | None => self.get_timestamp(),
246 Some(ts) => ts,
247 };
248
249 let initial_val = delta;
250 let value = initial_val.to_le_bytes().to_vec();
251
252 // Create new record with TTL if specified
253 let new_record = if ttl_seconds > 0 {
254 let ttl_expiry = timestamp + (ttl_seconds * 1_000_000_000); // Convert to nanoseconds
255 Arc::new(Record::new_with_timestamp_ttl(
256 key_vec.clone(),
257 value,
258 timestamp,
259 ttl_expiry,
260 ))
261 } else {
262 Arc::new(Record::new(key_vec.clone(), value, timestamp))
263 };
264
265 let _ = entry.insert_entry(Arc::clone(&new_record));
266
267 // Update skip list
268 self.tree.insert(key_vec.clone(), Arc::clone(&new_record));
269
270 // Update statistics
271 self.stats.record_count.fetch_add(1, Ordering::AcqRel);
272 let record_size = self.calculate_record_size(key.len(), 8);
273 self.stats
274 .memory_usage
275 .fetch_add(record_size, Ordering::AcqRel);
276
277 // Handle persistence if needed
278 if !self.memory_only {
279 if let Some(ref wb) = self.write_buffer {
280 wb.add_write(Operation::Insert, Arc::clone(&new_record), 0)?;
281 }
282 }
283
284 Ok(initial_val)
285 }
286 };
287
288 result
289 }
290
291 /// Insert a key only when it does not already exist.
292 ///
293 /// The existence check and insertion happen under the hash-table entry guard, so
294 /// concurrent callers cannot both create the same key.
295 ///
296 /// # Returns
297 ///
298 /// Returns `Ok(true)` when this call inserted the key and `Ok(false)` when another
299 /// value already exists.
300 ///
301 /// # Example
302 ///
303 /// ```rust
304 /// # use feoxdb::FeoxStore;
305 /// # fn main() -> feoxdb::Result<()> {
306 /// let store = FeoxStore::new(None)?;
307 /// assert!(store.insert_if_absent(b"job:1", b"first")?);
308 /// assert!(!store.insert_if_absent(b"job:1", b"second")?);
309 /// assert_eq!(store.get(b"job:1")?, b"first");
310 /// # Ok(())
311 /// # }
312 /// ```
313 pub fn insert_if_absent(&self, key: &[u8], value: &[u8]) -> Result<bool> {
314 let start = std::time::Instant::now();
315 self.validate_key_value(key, value)?;
316 let key_vec = key.to_vec();
317
318 match self.hash_table.entry(key_vec.clone()) {
319 scc::hash_map::Entry::Occupied(_) => Ok(false),
320 scc::hash_map::Entry::Vacant(entry) => {
321 let record_size = self.calculate_record_size(key.len(), value.len());
322 if !self.check_memory_limit(record_size) {
323 return Err(FeoxError::OutOfMemory);
324 }
325
326 let record = Arc::new(Record::new(
327 key_vec.clone(),
328 value.to_vec(),
329 self.get_timestamp(),
330 ));
331 let _entry = entry.insert_entry(Arc::clone(&record));
332
333 self.tree.insert(key_vec, Arc::clone(&record));
334 self.stats.record_count.fetch_add(1, Ordering::AcqRel);
335 self.stats
336 .memory_usage
337 .fetch_add(record_size, Ordering::AcqRel);
338 self.stats
339 .record_insert(start.elapsed().as_nanos() as u64, false);
340
341 if !self.memory_only {
342 if let Some(ref write_buffer) = self.write_buffer {
343 write_buffer.add_write(Operation::Insert, record, 0)?;
344 }
345 }
346
347 Ok(true)
348 }
349 }
350 }
351
352 /// Atomically compare and swap a value.
353 ///
354 /// Compares the current value of a key with an expected value, and if they match,
355 /// atomically replaces it with a new value. This operation is atomic within the
356 /// HashMap shard, preventing race conditions.
357 ///
358 /// # Arguments
359 ///
360 /// * `key` - The key to check and potentially update
361 /// * `expected` - The expected current value
362 /// * `new_value` - The new value to set if comparison succeeds
363 ///
364 /// # Returns
365 ///
366 /// Returns `Ok(true)` if the swap succeeded (current value matched expected).
367 /// Returns `Ok(false)` if the current value didn't match or key doesn't exist.
368 ///
369 /// # Errors
370 ///
371 /// * `InvalidKeySize` - Key is invalid
372 /// * `InvalidValueSize` - New value is too large
373 /// * `OutOfMemory` - Memory limit exceeded
374 /// * `IoError` - Failed to read value from disk
375 ///
376 /// # Example
377 ///
378 /// ```rust
379 /// # use feoxdb::FeoxStore;
380 /// # fn main() -> feoxdb::Result<()> {
381 /// # let store = FeoxStore::new(None)?;
382 /// store.insert(b"config", b"v1")?;
383 ///
384 /// // Successful CAS - value matches
385 /// let swapped = store.compare_and_swap(b"config", b"v1", b"v2")?;
386 /// assert_eq!(swapped, true);
387 ///
388 /// // Failed CAS - value doesn't match
389 /// let swapped = store.compare_and_swap(b"config", b"v1", b"v3")?;
390 /// assert_eq!(swapped, false); // Value is now "v2", not "v1"
391 ///
392 /// // CAS on non-existent key
393 /// let swapped = store.compare_and_swap(b"missing", b"any", b"new")?;
394 /// assert_eq!(swapped, false);
395 /// # Ok(())
396 /// # }
397 /// ```
398 pub fn compare_and_swap(&self, key: &[u8], expected: &[u8], new_value: &[u8]) -> Result<bool> {
399 self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, None, 0)
400 }
401
402 /// Compare and swap with explicit timestamp.
403 ///
404 /// This is the advanced version that allows manual timestamp control for
405 /// conflict resolution. Most users should use `compare_and_swap()` instead.
406 ///
407 /// # Arguments
408 ///
409 /// * `key` - The key to check and potentially update
410 /// * `expected` - The expected current value
411 /// * `new_value` - The new value to set if comparison succeeds
412 /// * `timestamp` - Optional timestamp. If `None`, uses current time.
413 ///
414 /// # Errors
415 ///
416 /// * `OlderTimestamp` - Timestamp is not newer than existing record
417 pub fn compare_and_swap_with_timestamp(
418 &self,
419 key: &[u8],
420 expected: &[u8],
421 new_value: &[u8],
422 timestamp: Option<u64>,
423 ) -> Result<bool> {
424 self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, timestamp, 0)
425 }
426
427 /// Compare and swap with TTL support.
428 ///
429 /// # Arguments
430 ///
431 /// * `key` - The key to check and potentially update
432 /// * `expected` - The expected current value
433 /// * `new_value` - The new value to set if comparison succeeds
434 /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
435 ///
436 /// # Errors
437 ///
438 /// * `InvalidKeySize` - Key is invalid
439 /// * `InvalidValueSize` - New value is too large
440 pub fn compare_and_swap_with_ttl(
441 &self,
442 key: &[u8],
443 expected: &[u8],
444 new_value: &[u8],
445 ttl_seconds: u64,
446 ) -> Result<bool> {
447 self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, None, ttl_seconds)
448 }
449
450 /// Compare and swap with explicit timestamp and TTL.
451 ///
452 /// # Arguments
453 ///
454 /// * `key` - The key to check and potentially update
455 /// * `expected` - The expected current value
456 /// * `new_value` - The new value to set if comparison succeeds
457 /// * `timestamp` - Optional timestamp. If `None`, uses current time.
458 /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
459 ///
460 /// # Errors
461 ///
462 /// * `OlderTimestamp` - Timestamp is not newer than existing record
463 pub fn compare_and_swap_with_timestamp_and_ttl(
464 &self,
465 key: &[u8],
466 expected: &[u8],
467 new_value: &[u8],
468 timestamp: Option<u64>,
469 ttl_seconds: u64,
470 ) -> Result<bool> {
471 let start = std::time::Instant::now();
472 self.validate_key_value(key, new_value)?;
473 let key_vec = key.to_vec();
474
475 // Phase 1: Check value and save record reference for version tracking
476 let initial_record = {
477 let entry = match self.hash_table.read(&key_vec, |_, v| v.clone()) {
478 Some(e) => e,
479 None => return Ok(false), // Key doesn't exist
480 };
481
482 let record_arc = entry;
483
484 // Check if value matches expected
485 let value_matches = if let Some(val) = record_arc.get_value() {
486 // Fast path: value in memory
487 val.as_ref() == expected
488 } else if let Some(value) = self
489 .cache
490 .as_ref()
491 .and_then(|cache| cache.get_for_record(key, &record_arc))
492 {
493 value.as_ref() == expected
494 } else {
495 let disk_value = self.load_value_from_disk(&record_arc)?;
496 if let Some(ref cache) = self.cache {
497 cache.insert_for_record(
498 key_vec.clone(),
499 Bytes::from(disk_value.clone()),
500 Arc::clone(&record_arc),
501 );
502 }
503 disk_value == expected
504 };
505
506 if !value_matches {
507 return Ok(false); // Value doesn't match expected
508 }
509
510 // Return the Arc pointer itself as our version identifier
511 // NOTE: We can't use the stored timestamp for verification here because
512 // SystemTime::now() resolution is 1us, which is too coarse for CAS operations.
513 record_arc
514 };
515
516 // Phase 2: Acquire write lock and verify record hasn't changed
517 match self.hash_table.entry(key_vec.clone()) {
518 scc::hash_map::Entry::Occupied(mut entry) => {
519 let old_record = entry.get();
520
521 // Check if the record is still the same one we read earlier
522 if !Arc::ptr_eq(old_record, &initial_record) {
523 // Record was modified between our check and acquiring lock
524 return Ok(false);
525 }
526
527 let timestamp = match timestamp {
528 Some(0) | None => self.get_timestamp(),
529 Some(ts) => ts,
530 };
531
532 if timestamp < old_record.timestamp {
533 return Err(FeoxError::OlderTimestamp);
534 }
535
536 let old_size = old_record.calculate_size();
537 let new_size = self.calculate_record_size(key.len(), new_value.len());
538 let old_value_len = old_record.value_len;
539 let old_record_arc = Arc::clone(old_record);
540
541 // Pre-check memory limit
542 if new_size > old_size && !self.check_memory_limit(new_size - old_size) {
543 return Err(FeoxError::OutOfMemory);
544 }
545
546 // Create new record with TTL if specified
547 let new_record = if ttl_seconds > 0 {
548 let ttl_expiry = timestamp + (ttl_seconds * 1_000_000_000); // Convert to nanoseconds
549 Arc::new(Record::new_with_timestamp_ttl(
550 key.to_vec(),
551 new_value.to_vec(),
552 timestamp,
553 ttl_expiry,
554 ))
555 } else {
556 Arc::new(Record::new(key.to_vec(), new_value.to_vec(), timestamp))
557 };
558
559 old_record_arc.refcount.store(0, Ordering::Release);
560 entry.insert(Arc::clone(&new_record));
561
562 self.tree.insert(key_vec.clone(), Arc::clone(&new_record));
563
564 if new_size > old_size {
565 self.stats
566 .memory_usage
567 .fetch_add(new_size - old_size, Ordering::AcqRel);
568 } else {
569 self.stats
570 .memory_usage
571 .fetch_sub(old_size - new_size, Ordering::AcqRel);
572 }
573
574 self.stats
575 .record_insert(start.elapsed().as_nanos() as u64, true);
576
577 if !self.memory_only {
578 if self.enable_caching {
579 if let Some(ref cache) = self.cache {
580 cache.remove(&key_vec);
581 }
582 }
583
584 if let Some(ref wb) = self.write_buffer {
585 wb.add_write(Operation::Update, new_record, old_value_len)?;
586 wb.add_write(Operation::Delete, old_record_arc, old_value_len)?;
587 }
588 }
589
590 Ok(true)
591 }
592 scc::hash_map::Entry::Vacant(_) => Ok(false),
593 }
594 }
595}