Skip to main content

alien_bindings/providers/kv/
local.rs

1//! Local disk-persisted KV backed by turso (`localkv.v2`), multi-process safe.
2//!
3//! # Connection strategy
4//!
5//! turso is async-native and its `Connection` is `Send + Sync`, so there is no
6//! `spawn_blocking` boundary and no `Mutex<Connection>` anywhere. `LocalKv`
7//! holds one `turso::Database` handle on `<dataDir>/localkv.sqlite`; each
8//! operation opens its **own** short-lived connection from it and drops it
9//! when the operation completes, so no statement state leaks between
10//! operations.
11//!
12//! Correctness under concurrent access (multiple handles on one file, i.e.
13//! multiple processes) comes from turso's multi-process WAL mode — enabled
14//! explicitly, experimental upstream, and gated by the multi-handle tests
15//! below — plus a `busy_timeout` (writers wait for the write lock instead of
16//! failing with `Busy`). The schema is created once in [`LocalKv::new`]. Reads
17//! of live rows never take the write lock; a read that encounters an expired
18//! row escalates to a short delete (see `get` and `exists`). Conditional puts
19//! are a single atomic `INSERT ... ON CONFLICT DO UPDATE ... WHERE` so the
20//! race is resolved by the database, not by application-level locking.
21//!
22//! See `crates/alien-bindings/FORMAT.md` for the on-disk `localkv.v2` contract.
23use 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
58/// Build the standard KV operation error context.
59fn 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
67/// Delete a row that is already expired (`expires_at <= now`), used by the
68/// lazy-expiry paths of `get` and `exists`.
69async 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    /// Get the data directory path (the directory that holds `localkv.sqlite`).
88    pub fn data_dir(&self) -> &PathBuf {
89        self.store.data_dir()
90    }
91
92    /// Get the number of items currently stored (including expired items).
93    /// Useful for testing.
94    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    /// Check if the store is empty (including expired items).
114    /// Useful for testing.
115    pub async fn is_empty(&self) -> Result<bool> {
116        Ok(self.len().await? == 0)
117    }
118
119    /// Clear all data from the store.
120    /// Useful for testing.
121    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    /// Get all keys currently in the store (including expired ones).
134    /// Useful for testing and debugging.
135    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    /// Validate key constraints using global KV validation.
154    fn validate_key(key: &str) -> Result<()> {
155        crate::providers::kv::validate_key(key)
156    }
157
158    /// Validate value constraints using global KV validation.
159    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                    // One atomic statement: insert if absent, otherwise overwrite
265                    // ONLY when the existing row is already expired. The changed
266                    // row count (returned by `execute`) is 1 for the winner, 0
267                    // for a loser.
268                    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                    // Lazily remove the expired row.
372                    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        // Collect matching, non-expired items after the last key in sorted order.
402        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                        // Keys are ordered ascending starting at `prefix`; once a key
439                        // stops matching the prefix, no later key can match either.
440                        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; // expired: treat as absent
459                        }
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        // Create a valid store, then rewrite its format marker to a future version.
779        {
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        // Reopening must fail fast, naming both the found and expected formats.
797        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    // ---- Multi-process-safety proofs ----
812
813    #[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        // Two independent handles on the SAME data_dir == two processes sharing the file.
818        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        // Seed a short-lived key with a conditional put.
876        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        // While it is still live, a conditional put must lose.
889        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        // Wait for the seeded key to expire.
902        time::sleep(Duration::from_millis(450)).await;
903
904        // Race the takeover: many conditional puts against the now-expired key.
905        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                // No busy errors expected under multi-process WAL + busy_timeout.
968                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}