feoxdb/core/store/atomic.rs
1use bytes::{BufMut, BytesMut};
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 if ttl_seconds > 0 {
137 self.ensure_ttl_write_supported()?;
138 }
139 self.validate_new_key(key)?;
140
141 let key_vec = key.to_vec();
142 let explicit_timestamp = timestamp.filter(|timestamp| *timestamp != 0);
143 let mut observed = None;
144
145 loop {
146 let Some(current) = self.hash_table.read(key, |_, record| Arc::clone(record)) else {
147 let retired_at = observed
148 .as_ref()
149 .map_or(0, |record: &Arc<Record>| record.retirement_timestamp());
150 let timestamp = match explicit_timestamp {
151 Some(timestamp) => timestamp,
152 None => {
153 let minimum = retired_at.checked_add(1).ok_or(FeoxError::OlderTimestamp)?;
154 self.get_timestamp(key).max(minimum)
155 }
156 };
157 if timestamp <= retired_at {
158 return Err(FeoxError::OlderTimestamp);
159 }
160
161 match self.hash_table.entry(key_vec.clone()) {
162 scc::hash_map::Entry::Occupied(_) => continue,
163 scc::hash_map::Entry::Vacant(entry) => {
164 let record_size = self.calculate_record_size(key.len(), 8);
165 let reservation = self.reserve_memory(record_size)?;
166 let new_record =
167 counter_record(key_vec.clone(), delta, timestamp, ttl_seconds);
168 let ttl_expiry = new_record.ttl_expiry.load(Ordering::Acquire);
169 let buffered_record = self
170 .write_buffer
171 .as_ref()
172 .filter(|_| !self.memory_only)
173 .map(|_| Arc::clone(&new_record));
174 let entry_guard = entry.insert_entry(Arc::clone(&new_record));
175 self.insert_into_tree(key_vec.clone(), new_record);
176 self.observe_published_timestamp(
177 key,
178 timestamp,
179 explicit_timestamp.is_some(),
180 );
181 reservation.commit();
182 self.note_ttl_transition(0, ttl_expiry);
183 self.stats.record_count.fetch_add(1, Ordering::Relaxed);
184 drop(entry_guard);
185
186 if let (Some(write_buffer), Some(record)) =
187 (&self.write_buffer, buffered_record)
188 {
189 write_buffer.add_write(Operation::Insert, record, 0)?;
190 }
191 return Ok(delta);
192 }
193 }
194 };
195
196 let root = observed.get_or_insert_with(|| Arc::clone(¤t));
197 if explicit_timestamp.is_some_and(|timestamp| timestamp <= current.timestamp) {
198 return Err(FeoxError::OlderTimestamp);
199 }
200
201 let expiry = current.ttl_expiry.load(Ordering::Acquire);
202 if self.enable_ttl && expiry > 0 {
203 let now = self.get_timestamp_pub();
204 if now > expiry {
205 self.retire_expired_if_current(key, ¤t, now)?;
206 continue;
207 }
208 }
209
210 let (value, _, source) = match self.resolve_value(key, Arc::clone(¤t)) {
211 Ok(resolved) => resolved,
212 Err(FeoxError::KeyNotFound) => {
213 let now = self.get_timestamp_pub();
214 self.retire_expired_if_current(key, ¤t, now)?;
215 continue;
216 }
217 Err(error) => return Err(error),
218 };
219 if value.len() != 8 {
220 return Err(FeoxError::InvalidOperation);
221 }
222 let current_value = i64::from_le_bytes(
223 value
224 .as_ref()
225 .try_into()
226 .map_err(|_| FeoxError::InvalidNumericValue)?,
227 );
228 let new_value = current_value.saturating_add(delta);
229 let timestamp = explicit_timestamp.unwrap_or_else(|| self.get_timestamp(key));
230
231 match self.hash_table.entry(key_vec.clone()) {
232 scc::hash_map::Entry::Occupied(mut entry) => {
233 let old_record = entry.get();
234 if !Arc::ptr_eq(old_record, &source) {
235 if explicit_timestamp
236 .is_some_and(|timestamp| timestamp <= root.retirement_timestamp())
237 {
238 return Err(FeoxError::OlderTimestamp);
239 }
240 continue;
241 }
242 if timestamp <= old_record.timestamp {
243 return Err(FeoxError::OlderTimestamp);
244 }
245
246 let old_size = old_record.calculate_size();
247 let new_size = self.calculate_record_size(old_record.key.len(), 8);
248 let reservation = self.reserve_memory(new_size.saturating_sub(old_size))?;
249 let old_record = Arc::clone(old_record);
250 let record = counter_record(key_vec.clone(), new_value, timestamp, ttl_seconds);
251 let old_expiry = old_record.ttl_expiry.load(Ordering::Acquire);
252 let new_expiry = record.ttl_expiry.load(Ordering::Acquire);
253
254 old_record.link_successor(&record);
255 old_record.refcount.store(0, Ordering::Release);
256 entry.insert(Arc::clone(&record));
257 self.publish_to_tree(&key_vec, Arc::clone(&record));
258 self.observe_published_timestamp(key, timestamp, explicit_timestamp.is_some());
259 reservation.commit();
260 self.note_ttl_transition(old_expiry, new_expiry);
261 if old_size > new_size {
262 self.release_memory(old_size - new_size);
263 }
264 drop(entry);
265
266 if !self.memory_only {
267 if self.enable_caching {
268 if let Some(cache) = &self.cache {
269 cache.remove_for_record(&key_vec, &old_record);
270 }
271 }
272
273 if let Some(write_buffer) = &self.write_buffer {
274 write_buffer.add_replacement(record, old_record)?;
275 }
276 }
277 return Ok(new_value);
278 }
279 scc::hash_map::Entry::Vacant(_) => {
280 if explicit_timestamp
281 .is_some_and(|timestamp| timestamp <= root.retirement_timestamp())
282 {
283 return Err(FeoxError::OlderTimestamp);
284 }
285 }
286 }
287 }
288 }
289
290 /// Insert a key only when it does not already exist.
291 ///
292 /// The existence check and insertion happen under the hash-table entry guard, so
293 /// concurrent callers cannot both create the same key.
294 ///
295 /// # Returns
296 ///
297 /// Returns `Ok(true)` when this call inserted the key and `Ok(false)` when another
298 /// value already exists.
299 ///
300 /// # Example
301 ///
302 /// ```rust
303 /// # use feoxdb::FeoxStore;
304 /// # fn main() -> feoxdb::Result<()> {
305 /// let store = FeoxStore::new(None)?;
306 /// assert!(store.insert_if_absent(b"job:1", b"first")?);
307 /// assert!(!store.insert_if_absent(b"job:1", b"second")?);
308 /// assert_eq!(store.get(b"job:1")?, b"first");
309 /// # Ok(())
310 /// # }
311 /// ```
312 pub fn insert_if_absent(&self, key: &[u8], value: &[u8]) -> Result<bool> {
313 let start = std::time::Instant::now();
314 self.validate_key_value(key, value)?;
315 let key_vec = key.to_vec();
316
317 match self.hash_table.entry(key_vec.clone()) {
318 scc::hash_map::Entry::Occupied(_) => Ok(false),
319 scc::hash_map::Entry::Vacant(entry) => {
320 let record_size = self.calculate_record_size(key.len(), value.len());
321 let reservation = self.reserve_memory(record_size)?;
322
323 let timestamp = self.get_timestamp(key);
324 let record = Arc::new(Record::new(key_vec.clone(), value.to_vec(), timestamp));
325 let buffered_record = self
326 .write_buffer
327 .as_ref()
328 .filter(|_| !self.memory_only)
329 .map(|_| Arc::clone(&record));
330 let entry_guard = entry.insert_entry(Arc::clone(&record));
331
332 self.insert_into_tree(key_vec, record);
333 reservation.commit();
334 self.stats.record_count.fetch_add(1, Ordering::Relaxed);
335 drop(entry_guard);
336 self.stats
337 .record_insert(start.elapsed().as_nanos() as u64, false);
338
339 if let (Some(write_buffer), Some(record)) = (&self.write_buffer, buffered_record) {
340 write_buffer.add_write(Operation::Insert, record, 0)?;
341 }
342
343 Ok(true)
344 }
345 }
346 }
347
348 /// Atomically compare and swap a value.
349 ///
350 /// Compares the current value of a key with an expected value, and if they match,
351 /// atomically replaces it with a new value. This operation is atomic within the
352 /// HashMap shard, preventing race conditions.
353 ///
354 /// # Arguments
355 ///
356 /// * `key` - The key to check and potentially update
357 /// * `expected` - The expected current value
358 /// * `new_value` - The new value to set if comparison succeeds
359 ///
360 /// # Returns
361 ///
362 /// Returns `Ok(true)` if the swap succeeded (current value matched expected).
363 /// Returns `Ok(false)` if the current value didn't match or key doesn't exist.
364 ///
365 /// # Errors
366 ///
367 /// * `InvalidKeySize` - Key is invalid
368 /// * `InvalidValueSize` - New value is too large
369 /// * `OutOfMemory` - Memory limit exceeded
370 /// * `IoError` - Failed to read value from disk
371 ///
372 /// # Example
373 ///
374 /// ```rust
375 /// # use feoxdb::FeoxStore;
376 /// # fn main() -> feoxdb::Result<()> {
377 /// # let store = FeoxStore::new(None)?;
378 /// store.insert(b"config", b"v1")?;
379 ///
380 /// // Successful CAS - value matches
381 /// let swapped = store.compare_and_swap(b"config", b"v1", b"v2")?;
382 /// assert_eq!(swapped, true);
383 ///
384 /// // Failed CAS - value doesn't match
385 /// let swapped = store.compare_and_swap(b"config", b"v1", b"v3")?;
386 /// assert_eq!(swapped, false); // Value is now "v2", not "v1"
387 ///
388 /// // CAS on non-existent key
389 /// let swapped = store.compare_and_swap(b"missing", b"any", b"new")?;
390 /// assert_eq!(swapped, false);
391 /// # Ok(())
392 /// # }
393 /// ```
394 pub fn compare_and_swap(&self, key: &[u8], expected: &[u8], new_value: &[u8]) -> Result<bool> {
395 self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, None, 0)
396 }
397
398 /// Compare and swap with explicit timestamp.
399 ///
400 /// This is the advanced version that allows manual timestamp control for
401 /// conflict resolution. Most users should use `compare_and_swap()` instead.
402 ///
403 /// # Arguments
404 ///
405 /// * `key` - The key to check and potentially update
406 /// * `expected` - The expected current value
407 /// * `new_value` - The new value to set if comparison succeeds
408 /// * `timestamp` - Optional timestamp. If `None`, uses current time.
409 ///
410 /// # Errors
411 ///
412 /// * `OlderTimestamp` - Timestamp is not newer than existing record
413 pub fn compare_and_swap_with_timestamp(
414 &self,
415 key: &[u8],
416 expected: &[u8],
417 new_value: &[u8],
418 timestamp: Option<u64>,
419 ) -> Result<bool> {
420 self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, timestamp, 0)
421 }
422
423 /// Compare and swap with TTL support.
424 ///
425 /// # Arguments
426 ///
427 /// * `key` - The key to check and potentially update
428 /// * `expected` - The expected current value
429 /// * `new_value` - The new value to set if comparison succeeds
430 /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
431 ///
432 /// # Errors
433 ///
434 /// * `InvalidKeySize` - Key is invalid
435 /// * `InvalidValueSize` - New value is too large
436 pub fn compare_and_swap_with_ttl(
437 &self,
438 key: &[u8],
439 expected: &[u8],
440 new_value: &[u8],
441 ttl_seconds: u64,
442 ) -> Result<bool> {
443 self.compare_and_swap_with_timestamp_and_ttl(key, expected, new_value, None, ttl_seconds)
444 }
445
446 /// Compare and swap with explicit timestamp and TTL.
447 ///
448 /// # Arguments
449 ///
450 /// * `key` - The key to check and potentially update
451 /// * `expected` - The expected current value
452 /// * `new_value` - The new value to set if comparison succeeds
453 /// * `timestamp` - Optional timestamp. If `None`, uses current time.
454 /// * `ttl_seconds` - Time-to-live in seconds (0 for no expiry)
455 ///
456 /// # Errors
457 ///
458 /// * `OlderTimestamp` - Timestamp is not newer than existing record
459 pub fn compare_and_swap_with_timestamp_and_ttl(
460 &self,
461 key: &[u8],
462 expected: &[u8],
463 new_value: &[u8],
464 timestamp: Option<u64>,
465 ttl_seconds: u64,
466 ) -> Result<bool> {
467 if ttl_seconds > 0 {
468 self.ensure_ttl_write_supported()?;
469 }
470 let start = std::time::Instant::now();
471 self.validate_key_value(key, new_value)?;
472 let key_vec = key.to_vec();
473
474 let initial_record = {
475 let record = match self
476 .hash_table
477 .read(&key_vec, |_, record| Arc::clone(record))
478 {
479 Some(record) => record,
480 None => return Ok(false),
481 };
482 let (value, cache_hit, source) = match self.resolve_value(key, record) {
483 Ok(resolved) => resolved,
484 Err(FeoxError::KeyNotFound | FeoxError::StaleExtent) => return Ok(false),
485 Err(error) => return Err(error),
486 };
487
488 if !cache_hit {
489 if let Some(cache) = &self.cache {
490 cache.insert_for_record(key_vec.clone(), value.clone(), &source);
491 }
492 }
493 if value.as_ref() != expected {
494 return Ok(false);
495 }
496 source
497 };
498
499 let timestamp = self.resolve_timestamp(key, timestamp);
500 self.replace_record_if_current(
501 &key_vec,
502 &initial_record,
503 new_value,
504 timestamp,
505 ttl_seconds,
506 start,
507 )
508 }
509
510 pub(super) fn replace_record_if_current(
511 &self,
512 key: &[u8],
513 expected: &Arc<Record>,
514 new_value: &[u8],
515 timestamp: (u64, bool),
516 ttl_seconds: u64,
517 start: std::time::Instant,
518 ) -> Result<bool> {
519 let (timestamp, explicit_timestamp) = timestamp;
520 match self.hash_table.entry(key.to_vec()) {
521 scc::hash_map::Entry::Occupied(mut entry) => {
522 let old_record = entry.get();
523
524 if !Arc::ptr_eq(old_record, expected) {
525 return Ok(false);
526 }
527
528 if timestamp <= old_record.timestamp {
529 return Err(FeoxError::OlderTimestamp);
530 }
531
532 let old_size = old_record.calculate_size();
533 let new_size = self.calculate_record_size(key.len(), new_value.len());
534 let reservation = self.reserve_memory(new_size.saturating_sub(old_size))?;
535 let old_record_arc = Arc::clone(old_record);
536 let old_expiry = old_record_arc.ttl_expiry.load(Ordering::Acquire);
537
538 // Create new record with TTL if specified
539 let new_record = if ttl_seconds > 0 {
540 let ttl_expiry =
541 timestamp.saturating_add(ttl_seconds.saturating_mul(1_000_000_000));
542 Arc::new(Record::new_with_timestamp_ttl(
543 key.to_vec(),
544 new_value.to_vec(),
545 timestamp,
546 ttl_expiry,
547 ))
548 } else {
549 Arc::new(Record::new(key.to_vec(), new_value.to_vec(), timestamp))
550 };
551
552 old_record_arc.link_successor(&new_record);
553 old_record_arc.refcount.store(0, Ordering::Release);
554 entry.insert(Arc::clone(&new_record));
555 self.publish_to_tree(key, Arc::clone(&new_record));
556 self.observe_published_timestamp(key, timestamp, explicit_timestamp);
557 reservation.commit();
558 self.note_ttl_transition(old_expiry, new_record.ttl_expiry.load(Ordering::Acquire));
559 if old_size > new_size {
560 self.release_memory(old_size - new_size);
561 }
562 drop(entry);
563
564 self.stats
565 .record_insert(start.elapsed().as_nanos() as u64, true);
566
567 if !self.memory_only {
568 if self.enable_caching {
569 if let Some(ref cache) = self.cache {
570 cache.remove_for_record(key, &old_record_arc);
571 }
572 }
573
574 if let Some(ref wb) = self.write_buffer {
575 wb.add_replacement(new_record, old_record_arc)?;
576 }
577 }
578
579 Ok(true)
580 }
581 scc::hash_map::Entry::Vacant(_) => Ok(false),
582 }
583 }
584}
585
586fn counter_record(key: Vec<u8>, value: i64, timestamp: u64, ttl_seconds: u64) -> Arc<Record> {
587 let mut bytes = BytesMut::with_capacity(std::mem::size_of::<i64>());
588 bytes.put_i64_le(value);
589 let value = bytes.freeze();
590
591 if ttl_seconds > 0 {
592 Arc::new(Record::new_from_bytes_with_ttl(
593 key,
594 value,
595 timestamp,
596 timestamp.saturating_add(ttl_seconds.saturating_mul(1_000_000_000)),
597 ))
598 } else {
599 Arc::new(Record::new_from_bytes(key, value, timestamp))
600 }
601}