1use crate::error::{ErrorData, Result};
24use crate::providers::local_store::{
25 as_blob, as_i64, as_opt_i64, as_text, opt_i64_value, query_all, LocalStore, StoreSpec,
26};
27use crate::traits::{Binding, Kv, KvEntry, PutCondition, PutOptions, ScanResult};
28use alien_error::{AlienError, Context as _, IntoAlienError as _};
29use async_trait::async_trait;
30use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
31use chrono::Utc;
32use serde::{Deserialize, Serialize};
33use std::path::PathBuf;
34use turso::Connection;
35use uuid::Uuid;
36
37use super::{decode_version, encode_version};
38
39static KV_SPEC: StoreSpec = StoreSpec {
40 db_filename: "localkv.sqlite",
41 format_version: "localkv.v2",
42 binding_type: "local KV",
43 schema_ddl: "CREATE TABLE IF NOT EXISTS kv (key TEXT PRIMARY KEY, value BLOB NOT NULL, expires_at INTEGER, version TEXT NOT NULL);",
44};
45
46#[derive(Debug)]
47pub struct LocalKv {
48 store: LocalStore,
49}
50
51#[derive(Debug, Serialize, Deserialize)]
52struct CursorState {
53 version: u8,
54 prefix: String,
55 last_key: String,
56}
57
58fn kv_error(operation: &str, key: &str, reason: &str) -> ErrorData {
60 ErrorData::KvOperationFailed {
61 operation: operation.to_string(),
62 key: key.to_string(),
63 reason: reason.to_string(),
64 }
65}
66
67async fn delete_expired(conn: &Connection, operation: &str, key: &str, now: i64) -> Result<()> {
70 conn.execute(
71 "DELETE FROM kv WHERE key = ?1 AND expires_at IS NOT NULL AND expires_at <= ?2",
72 (key, now),
73 )
74 .await
75 .into_alien_error()
76 .context(kv_error(operation, key, "failed to delete expired row"))?;
77 Ok(())
78}
79
80impl LocalKv {
81 pub async fn new(data_dir: PathBuf) -> Result<Self> {
82 Ok(Self {
83 store: LocalStore::open(data_dir, &KV_SPEC).await?,
84 })
85 }
86
87 pub fn data_dir(&self) -> &PathBuf {
89 self.store.data_dir()
90 }
91
92 pub async fn len(&self) -> Result<usize> {
95 self.store
96 .with_conn(|conn| async move {
97 let rows = query_all(&conn, "SELECT COUNT(*) FROM kv", ())
98 .await
99 .into_alien_error()
100 .context(kv_error("len", "*", "failed to count rows"))?;
101 let count = rows
102 .first()
103 .and_then(|row| row.first())
104 .and_then(as_i64)
105 .ok_or_else(|| {
106 AlienError::new(kv_error("len", "*", "count query returned no value"))
107 })?;
108 Ok(count as usize)
109 })
110 .await
111 }
112
113 pub async fn is_empty(&self) -> Result<bool> {
116 Ok(self.len().await? == 0)
117 }
118
119 pub async fn clear(&self) -> Result<()> {
122 self.store
123 .with_conn(|conn| async move {
124 conn.execute("DELETE FROM kv", ())
125 .await
126 .into_alien_error()
127 .context(kv_error("clear", "*", "failed to clear local KV store"))?;
128 Ok(())
129 })
130 .await
131 }
132
133 pub async fn keys(&self) -> Result<Vec<String>> {
136 self.store
137 .with_conn(|conn| async move {
138 let rows = query_all(&conn, "SELECT key FROM kv", ())
139 .await
140 .into_alien_error()
141 .context(kv_error("keys", "*", "failed to scan keys"))?;
142 let mut keys = Vec::with_capacity(rows.len());
143 for row in &rows {
144 keys.push(row.first().and_then(as_text).ok_or_else(|| {
145 AlienError::new(kv_error("keys", "*", "failed to read key row"))
146 })?);
147 }
148 Ok(keys)
149 })
150 .await
151 }
152
153 fn validate_key(key: &str) -> Result<()> {
155 crate::providers::kv::validate_key(key)
156 }
157
158 fn validate_value(value: &[u8]) -> Result<()> {
160 crate::providers::kv::validate_value(value)
161 }
162
163 fn encode_cursor(state: &CursorState) -> Result<String> {
164 let bytes =
165 serde_json::to_vec(state)
166 .into_alien_error()
167 .context(ErrorData::InvalidInput {
168 operation_context: "Local KV cursor encoding".to_string(),
169 details: "Failed to serialize cursor state".to_string(),
170 field_name: Some("cursor".to_string()),
171 })?;
172 Ok(URL_SAFE_NO_PAD.encode(bytes))
173 }
174
175 fn decode_cursor(prefix: &str, cursor: &str) -> Result<CursorState> {
176 let bytes =
177 URL_SAFE_NO_PAD
178 .decode(cursor)
179 .into_alien_error()
180 .context(ErrorData::InvalidInput {
181 operation_context: "Local KV cursor decoding".to_string(),
182 details: "Invalid cursor encoding".to_string(),
183 field_name: Some("cursor".to_string()),
184 })?;
185 let state: CursorState =
186 serde_json::from_slice(&bytes)
187 .into_alien_error()
188 .context(ErrorData::InvalidInput {
189 operation_context: "Local KV cursor decoding".to_string(),
190 details: "Invalid cursor data".to_string(),
191 field_name: Some("cursor".to_string()),
192 })?;
193 if state.version != 1 || state.prefix != prefix || !state.last_key.starts_with(prefix) {
194 return Err(AlienError::new(ErrorData::InvalidInput {
195 operation_context: "Local KV cursor validation".to_string(),
196 details: "Cursor does not belong to this prefix scan".to_string(),
197 field_name: Some("cursor".to_string()),
198 }));
199 }
200 Ok(state)
201 }
202}
203
204impl Binding for LocalKv {}
205
206#[async_trait]
207impl Kv for LocalKv {
208 async fn get(&self, key: &str) -> Result<Option<KvEntry>> {
209 Self::validate_key(key)?;
210
211 self.store
212 .with_conn(|conn| async move {
213 let now = Utc::now().timestamp_millis();
214 let rows = query_all(
215 &conn,
216 "SELECT value, expires_at, version FROM kv WHERE key = ?1",
217 (key,),
218 )
219 .await
220 .into_alien_error()
221 .context(kv_error("get", key, "failed to read entry"))?;
222
223 let Some(row) = rows.first() else {
224 return Ok(None);
225 };
226 let value = row.first().and_then(as_blob).ok_or_else(|| {
227 AlienError::new(kv_error("get", key, "stored value is not a blob"))
228 })?;
229 let expires_at = row.get(1).and_then(as_opt_i64).ok_or_else(|| {
230 AlienError::new(kv_error("get", key, "stored expires_at is not an integer"))
231 })?;
232 let backend_version = row.get(2).and_then(as_text).ok_or_else(|| {
233 AlienError::new(kv_error("get", key, "stored version is not text"))
234 })?;
235
236 if matches!(expires_at, Some(exp) if exp <= now) {
237 delete_expired(&conn, "get", key, now).await?;
238 return Ok(None);
239 }
240
241 Ok(Some(KvEntry {
242 key: key.to_string(),
243 value,
244 version: encode_version(key, backend_version, expires_at)?,
245 }))
246 })
247 .await
248 }
249
250 async fn put(&self, key: &str, value: Vec<u8>, options: Option<PutOptions>) -> Result<bool> {
251 Self::validate_key(key)?;
252 Self::validate_value(&value)?;
253 let options = options.unwrap_or_default();
254
255 self.store
256 .with_conn(|conn| async move {
257 let now = Utc::now().timestamp_millis();
258 let expires_at: Option<i64> = options
259 .ttl
260 .map(|d| now.saturating_add(i64::try_from(d.as_millis()).unwrap_or(i64::MAX)));
261 let new_version = Uuid::new_v4().simple().to_string();
262
263 if matches!(options.condition, PutCondition::Absent) {
264 let changed = conn
269 .execute(
270 "INSERT INTO kv (key, value, expires_at, version) VALUES (?1, ?2, ?3, ?4) \
271 ON CONFLICT(key) DO UPDATE SET value = ?2, expires_at = ?3, version = ?4 \
272 WHERE kv.expires_at IS NOT NULL AND kv.expires_at <= ?5",
273 (key, value, opt_i64_value(expires_at), new_version, now),
274 )
275 .await
276 .into_alien_error()
277 .context(kv_error("put", key, "failed conditional put"))?;
278 Ok(changed == 1)
279 } else if let PutCondition::Version(if_version) = options.condition {
280 let expected = decode_version(key, &if_version)?;
281 if expected.expired {
282 return Ok(false);
283 }
284 let changed = conn
285 .execute(
286 "UPDATE kv SET value = ?2, expires_at = ?3, version = ?4 \
287 WHERE key = ?1 AND version = ?5 \
288 AND (expires_at IS NULL OR expires_at > ?6)",
289 (
290 key,
291 value,
292 opt_i64_value(expires_at),
293 new_version,
294 expected.backend_version,
295 now,
296 ),
297 )
298 .await
299 .into_alien_error()
300 .context(kv_error("put", key, "failed versioned put"))?;
301 Ok(changed == 1)
302 } else {
303 conn.execute(
304 "INSERT INTO kv (key, value, expires_at, version) VALUES (?1, ?2, ?3, ?4) \
305 ON CONFLICT(key) DO UPDATE SET value = ?2, expires_at = ?3, version = ?4",
306 (key, value, opt_i64_value(expires_at), new_version),
307 )
308 .await
309 .into_alien_error()
310 .context(kv_error("put", key, "failed to upsert value"))?;
311 Ok(true)
312 }
313 })
314 .await
315 }
316
317 async fn delete(&self, key: &str, if_version: Option<&str>) -> Result<bool> {
318 Self::validate_key(key)?;
319
320 self.store
321 .with_conn(|conn| async move {
322 if let Some(version) = if_version {
323 let expected = decode_version(key, version)?;
324 if expected.expired {
325 return Ok(false);
326 }
327 let changed = conn
328 .execute(
329 "DELETE FROM kv WHERE key = ?1 AND version = ?2 \
330 AND (expires_at IS NULL OR expires_at > ?3)",
331 (key, expected.backend_version, Utc::now().timestamp_millis()),
332 )
333 .await
334 .into_alien_error()
335 .context(kv_error("delete", key, "failed conditional delete"))?;
336 Ok(changed == 1)
337 } else {
338 conn.execute("DELETE FROM kv WHERE key = ?1", (key,))
339 .await
340 .into_alien_error()
341 .context(kv_error("delete", key, "failed to delete key"))?;
342 Ok(true)
343 }
344 })
345 .await
346 }
347
348 async fn exists(&self, key: &str) -> Result<bool> {
349 Self::validate_key(key)?;
350
351 self.store
352 .with_conn(|conn| async move {
353 let now = Utc::now().timestamp_millis();
354 let rows = query_all(&conn, "SELECT expires_at FROM kv WHERE key = ?1", (key,))
355 .await
356 .into_alien_error()
357 .context(kv_error("exists", key, "failed to check existence"))?;
358
359 let Some(row) = rows.first() else {
360 return Ok(false);
361 };
362 let expires_at = row.first().and_then(as_opt_i64).ok_or_else(|| {
363 AlienError::new(kv_error(
364 "exists",
365 key,
366 "stored expires_at is not an integer",
367 ))
368 })?;
369
370 if matches!(expires_at, Some(exp) if exp <= now) {
371 delete_expired(&conn, "exists", key, now).await?;
373 Ok(false)
374 } else {
375 Ok(true)
376 }
377 })
378 .await
379 }
380
381 async fn scan_prefix(
382 &self,
383 prefix: &str,
384 limit: Option<usize>,
385 cursor: Option<String>,
386 ) -> Result<ScanResult> {
387 Self::validate_key(prefix)?;
388
389 let limit = limit.unwrap_or(1000);
390 let last_key = cursor
391 .as_deref()
392 .map(|cursor| Self::decode_cursor(prefix, cursor))
393 .transpose()?
394 .map(|state| state.last_key);
395 if limit == 0 {
396 return Ok(ScanResult {
397 items: Vec::new(),
398 next_cursor: cursor,
399 });
400 }
401 let matching: Vec<KvEntry> =
403 self.store
404 .with_conn(|conn| async move {
405 let now = Utc::now().timestamp_millis();
406 let rows = match last_key.as_deref() {
407 Some(last_key) => {
408 query_all(
409 &conn,
410 "SELECT key, value, expires_at, version FROM kv WHERE key > ?1 ORDER BY key",
411 (last_key,),
412 )
413 .await
414 }
415 None => query_all(
416 &conn,
417 "SELECT key, value, expires_at, version FROM kv WHERE key >= ?1 ORDER BY key",
418 (prefix,),
419 )
420 .await,
421 }
422 .into_alien_error()
423 .context(kv_error(
424 "scan_prefix",
425 prefix,
426 "failed to scan prefix",
427 ))?;
428
429 let mut matching = Vec::new();
430 for row in &rows {
431 let k = row.first().and_then(as_text).ok_or_else(|| {
432 AlienError::new(kv_error(
433 "scan_prefix",
434 prefix,
435 "failed to read scan row key",
436 ))
437 })?;
438 if !k.starts_with(prefix) {
441 break;
442 }
443 let v = row.get(1).and_then(as_blob).ok_or_else(|| {
444 AlienError::new(kv_error(
445 "scan_prefix",
446 prefix,
447 "stored value is not a blob",
448 ))
449 })?;
450 let exp = row.get(2).and_then(as_opt_i64).ok_or_else(|| {
451 AlienError::new(kv_error(
452 "scan_prefix",
453 prefix,
454 "stored expires_at is not an integer",
455 ))
456 })?;
457 if matches!(exp, Some(e) if e <= now) {
458 continue; }
460 let backend_version = row.get(3).and_then(as_text).ok_or_else(|| {
461 AlienError::new(kv_error(
462 "scan_prefix",
463 prefix,
464 "stored version is not text",
465 ))
466 })?;
467 matching.push(KvEntry {
468 version: encode_version(&k, backend_version, exp)?,
469 key: k,
470 value: v,
471 });
472 }
473 Ok(matching)
474 })
475 .await?;
476
477 let has_more = matching.len() > limit;
478 let items = matching.into_iter().take(limit).collect::<Vec<_>>();
479 let next_cursor = if has_more {
480 items
481 .last()
482 .map(|entry| {
483 Self::encode_cursor(&CursorState {
484 version: 1,
485 prefix: prefix.to_string(),
486 last_key: entry.key.clone(),
487 })
488 })
489 .transpose()?
490 } else {
491 None
492 };
493
494 Ok(ScanResult { items, next_cursor })
495 }
496}
497
498#[cfg(test)]
499mod tests {
500 use super::*;
501 use crate::providers::local_store::open_database;
502 use std::sync::Arc;
503 use std::time::Duration;
504 use tempfile::TempDir;
505 use tokio::time;
506
507 async fn create_test_kv() -> (LocalKv, TempDir) {
508 let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
509 let kv = LocalKv::new(temp_dir.path().join("kv.db"))
510 .await
511 .expect("Failed to create LocalKv");
512 (kv, temp_dir)
513 }
514
515 #[tokio::test]
516 async fn test_basic_operations() {
517 let (kv, _temp_dir) = create_test_kv().await;
518
519 assert!(kv
520 .put("test_key", b"test_value".to_vec(), None)
521 .await
522 .unwrap());
523 let value = kv.get("test_key").await.unwrap().map(|entry| entry.value);
524 assert_eq!(value, Some(b"test_value".to_vec()));
525
526 assert!(kv.exists("test_key").await.unwrap());
527 assert!(!kv.exists("nonexistent").await.unwrap());
528
529 kv.delete("test_key", None).await.unwrap();
530 assert!(!kv.exists("test_key").await.unwrap());
531 assert_eq!(kv.get("test_key").await.unwrap(), None);
532 }
533
534 #[tokio::test]
535 async fn test_conditional_put() {
536 let (kv, _temp_dir) = create_test_kv().await;
537
538 let options = Some(PutOptions {
539 ttl: None,
540 condition: PutCondition::Absent,
541 });
542 assert!(kv
543 .put("key", b"value1".to_vec(), options.clone())
544 .await
545 .unwrap());
546
547 assert!(!kv.put("key", b"value2".to_vec(), options).await.unwrap());
548
549 assert_eq!(
550 kv.get("key").await.unwrap().map(|entry| entry.value),
551 Some(b"value1".to_vec())
552 );
553
554 assert!(kv.put("key", b"value3".to_vec(), None).await.unwrap());
555 assert_eq!(
556 kv.get("key").await.unwrap().map(|entry| entry.value),
557 Some(b"value3".to_vec())
558 );
559 }
560
561 #[tokio::test]
562 async fn test_ttl_expiration() {
563 let (kv, _temp_dir) = create_test_kv().await;
564
565 let options = Some(PutOptions {
566 ttl: Some(Duration::from_millis(500)),
567 condition: PutCondition::None,
568 });
569
570 kv.put("expiring_key", b"value".to_vec(), options)
571 .await
572 .unwrap();
573
574 assert!(kv.exists("expiring_key").await.unwrap());
575 assert_eq!(
576 kv.get("expiring_key")
577 .await
578 .unwrap()
579 .map(|entry| entry.value),
580 Some(b"value".to_vec())
581 );
582
583 time::sleep(Duration::from_millis(750)).await;
584
585 assert!(!kv.exists("expiring_key").await.unwrap());
586 assert_eq!(kv.get("expiring_key").await.unwrap(), None);
587 }
588
589 #[tokio::test]
590 async fn test_prefix_scanning() {
591 let (kv, _temp_dir) = create_test_kv().await;
592
593 kv.put("prefix:key1", b"value1".to_vec(), None)
594 .await
595 .unwrap();
596 kv.put("prefix:key2", b"value2".to_vec(), None)
597 .await
598 .unwrap();
599 kv.put("prefix:key3", b"value3".to_vec(), None)
600 .await
601 .unwrap();
602 kv.put("other:key", b"other".to_vec(), None).await.unwrap();
603
604 let result = kv.scan_prefix("prefix:", None, None).await.unwrap();
605 assert_eq!(result.items.len(), 3);
606 assert!(result.next_cursor.is_none());
607
608 assert_eq!(result.items[0].key, "prefix:key1");
609 assert_eq!(result.items[1].key, "prefix:key2");
610 assert_eq!(result.items[2].key, "prefix:key3");
611
612 let result = kv.scan_prefix("prefix:", Some(2), None).await.unwrap();
613 assert_eq!(result.items.len(), 2);
614 assert!(result.next_cursor.is_some());
615
616 let cursor = result.next_cursor.unwrap();
617 let result = kv
618 .scan_prefix("prefix:", Some(2), Some(cursor))
619 .await
620 .unwrap();
621 assert_eq!(result.items.len(), 1);
622 assert_eq!(result.items[0].key, "prefix:key3");
623 assert!(result.next_cursor.is_none());
624 }
625
626 #[tokio::test]
627 async fn prefix_cursor_does_not_skip_after_an_earlier_key_is_deleted() {
628 let (kv, _temp_dir) = create_test_kv().await;
629 for key in ["prefix:key1", "prefix:key2", "prefix:key3"] {
630 kv.put(key, key.as_bytes().to_vec(), None).await.unwrap();
631 }
632
633 let first = kv.scan_prefix("prefix:", Some(2), None).await.unwrap();
634 assert_eq!(
635 first
636 .items
637 .iter()
638 .map(|entry| entry.key.as_str())
639 .collect::<Vec<_>>(),
640 ["prefix:key1", "prefix:key2"]
641 );
642
643 kv.delete("prefix:key1", None).await.unwrap();
644 let second = kv
645 .scan_prefix("prefix:", Some(2), first.next_cursor)
646 .await
647 .unwrap();
648 assert_eq!(second.items[0].key, "prefix:key3");
649 assert!(second.next_cursor.is_none());
650 }
651
652 #[tokio::test]
653 async fn prefix_cursor_is_rejected_for_another_prefix() {
654 let (kv, _temp_dir) = create_test_kv().await;
655 for key in ["first:key1", "first:key2"] {
656 kv.put(key, key.as_bytes().to_vec(), None).await.unwrap();
657 }
658
659 let cursor = kv
660 .scan_prefix("first:", Some(1), None)
661 .await
662 .unwrap()
663 .next_cursor
664 .expect("the first page should have a cursor");
665 assert!(kv
666 .scan_prefix("second:", Some(1), Some(cursor.clone()))
667 .await
668 .is_err());
669 assert!(kv
670 .scan_prefix("second:", Some(0), Some(cursor))
671 .await
672 .is_err());
673 assert!(kv
674 .scan_prefix("first:", Some(0), Some("not-a-cursor".to_string()))
675 .await
676 .is_err());
677 }
678
679 #[tokio::test]
680 async fn test_persistence_across_reopens() {
681 let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
682 let db_path = temp_dir.path().join("kv.db");
683
684 {
685 let kv = LocalKv::new(db_path.clone())
686 .await
687 .expect("Failed to create LocalKv");
688 kv.put("persistent_key", b"persistent_value".to_vec(), None)
689 .await
690 .unwrap();
691 }
692
693 {
694 let kv = LocalKv::new(db_path)
695 .await
696 .expect("Failed to reopen LocalKv");
697 let value = kv
698 .get("persistent_key")
699 .await
700 .unwrap()
701 .map(|entry| entry.value);
702 assert_eq!(value, Some(b"persistent_value".to_vec()));
703 }
704 }
705
706 #[tokio::test]
707 async fn test_key_validation() {
708 let (kv, _temp_dir) = create_test_kv().await;
709
710 assert!(kv.put("", b"value".to_vec(), None).await.is_err());
711 assert!(kv.get("").await.is_err());
712
713 let long_key = "a".repeat(513);
714 assert!(kv.put(&long_key, b"value".to_vec(), None).await.is_err());
715
716 assert!(kv
717 .put("key with spaces", b"value".to_vec(), None)
718 .await
719 .is_err());
720 assert!(kv
721 .put("key\nwith\nnewlines", b"value".to_vec(), None)
722 .await
723 .is_err());
724 assert!(kv
725 .put("key/with/slashes", b"value".to_vec(), None)
726 .await
727 .is_err());
728
729 assert!(kv
730 .put("valid_key-123", b"value".to_vec(), None)
731 .await
732 .is_ok());
733 assert!(kv
734 .put("domain.com:8080", b"value".to_vec(), None)
735 .await
736 .is_ok());
737 }
738
739 #[tokio::test]
740 async fn test_value_validation() {
741 let (kv, _temp_dir) = create_test_kv().await;
742
743 let large_value = vec![0u8; 24_577];
744 assert!(kv.put("key", large_value, None).await.is_err());
745
746 let max_value = vec![0u8; 24_576];
747 assert!(kv.put("key", max_value, None).await.is_ok());
748 }
749
750 #[tokio::test]
751 async fn test_utility_methods() {
752 let (kv, _temp_dir) = create_test_kv().await;
753
754 assert!(kv.is_empty().await.unwrap());
755 assert_eq!(kv.len().await.unwrap(), 0);
756 assert_eq!(kv.keys().await.unwrap(), Vec::<String>::new());
757
758 kv.put("key1", b"value1".to_vec(), None).await.unwrap();
759 kv.put("key2", b"value2".to_vec(), None).await.unwrap();
760
761 assert!(!kv.is_empty().await.unwrap());
762 assert_eq!(kv.len().await.unwrap(), 2);
763
764 let mut keys = kv.keys().await.unwrap();
765 keys.sort();
766 assert_eq!(keys, vec!["key1", "key2"]);
767
768 kv.clear().await.unwrap();
769 assert!(kv.is_empty().await.unwrap());
770 assert_eq!(kv.len().await.unwrap(), 0);
771 }
772
773 #[tokio::test]
774 async fn test_unknown_format_rejected_on_open() {
775 let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
776 let dir = temp_dir.path().join("kv");
777
778 {
780 let kv = LocalKv::new(dir.clone()).await.expect("initial open");
781 kv.put("k", b"v".to_vec(), None).await.unwrap();
782 }
783 {
784 let db = open_database(&dir.join("localkv.sqlite"), "test")
785 .await
786 .expect("raw open");
787 let conn = db.connect().expect("raw connect");
788 conn.execute(
789 "UPDATE meta SET value = 'localkv.v99' WHERE key = 'format'",
790 (),
791 )
792 .await
793 .expect("format overwrite");
794 }
795
796 let err = LocalKv::new(dir)
798 .await
799 .expect_err("unknown format must be rejected");
800 let msg = err.to_string();
801 assert!(
802 msg.contains("localkv.v99"),
803 "error must name the found format, got: {msg}"
804 );
805 assert!(
806 msg.contains("localkv.v2"),
807 "error must name the expected format, got: {msg}"
808 );
809 }
810
811 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
814 async fn test_conditional_put_atomicity_across_handles() {
815 let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
816 let dir = temp_dir.path().join("kv");
817 let kv_a = Arc::new(LocalKv::new(dir.clone()).await.expect("open handle a"));
819 let kv_b = Arc::new(LocalKv::new(dir.clone()).await.expect("open handle b"));
820
821 let n = 16;
822 let mut handles = Vec::new();
823 for i in 0..n {
824 let kv = if i % 2 == 0 {
825 kv_a.clone()
826 } else {
827 kv_b.clone()
828 };
829 handles.push(tokio::spawn(async move {
830 let val = format!("val-{i}").into_bytes();
831 let opts = Some(PutOptions {
832 ttl: None,
833 condition: PutCondition::Absent,
834 });
835 let won = kv.put("race", val.clone(), opts).await.expect("put ok");
836 (won, val)
837 }));
838 }
839
840 let mut winners = Vec::new();
841 for h in handles {
842 let (won, val) = h.await.expect("task join");
843 if won {
844 winners.push(val);
845 }
846 }
847
848 assert_eq!(
849 winners.len(),
850 1,
851 "exactly one conditional put must win across both handles"
852 );
853 let stored = kv_a.get("race").await.unwrap().expect("key present");
854 assert_eq!(
855 stored.value, winners[0],
856 "stored value must equal the winner"
857 );
858 assert_eq!(
859 kv_b.get("race")
860 .await
861 .unwrap()
862 .expect("key present via b")
863 .value,
864 winners[0]
865 );
866 }
867
868 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
869 async fn test_ttl_expiry_takeover_conditional_put() {
870 let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
871 let dir = temp_dir.path().join("kv");
872 let kv_a = Arc::new(LocalKv::new(dir.clone()).await.expect("open handle a"));
873 let kv_b = Arc::new(LocalKv::new(dir.clone()).await.expect("open handle b"));
874
875 assert!(kv_a
877 .put(
878 "k",
879 b"initial".to_vec(),
880 Some(PutOptions {
881 ttl: Some(Duration::from_millis(300)),
882 condition: PutCondition::Absent,
883 }),
884 )
885 .await
886 .unwrap());
887
888 assert!(!kv_b
890 .put(
891 "k",
892 b"early".to_vec(),
893 Some(PutOptions {
894 ttl: None,
895 condition: PutCondition::Absent,
896 }),
897 )
898 .await
899 .unwrap());
900
901 time::sleep(Duration::from_millis(450)).await;
903
904 let n = 12;
906 let mut handles = Vec::new();
907 for i in 0..n {
908 let kv = if i % 2 == 0 {
909 kv_a.clone()
910 } else {
911 kv_b.clone()
912 };
913 handles.push(tokio::spawn(async move {
914 let val = format!("takeover-{i}").into_bytes();
915 let won = kv
916 .put(
917 "k",
918 val.clone(),
919 Some(PutOptions {
920 ttl: None,
921 condition: PutCondition::Absent,
922 }),
923 )
924 .await
925 .expect("put ok");
926 (won, val)
927 }));
928 }
929
930 let mut winners = Vec::new();
931 for h in handles {
932 let (won, val) = h.await.expect("task join");
933 if won {
934 winners.push(val);
935 }
936 }
937
938 assert_eq!(
939 winners.len(),
940 1,
941 "exactly one takeover conditional put must win after expiry"
942 );
943 let stored = kv_b.get("k").await.unwrap().expect("key present");
944 assert_eq!(
945 stored.value, winners[0],
946 "stored value must equal the takeover winner"
947 );
948 }
949
950 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
951 async fn test_multi_handle_concurrent_smoke() {
952 let temp_dir = tempfile::tempdir().expect("Failed to create temp dir");
953 let dir = temp_dir.path().join("kv");
954 let kv_a = Arc::new(LocalKv::new(dir.clone()).await.expect("open handle a"));
955 let kv_b = Arc::new(LocalKv::new(dir.clone()).await.expect("open handle b"));
956
957 let mut handles = Vec::new();
958 for i in 0..50 {
959 let kv = if i % 2 == 0 {
960 kv_a.clone()
961 } else {
962 kv_b.clone()
963 };
964 handles.push(tokio::spawn(async move {
965 let key = format!("key_{i}");
966 let val = format!("v{i}").into_bytes();
967 kv.put(&key, val.clone(), None).await.expect("put ok");
969 let got = kv.get(&key).await.expect("get ok");
970 assert_eq!(got.map(|entry| entry.value), Some(val));
971 }));
972 }
973 for h in handles {
974 h.await.expect("task join");
975 }
976
977 assert_eq!(kv_a.len().await.unwrap(), 50, "handle a sees all keys");
978 assert_eq!(kv_b.len().await.unwrap(), 50, "handle b sees all keys");
979 }
980}