1use futures::channel::mpsc::{self, Receiver, Sender};
2use 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#[derive(Debug)]
13pub enum CommitEvent {
14 AutoStart((Sender<CommitMessage>, QueryId)),
16 CustomStart((Sender<CommitMessage>, QueryId)),
18 SendRequests((Vec<CommitMessage>, QueryId)),
20 Done,
23}
24
25#[derive(Debug)]
27pub enum CommitMessage {
28 Send(CommitRequestParams),
30 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#[derive(Debug)]
47pub enum Commit {
48 No,
49 Auto(Progress),
50 Custom(Progress),
53}
54
55#[derive(Debug, Default)]
56pub enum Progress {
57 #[default]
58 BeforeStart,
59 Sending((Receiver<CommitMessage>, BTreeSet<Tid>)),
62 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 }
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}