pdk-lock-lib 1.7.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use super::*;
use proxy_wasm_stub::types::Status;
use std::cell::RefCell;
use std::collections::HashMap;

#[test]
fn expires_with_time() {
    let clock = Rc::new(ManualClock::default());
    let shared_data = Rc::new(MapSharedData::new());

    let lock = TryLock::new(
        "key".to_string(),
        Duration::from_secs(10),
        Rc::clone(&clock) as Rc<dyn Clock>,
        shared_data,
    );

    let lock1 = lock.try_lock();
    let lock2 = lock.try_lock();

    assert!(lock1.is_some());
    assert!(lock2.is_none());
    clock.move_forward(Duration::from_millis(10001));
    let lock3 = lock.try_lock();
    assert!(lock3.is_some());
}

#[test]
fn release_after_acquire() {
    let clock = Rc::new(ManualClock::default());
    let shared_data = Rc::new(MapSharedData::new());

    let lock = TryLock::new(
        "key".to_string(),
        Duration::from_secs(10),
        Rc::clone(&clock) as Rc<dyn Clock>,
        shared_data,
    );

    let lock1 = lock.try_lock();
    assert!(lock1.is_some());
    drop(lock1);

    let lock1 = lock.try_lock();
    assert!(lock1.is_some());
    clock.move_forward(Duration::from_millis(5000));
    drop(lock1);

    let lock1 = lock.try_lock();
    assert!(lock1.is_some());
    clock.move_forward(Duration::from_millis(20000));
    drop(lock1);
}

#[test]
fn refresh_lock() {
    let clock = Rc::new(ManualClock::default());
    let shared_data = Rc::new(MapSharedData::new());

    let lock = TryLock::new(
        "key".to_string(),
        Duration::from_secs(10),
        Rc::clone(&clock) as Rc<dyn Clock>,
        shared_data,
    );

    if let Some(lock1) = lock.try_lock() {
        assert!(lock.try_lock().is_none());
        clock.move_forward(Duration::from_secs(5));
        assert!(lock1.refresh_lock());
        assert!(lock.try_lock().is_none());
        clock.move_forward(Duration::from_secs(6));
        assert!(lock.try_lock().is_none());
        clock.move_forward(Duration::from_secs(6));
        let lock2 = lock.try_lock();
        assert!(lock2.is_some());
        assert!(!lock1.refresh_lock());
    };
}

// ----- Clock -----
pub struct ManualClock {
    current: RefCell<SystemTime>,
}

impl ManualClock {
    fn new() -> Self {
        Self {
            current: RefCell::new(SystemTime::now()),
        }
    }
}

impl Default for ManualClock {
    /// Returns the manual clock implementation that is useful for testing
    fn default() -> Self {
        Self::new()
    }
}

impl Clock for ManualClock {
    fn get_current_time(&self) -> SystemTime {
        *self.current.borrow()
    }
}

impl ManualClock {
    /// Uses interior mutability pattern to allow seamless modification
    pub fn move_forward(&self, duration: Duration) {
        self.current.replace(self.get_current_time().add(duration));
    }
}

// ----- Shared Data -----
#[derive(Default)]
pub struct MapSharedData {
    map: RefCell<HashMap<String, MapSharedDataEntry>>,
}

struct MapSharedDataEntry {
    data: Option<Vec<u8>>,
    lock: u32,
}

impl MapSharedData {
    pub fn new() -> Self {
        Self {
            map: RefCell::new(HashMap::new()),
        }
    }

    fn do_store(&self, key: &str, value: &[u8], version: u32) {
        self.map.borrow_mut().insert(
            key.to_string(),
            MapSharedDataEntry {
                data: Some(value.to_vec()),
                lock: version,
            },
        );
    }

    fn next_cas(cas: Option<u32>) -> u32 {
        match cas {
            Some(u32::MAX) | None => 1,
            Some(cas) => cas + 1,
        }
    }
}

impl SharedData for MapSharedData {
    fn shared_data_get(&self, key: &str) -> (Option<Vec<u8>>, Option<u32>) {
        let key = key.to_string();
        self.map
            .borrow()
            .get(&key)
            .map(|entry| (entry.data.clone(), Some(entry.lock)))
            .unwrap_or_else(|| (None, None))
    }

    fn shared_data_set(&self, key: &str, value: &[u8], version: Option<u32>) -> Result<(), Status> {
        let (_, lock) = self.shared_data_get(key);
        if version.is_none() || lock.is_none() || version.unwrap() == lock.unwrap() {
            self.do_store(key, value, Self::next_cas(version));
            Ok(())
        } else {
            Err(Status::CasMismatch)
        }
    }

    fn shared_data_remove(
        &self,
        key: &str,
        version: Option<u32>,
    ) -> Result<Option<Vec<u8>>, Status> {
        let (val, _) = self.shared_data_get(key);
        self.shared_data_set(key, &[], version)?;
        Ok(val)
    }

    fn shared_data_keys(&self) -> Vec<String> {
        self.map.borrow().keys().cloned().collect()
    }
}