pg_blast_radius/
locks.rs

1use serde::{Deserialize, Serialize};
2
3use crate::types::LockMode;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
6#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
7pub enum DmlKind {
8    #[default]
9    Select,
10    SelectForUpdate,
11    Insert,
12    Update,
13    Delete,
14}
15
16impl DmlKind {
17    pub fn lock_mode(self) -> LockMode {
18        match self {
19            Self::Select => LockMode::AccessShare,
20            Self::SelectForUpdate => LockMode::RowShare,
21            Self::Insert | Self::Update | Self::Delete => LockMode::RowExclusive,
22        }
23    }
24}
25
26impl std::fmt::Display for DmlKind {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        match self {
29            Self::Select => write!(f, "SELECT"),
30            Self::SelectForUpdate => write!(f, "SELECT FOR UPDATE"),
31            Self::Insert => write!(f, "INSERT"),
32            Self::Update => write!(f, "UPDATE"),
33            Self::Delete => write!(f, "DELETE"),
34        }
35    }
36}
37
38const CONFLICT: [[bool; 8]; 8] = [
39    //  AS     RS     RE     SUE    S      SRE    E      AE
40    [false, false, false, false, false, false, false, true ],  // AccessShare
41    [false, false, false, false, false, false, true,  true ],  // RowShare
42    [false, false, false, false, true,  true,  true,  true ],  // RowExclusive
43    [false, false, false, true,  true,  true,  true,  true ],  // ShareUpdateExclusive
44    [false, false, true,  true,  false, true,  true,  true ],  // Share
45    [false, false, true,  true,  true,  true,  true,  true ],  // ShareRowExclusive
46    [false, true,  true,  true,  true,  true,  true,  true ],  // Exclusive
47    [true,  true,  true,  true,  true,  true,  true,  true ],  // AccessExclusive
48];
49
50fn lock_index(mode: LockMode) -> usize {
51    match mode {
52        LockMode::AccessShare => 0,
53        LockMode::RowShare => 1,
54        LockMode::RowExclusive => 2,
55        LockMode::ShareUpdateExclusive => 3,
56        LockMode::Share => 4,
57        LockMode::ShareRowExclusive => 5,
58        LockMode::Exclusive => 6,
59        LockMode::AccessExclusive => 7,
60    }
61}
62
63pub fn conflicts(requested: LockMode, held: LockMode) -> bool {
64    CONFLICT[lock_index(requested)][lock_index(held)]
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn access_share_only_conflicts_with_access_exclusive() {
73        assert!(!conflicts(LockMode::AccessShare, LockMode::AccessShare));
74        assert!(!conflicts(LockMode::AccessShare, LockMode::RowShare));
75        assert!(!conflicts(LockMode::AccessShare, LockMode::RowExclusive));
76        assert!(!conflicts(LockMode::AccessShare, LockMode::ShareUpdateExclusive));
77        assert!(!conflicts(LockMode::AccessShare, LockMode::Share));
78        assert!(!conflicts(LockMode::AccessShare, LockMode::ShareRowExclusive));
79        assert!(!conflicts(LockMode::AccessShare, LockMode::Exclusive));
80        assert!(conflicts(LockMode::AccessShare, LockMode::AccessExclusive));
81    }
82
83    #[test]
84    fn access_exclusive_conflicts_with_everything() {
85        assert!(conflicts(LockMode::AccessExclusive, LockMode::AccessShare));
86        assert!(conflicts(LockMode::AccessExclusive, LockMode::RowShare));
87        assert!(conflicts(LockMode::AccessExclusive, LockMode::RowExclusive));
88        assert!(conflicts(LockMode::AccessExclusive, LockMode::ShareUpdateExclusive));
89        assert!(conflicts(LockMode::AccessExclusive, LockMode::Share));
90        assert!(conflicts(LockMode::AccessExclusive, LockMode::ShareRowExclusive));
91        assert!(conflicts(LockMode::AccessExclusive, LockMode::Exclusive));
92        assert!(conflicts(LockMode::AccessExclusive, LockMode::AccessExclusive));
93    }
94
95    #[test]
96    fn matrix_is_symmetric() {
97        let modes = [
98            LockMode::AccessShare,
99            LockMode::RowShare,
100            LockMode::RowExclusive,
101            LockMode::ShareUpdateExclusive,
102            LockMode::Share,
103            LockMode::ShareRowExclusive,
104            LockMode::Exclusive,
105            LockMode::AccessExclusive,
106        ];
107        for &a in &modes {
108            for &b in &modes {
109                assert_eq!(
110                    conflicts(a, b),
111                    conflicts(b, a),
112                    "conflict matrix is not symmetric for {a} vs {b}"
113                );
114            }
115        }
116    }
117
118    #[test]
119    fn row_exclusive_conflicts_with_share() {
120        assert!(conflicts(LockMode::RowExclusive, LockMode::Share));
121        assert!(conflicts(LockMode::Share, LockMode::RowExclusive));
122    }
123
124    #[test]
125    fn share_does_not_conflict_with_share() {
126        assert!(!conflicts(LockMode::Share, LockMode::Share));
127    }
128
129    #[test]
130    fn dml_lock_modes() {
131        assert_eq!(DmlKind::Select.lock_mode(), LockMode::AccessShare);
132        assert_eq!(DmlKind::SelectForUpdate.lock_mode(), LockMode::RowShare);
133        assert_eq!(DmlKind::Insert.lock_mode(), LockMode::RowExclusive);
134        assert_eq!(DmlKind::Update.lock_mode(), LockMode::RowExclusive);
135        assert_eq!(DmlKind::Delete.lock_mode(), LockMode::RowExclusive);
136    }
137}
pg_blast_radius/locks.rs

pg_blast_radius/
locks.rs
