athena-worker 3.18.0

Shared worker runtime primitives for Athena background jobs
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use chrono::{DateTime, Duration, Utc};
use uuid::Uuid;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LeaseRecord {
    pub owner_token: Uuid,
    pub expires_at: DateTime<Utc>,
}

pub fn can_claim(existing: Option<&LeaseRecord>, now: DateTime<Utc>) -> bool {
    match existing {
        Some(lease) => lease.expires_at <= now,
        None => true,
    }
}

pub fn lease_owned_by(
    existing: Option<&LeaseRecord>,
    owner_token: Uuid,
    now: DateTime<Utc>,
) -> bool {
    existing.is_some_and(|lease| lease.owner_token == owner_token && lease.expires_at > now)
}

pub fn renew_lease(
    existing: &LeaseRecord,
    owner_token: Uuid,
    now: DateTime<Utc>,
    ttl: Duration,
) -> Option<LeaseRecord> {
    if existing.owner_token != owner_token || existing.expires_at <= now {
        return None;
    }

    Some(LeaseRecord {
        owner_token,
        expires_at: now + ttl,
    })
}

pub fn can_finalize(existing: Option<&LeaseRecord>, owner_token: Uuid, now: DateTime<Utc>) -> bool {
    lease_owned_by(existing, owner_token, now)
}

#[cfg(test)]
mod tests {
    use super::{LeaseRecord, can_claim, can_finalize, lease_owned_by, renew_lease};
    use chrono::{Duration, Utc};
    use uuid::Uuid;

    #[test]
    fn claim_is_allowed_when_lease_missing_or_expired() {
        let now = Utc::now();
        assert!(can_claim(None, now));

        let old = LeaseRecord {
            owner_token: Uuid::new_v4(),
            expires_at: now - Duration::seconds(1),
        };
        assert!(can_claim(Some(&old), now));

        let active = LeaseRecord {
            owner_token: Uuid::new_v4(),
            expires_at: now + Duration::seconds(30),
        };
        assert!(!can_claim(Some(&active), now));
    }

    #[test]
    fn lease_ownership_requires_same_owner_and_unexpired_lease() {
        let now = Utc::now();
        let owner = Uuid::new_v4();
        let lease = LeaseRecord {
            owner_token: owner,
            expires_at: now + Duration::seconds(10),
        };
        assert!(lease_owned_by(Some(&lease), owner, now));
        assert!(!lease_owned_by(Some(&lease), Uuid::new_v4(), now));
        assert!(!lease_owned_by(
            Some(&lease),
            owner,
            now + Duration::seconds(11)
        ));
    }

    #[test]
    fn renew_and_finalize_are_owner_guarded() {
        let now = Utc::now();
        let owner = Uuid::new_v4();
        let lease = LeaseRecord {
            owner_token: owner,
            expires_at: now + Duration::seconds(3),
        };

        let renewed = renew_lease(&lease, owner, now, Duration::seconds(20)).expect("renewed");
        assert!(renewed.expires_at > lease.expires_at);
        assert!(can_finalize(
            Some(&renewed),
            owner,
            now + Duration::seconds(1)
        ));
        assert!(!can_finalize(
            Some(&renewed),
            Uuid::new_v4(),
            now + Duration::seconds(1)
        ));
    }
}