ts_pd/
status.rs

1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! A room's transfer state, as the PD and EPD peers' status messages carry
22//! it.
23
24use std::collections::HashMap;
25
26/// A transfer's state, in the order it advances. The values are the wire
27/// encoding of the P/D status messages and of the replica's MIN agreement.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
29pub enum TransferPoll {
30    /// The transfer failed; sticky.
31    Failed = 0,
32    /// The peers are connecting.
33    Bootstrapping = 1,
34    /// The peers are connected.
35    Bootstrapped = 2,
36    /// Decode: the manifest went out; waiting for the KV.
37    WaitingForInput = 3,
38    /// Unused: a transfer is not observed mid-flight.
39    Transferring = 4,
40    /// The KV arrived (or left, on Prefill).
41    Success = 5,
42}
43
44impl TransferPoll {
45    /// The state a wire value names.
46    pub fn from_wire(value: i64) -> Option<Self> {
47        Some(match value {
48            0 => Self::Failed,
49            1 => Self::Bootstrapping,
50            2 => Self::Bootstrapped,
51            3 => Self::WaitingForInput,
52            4 => Self::Transferring,
53            5 => Self::Success,
54            _ => return None,
55        })
56    }
57}
58
59/// Room states and failure reasons.
60#[derive(Debug, Default)]
61pub struct RoomTable {
62    states: HashMap<u64, TransferPoll>,
63    failures: HashMap<u64, String>,
64}
65
66impl RoomTable {
67    /// A room's state, if the room is known.
68    pub fn get(&self, room: u64) -> Option<TransferPoll> {
69        self.states.get(&room).copied()
70    }
71
72    /// Advance a room: states only move forward, and `Failed` is sticky, so
73    /// a late success cannot resurrect a broken transfer.
74    pub fn update(&mut self, room: u64, status: TransferPoll) {
75        let entry = self.states.entry(room).or_insert(status);
76        *entry = if *entry == TransferPoll::Failed || status == TransferPoll::Failed {
77            TransferPoll::Failed
78        } else {
79            (*entry).max(status)
80        };
81    }
82
83    /// Record why a room failed; a later reason replaces an earlier one.
84    pub fn record_failure(&mut self, room: u64, reason: impl Into<String>) {
85        self.failures.insert(room, reason.into());
86    }
87
88    /// Whether a failure reason is on record.
89    pub fn has_failure(&self, room: u64) -> bool {
90        self.failures.contains_key(&room)
91    }
92
93    /// The failure reason on record.
94    pub fn failure(&self, room: u64) -> Option<&str> {
95        self.failures.get(&room).map(String::as_str)
96    }
97
98    /// Forget a room.
99    pub fn remove(&mut self, room: u64) {
100        self.states.remove(&room);
101        self.failures.remove(&room);
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn states_advance_and_failure_sticks() {
111        let mut table = RoomTable::default();
112        table.update(1, TransferPoll::Bootstrapped);
113        table.update(1, TransferPoll::Bootstrapping);
114        assert_eq!(table.get(1), Some(TransferPoll::Bootstrapped));
115        table.update(1, TransferPoll::Failed);
116        table.update(1, TransferPoll::Success);
117        assert_eq!(table.get(1), Some(TransferPoll::Failed));
118        assert_eq!(table.get(2), None);
119        assert_eq!(TransferPoll::from_wire(5), Some(TransferPoll::Success));
120        assert_eq!(TransferPoll::from_wire(9), None);
121    }
122}
ts_pd/status.rs

ts_pd/
status.rs
