Skip to main content

dht_rpc/
commit.rs

1use futures::channel::mpsc::{self, Receiver, Sender};
2/// Commit
3/// * commit request - An outgoing [`ReplyMsgData`] which includes a valid token. When the message
4///   is recieved and verified, the reciever does some mutation of state
5use std::{collections::BTreeSet, iter::FromIterator, net::SocketAddr};
6
7use crate::{constants::DEFAULT_COMMIT_CHANNEL_SIZE, io::Tid};
8
9use crate::{Command, IdBytes, query::QueryId};
10
11/// Emitted from a query's commit
12#[derive(Debug)]
13pub enum CommitEvent {
14    // Emitted when commit process starts for Commit::Auto. Progress is Sending
15    AutoStart((Sender<CommitMessage>, QueryId)),
16    // emitted when commit process starts for custom. Progess is set to Sending
17    CustomStart((Sender<CommitMessage>, QueryId)),
18    // Emitted when commit process has rquests it wants to send. Send them.
19    SendRequests((Vec<CommitMessage>, QueryId)),
20    /// Commit Done
21    /// TODO add info about commit, like successful replies, timeouts, etc
22    Done,
23}
24
25/// Sent by user with Commit::Custom to send commit requests
26#[derive(Debug)]
27pub enum CommitMessage {
28    /// User emit this when they want to send a commit request.
29    Send(CommitRequestParams),
30    /// User emits this when they've push all commit requests
31    /// Casuse  Progress to transition to AwaitingReplies
32    Done,
33}
34
35#[derive(Debug)]
36pub struct CommitRequestParams {
37    pub command: Command,
38    pub target: Option<IdBytes>,
39    pub value: Option<Vec<u8>>,
40    pub peer: SocketAddr,
41    pub query_id: QueryId,
42    pub token: [u8; 32],
43}
44
45/// The kinds of [`Commit`] a Query can have.
46#[derive(Debug)]
47pub enum Commit {
48    No,
49    Auto(Progress),
50    //Custom((Option<Receiver<CommitMessages>>, Progress)),
51    //Custom(Custom),
52    Custom(Progress),
53}
54
55#[derive(Debug, Default)]
56pub enum Progress {
57    #[default]
58    BeforeStart,
59    /// Sending commit requests in progress. tids are added as they are sent, but also can be
60    /// removed when responses are recieved
61    Sending((Receiver<CommitMessage>, BTreeSet<Tid>)),
62    /// awaiting replies to commit messages
63    AwaitingReplies(BTreeSet<Tid>),
64    Done,
65}
66
67use Progress as P;
68impl Progress {
69    pub fn start_sending(&mut self) -> Sender<CommitMessage> {
70        if matches!(self, P::BeforeStart) {
71            let (tx, rx) = mpsc::channel(DEFAULT_COMMIT_CHANNEL_SIZE);
72            *self = P::Sending((rx, Default::default()));
73            tx
74        } else {
75            panic!("Tried to start sending but already started");
76        }
77    }
78
79    pub fn all_replies_recieved(&self) -> bool {
80        match self {
81            P::BeforeStart => false,
82            P::Sending(_) => false,
83            P::AwaitingReplies(tids) => tids.is_empty(),
84            P::Done => true,
85        }
86    }
87    pub fn transition_to_awaiting(&mut self) {
88        *self = P::AwaitingReplies(match self {
89            P::Sending((_, tids)) => tids.clone(),
90            _ => panic!("not in sending"),
91        })
92    }
93
94    pub fn poll(&mut self) -> Option<CommitMessage> {
95        let P::Sending((rx, _tids)) = self else {
96            panic!("poll while not sending");
97        };
98        rx.try_next().ok().flatten()
99    }
100    pub fn sent_tid(&mut self, tid: Tid) -> bool {
101        match self {
102            P::Sending((_, tids)) => tids.insert(tid),
103            _ => panic!("only call while `Sending`"),
104        }
105        // insert to Sending.tids
106    }
107
108    pub fn start_awaiting(&mut self, tids: Vec<Tid>) {
109        if matches!(self, P::BeforeStart | P::Sending(_)) {
110            *self = P::AwaitingReplies(BTreeSet::from_iter(tids))
111        } else {
112            panic!("Tried to start commit that was already started");
113        }
114    }
115    pub fn recieved_tid(&mut self, tid: Tid) -> bool {
116        let done = match self {
117            Self::AwaitingReplies(tids) => {
118                tids.remove(&tid);
119                tids.is_empty()
120            }
121            Self::Sending((_, tids)) => {
122                tids.remove(&tid);
123                false
124            }
125            _ => false,
126        };
127        if done {
128            *self = P::Done;
129        }
130        matches!(self, P::Done)
131    }
132}