ternary_lease/
lib.rs

1//! # ternary-lease
2//!
3//! Distributed lease management for GPU resources with ternary states.
4
5use std::collections::HashMap;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum LeaseState { Held = 1, Expired = 0, Revoked = -1 }
9
10#[derive(Debug, Clone)]
11pub struct Lease {
12    pub id: u64,
13    pub resource: String,
14    pub holder: String,
15    pub ttl_ticks: u32,
16    pub remaining: u32,
17}
18
19pub struct LeaseManager {
20    leases: HashMap<u64, Lease>,
21    next_id: u64,
22    tick: u64,
23    revocations: u64,
24}
25
26impl LeaseManager {
27    pub fn new() -> Self {
28        Self { leases: HashMap::new(), next_id: 1, tick: 0, revocations: 0 }
29    }
30
31    pub fn acquire(&mut self, resource: &str, holder: &str, ttl: u32) -> u64 {
32        let id = self.next_id;
33        self.next_id += 1;
34        self.leases.insert(id, Lease { id, resource: resource.into(), holder: holder.into(), ttl_ticks: ttl, remaining: ttl });
35        id
36    }
37
38    pub fn renew(&mut self, lease_id: u64) -> bool {
39        if let Some(lease) = self.leases.get_mut(&lease_id) {
40            if lease.remaining > 0 {
41                lease.remaining = lease.ttl_ticks;
42                return true;
43            }
44        }
45        false
46    }
47
48    pub fn revoke(&mut self, lease_id: u64) -> bool {
49        if let Some(lease) = self.leases.get_mut(&lease_id) {
50            lease.remaining = 0;
51            self.revocations += 1;
52            return true;
53        }
54        false
55    }
56
57    pub fn state(&self, lease_id: u64) -> LeaseState {
58        match self.leases.get(&lease_id) {
59            None => LeaseState::Revoked,
60            Some(l) if l.remaining == 0 => LeaseState::Expired,
61            Some(_) => LeaseState::Held,
62        }
63    }
64
65    pub fn tick(&mut self) {
66        self.tick += 1;
67        for lease in self.leases.values_mut() {
68            if lease.remaining > 0 { lease.remaining -= 1; }
69        }
70    }
71
72    pub fn find_deadlocks(&self) -> Vec<(u64, u64)> {
73        // Simple cycle detection: A holds what B wants, B holds what A wants
74        let mut deadlocks = Vec::new();
75        let leases: Vec<&Lease> = self.leases.values().filter(|l| l.remaining > 0).collect();
76        for i in 0..leases.len() {
77            for j in (i+1)..leases.len() {
78                if leases[i].holder == leases[j].resource && leases[j].holder == leases[i].resource {
79                    deadlocks.push((leases[i].id, leases[j].id));
80                }
81            }
82        }
83        deadlocks
84    }
85
86    pub fn held_by(&self, holder: &str) -> Vec<&Lease> {
87        self.leases.values().filter(|l| l.remaining > 0 && l.holder == holder).collect()
88    }
89
90    pub fn active_count(&self) -> usize { self.leases.values().filter(|l| l.remaining > 0).count() }
91    pub fn revocations(&self) -> u64 { self.revocations }
92    pub fn tick_count(&self) -> u64 { self.tick }
93}
94
95impl Default for LeaseManager { fn default() -> Self { Self::new() } }
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_acquire_held() {
103        let mut lm = LeaseManager::new();
104        let id = lm.acquire("gpu0", "worker1", 10);
105        assert_eq!(lm.state(id), LeaseState::Held);
106    }
107
108    #[test]
109    fn test_expiry() {
110        let mut lm = LeaseManager::new();
111        let id = lm.acquire("gpu0", "worker1", 2);
112        lm.tick(); lm.tick();
113        assert_eq!(lm.state(id), LeaseState::Expired);
114    }
115
116    #[test]
117    fn test_renew() {
118        let mut lm = LeaseManager::new();
119        let id = lm.acquire("gpu0", "w", 3);
120        lm.tick(); lm.tick();
121        assert!(lm.renew(id));
122        lm.tick(); lm.tick();
123        assert_eq!(lm.state(id), LeaseState::Held);
124    }
125
126    #[test]
127    fn test_revoke() {
128        let mut lm = LeaseManager::new();
129        let id = lm.acquire("gpu0", "w", 10);
130        lm.revoke(id);
131        assert_eq!(lm.state(id), LeaseState::Expired);
132        assert_eq!(lm.revocations(), 1);
133    }
134
135    #[test]
136    fn test_deadlock_detection() {
137        let mut lm = LeaseManager::new();
138        // A holds "B", B holds "A" — circular wait
139        lm.acquire("B", "A", 10);
140        lm.acquire("A", "B", 10);
141        let deadlocks = lm.find_deadlocks();
142        assert_eq!(deadlocks.len(), 1);
143    }
144
145    #[test]
146    fn test_held_by() {
147        let mut lm = LeaseManager::new();
148        lm.acquire("gpu0", "w1", 10);
149        lm.acquire("gpu1", "w1", 10);
150        lm.acquire("gpu2", "w2", 10);
151        assert_eq!(lm.held_by("w1").len(), 2);
152    }
153
154    #[test]
155    fn test_active_count() {
156        let mut lm = LeaseManager::new();
157        lm.acquire("g0", "w", 1);
158        lm.acquire("g1", "w", 10);
159        lm.tick(); // g0 expires
160        assert_eq!(lm.active_count(), 1);
161    }
162
163    #[test]
164    fn test_renew_expired_fails() {
165        let mut lm = LeaseManager::new();
166        let id = lm.acquire("g0", "w", 1);
167        lm.tick(); // expired
168        assert!(!lm.renew(id));
169    }
170}
ternary_lease/lib.rs

ternary_lease/
lib.rs
