Skip to main content

orch/cluster/
command.rs

1use serde::{Deserialize, Serialize};
2use sqlx::SqliteConnection;
3
4use crate::store::control::MemoryEntry;
5use crate::store::dependencies::{self, NewDependency};
6use crate::store::jobs::{self, JobState};
7use crate::store::kv;
8use crate::store::tasks::{apply, NewTask, StoreRouteCmd, TaskStatus, UpsertExternalCmd};
9
10type Updates = Vec<(String, serde_json::Value)>;
11
12/// One replicated write. The proposer stamps `now`, so every node applies the same value
13#[derive(Serialize, Deserialize, Clone, Debug)]
14pub struct Request {
15    pub now: String,
16    pub cmd: Command,
17}
18
19#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
20#[serde(tag = "op", rename_all = "snake_case")]
21pub enum Command {
22    KvSet {
23        key: String,
24        value: String,
25    },
26    KvIncrement {
27        key: String,
28    },
29    KvInsertIfAbsent {
30        key: String,
31        value: String,
32    },
33    KvDelete {
34        key: String,
35    },
36    TaskCreate {
37        new: NewTask,
38        internal: bool,
39    },
40    TaskUpsertExternal(UpsertExternalCmd),
41    TaskUpdateStatusAndFields {
42        id: i64,
43        status: TaskStatus,
44        updates: Updates,
45    },
46    TaskUpdateStatus {
47        id: i64,
48        status: TaskStatus,
49        expected: Option<TaskStatus>,
50    },
51    TaskUpdateExternalId {
52        id: i64,
53        external_id: String,
54    },
55    TaskTouch {
56        id: i64,
57    },
58    TaskResetToNew {
59        id: i64,
60    },
61    TaskSetFields {
62        id: i64,
63        updates: Updates,
64        silent: bool,
65    },
66    TaskIncrement {
67        id: i64,
68        field: String,
69    },
70    TaskResetCounters {
71        id: i64,
72        failure_only: bool,
73    },
74    TaskStoreRoute(StoreRouteCmd),
75    TaskSetEstimateIfZero {
76        id: i64,
77        estimate: u8,
78    },
79    TaskAppendMemory {
80        id: i64,
81        entry: MemoryEntry,
82    },
83    TaskMarkCleaned {
84        id: i64,
85    },
86    MoveRepo {
87        old: String,
88        new: String,
89    },
90    DepReplace {
91        task_id: i64,
92        rows: Vec<NewDependency>,
93    },
94    DepLinkKnown {
95        task_id: i64,
96        repo: String,
97        number: i64,
98    },
99    DepBlockClaim {
100        task_id: i64,
101    },
102    JobStateUpsert(JobState),
103    ChannelSubscribe {
104        channel: String,
105        thread_id: String,
106        repo: String,
107        topic_id: String,
108    },
109    ChannelUnsubscribe {
110        channel: String,
111        thread_id: String,
112        repo: String,
113    },
114}
115
116/// A failed command is `Err` on every node, so one bad command cannot stop the state machine
117#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
118pub enum Reply {
119    Done,
120    Bool(bool),
121    Int(i64),
122    Count(u64),
123    Value(String),
124    Unmet(Option<Vec<String>>),
125    Err(String),
126}
127
128impl Command {
129    /// A write through the log may not touch this node's own state: a local kv key, such as its
130    /// Raft position, or a local task column, such as the worktree path that cleanup removes
131    pub(crate) fn check_replicable(&self) -> anyhow::Result<()> {
132        let key = match self {
133            Command::KvSet { key, .. }
134            | Command::KvIncrement { key }
135            | Command::KvInsertIfAbsent { key, .. }
136            | Command::KvDelete { key } => Some(key),
137            _ => None,
138        };
139        if let Some(key) = key {
140            anyhow::ensure!(
141                !kv::is_local(key),
142                "{key} is local to each node, never replicated"
143            );
144        }
145
146        let updates = match self {
147            Command::TaskSetFields { updates, .. }
148            | Command::TaskUpdateStatusAndFields { updates, .. } => &updates[..],
149            _ => &[],
150        };
151        if let Some((col, _)) = updates
152            .iter()
153            .find(|(col, _)| crate::store::tasks::is_local_column(col))
154        {
155            anyhow::bail!("{col} is local to each node, never replicated");
156        }
157
158        Ok(())
159    }
160
161    /// `replicated` is true when Raft applies the command, so new ids come from the replicated
162    /// rows. A direct write keeps SQLite's AUTOINCREMENT, as orch on one machine always has
163    pub(crate) async fn apply(
164        self,
165        conn: &mut SqliteConnection,
166        now: &str,
167        replicated: bool,
168    ) -> anyhow::Result<Reply> {
169        if replicated {
170            self.check_replicable()?;
171        }
172
173        Ok(match self {
174            Command::KvSet { key, value } => {
175                kv::set(conn, now, &key, &value).await?;
176                Reply::Done
177            }
178            Command::KvIncrement { key } => Reply::Count(kv::increment(conn, now, &key).await?),
179            Command::KvInsertIfAbsent { key, value } => {
180                Reply::Value(kv::insert_if_absent(conn, now, &key, &value).await?)
181            }
182            Command::KvDelete { key } => {
183                kv::delete(conn, &key).await?;
184                Reply::Done
185            }
186            Command::TaskCreate { new, internal } => {
187                Reply::Int(apply::create(conn, now, &new, internal, replicated).await?)
188            }
189            Command::TaskUpsertExternal(ext) => {
190                Reply::Int(apply::upsert_external(conn, now, &ext, replicated).await?)
191            }
192            Command::TaskUpdateStatusAndFields {
193                id,
194                status,
195                updates,
196            } => {
197                apply::update_status_and_fields(conn, now, id, status, &updates).await?;
198                Reply::Done
199            }
200            Command::TaskUpdateStatus {
201                id,
202                status,
203                expected,
204            } => Reply::Bool(apply::update_status(conn, now, id, status, expected).await?),
205            Command::TaskUpdateExternalId { id, external_id } => {
206                apply::update_external_id(conn, now, id, &external_id).await?;
207                Reply::Done
208            }
209            Command::TaskTouch { id } => {
210                apply::touch(conn, now, id).await?;
211                Reply::Done
212            }
213            Command::TaskResetToNew { id } => {
214                apply::reset_to_new(conn, now, id).await?;
215                Reply::Done
216            }
217            Command::TaskSetFields {
218                id,
219                updates,
220                silent,
221            } => {
222                apply::set_fields(conn, now, id, &updates, silent).await?;
223                Reply::Done
224            }
225            Command::TaskIncrement { id, field } => {
226                Reply::Int(apply::increment(conn, now, id, &field).await?)
227            }
228            Command::TaskResetCounters { id, failure_only } => {
229                apply::reset_counters(conn, now, id, failure_only).await?;
230                Reply::Done
231            }
232            Command::TaskStoreRoute(route) => {
233                apply::store_route(conn, now, &route).await?;
234                Reply::Done
235            }
236            Command::TaskSetEstimateIfZero { id, estimate } => {
237                Reply::Bool(apply::set_estimate_if_zero(conn, now, id, estimate).await?)
238            }
239            Command::TaskAppendMemory { id, entry } => {
240                apply::append_memory(conn, now, id, &entry).await?;
241                Reply::Done
242            }
243            Command::TaskMarkCleaned { id } => {
244                apply::mark_cleaned(conn, now, id).await?;
245                Reply::Done
246            }
247            Command::MoveRepo { old, new } => {
248                let moved = dependencies::apply::move_repo(conn, &old, &new).await?;
249                Reply::Int(i64::try_from(moved)?)
250            }
251            Command::DepReplace { task_id, rows } => {
252                dependencies::apply::replace_dependencies(conn, task_id, &rows).await?;
253                Reply::Done
254            }
255            Command::DepLinkKnown {
256                task_id,
257                repo,
258                number,
259            } => {
260                dependencies::apply::link_known_task(conn, task_id, &repo, number).await?;
261                Reply::Done
262            }
263            Command::DepBlockClaim { task_id } => {
264                Reply::Unmet(dependencies::apply::block_claim(conn, now, task_id).await?)
265            }
266            Command::JobStateUpsert(state) => {
267                jobs::apply::upsert_job_state(conn, &state).await?;
268                Reply::Done
269            }
270            Command::ChannelSubscribe {
271                channel,
272                thread_id,
273                repo,
274                topic_id,
275            } => {
276                jobs::apply::subscribe(conn, &channel, &thread_id, &repo, &topic_id).await?;
277                Reply::Done
278            }
279            Command::ChannelUnsubscribe {
280                channel,
281                thread_id,
282                repo,
283            } => {
284                jobs::apply::unsubscribe(conn, &channel, &thread_id, &repo).await?;
285                Reply::Done
286            }
287        })
288    }
289}