gevulot_rs/
builders.rs

1use derive_builder::Builder;
2
3use crate::{
4    error::{Error, Result},
5    proto::gevulot::gevulot::{self, InputContext, Label, OutputContext, TaskEnv},
6};
7
8/// Enum representing different units of bytes.
9#[derive(Clone)]
10pub enum ByteUnit {
11    Byte,
12    Kilobyte,
13    Megabyte,
14    Gigabyte,
15}
16
17impl ByteUnit {
18    /// Converts a value in the given ByteUnit to bytes.
19    fn to_bytes(&self, value: u64) -> u64 {
20        match self {
21            ByteUnit::Byte => value,
22            ByteUnit::Kilobyte => value * 1024,
23            ByteUnit::Megabyte => value * 1024 * 1024,
24            ByteUnit::Gigabyte => value * 1024 * 1024 * 1024,
25        }
26    }
27}
28
29/// Struct representing a size in bytes with a specific unit.
30#[derive(Clone)]
31pub struct ByteSize {
32    value: u64,
33    unit: ByteUnit,
34}
35
36impl ByteSize {
37    /// Creates a new ByteSize instance.
38    pub fn new(value: u64, unit: ByteUnit) -> Self {
39        Self { value, unit }
40    }
41
42    /// Converts the ByteSize to bytes.
43    pub fn to_bytes(&self) -> u64 {
44        self.unit.to_bytes(self.value)
45    }
46}
47
48impl std::fmt::Display for ByteSize {
49    /// Formats the ByteSize for display.
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        let unit_str = match self.unit {
52            ByteUnit::Byte => "B",
53            ByteUnit::Kilobyte => "KB",
54            ByteUnit::Megabyte => "MB",
55            ByteUnit::Gigabyte => "GB",
56        };
57        write!(f, "{} {}", self.value, unit_str)
58    }
59}
60
61impl From<(u64, ByteUnit)> for ByteSize {
62    /// Converts a tuple of (u64, ByteUnit) to a ByteSize.
63    fn from(value: (u64, ByteUnit)) -> Self {
64        Self {
65            value: value.0,
66            unit: value.1,
67        }
68    }
69}
70
71#[derive(Builder)]
72pub struct MsgCreateTask {
73    pub creator: String,
74    pub image: String,
75    #[builder(default = "Vec::new()")]
76    pub command: Vec<String>,
77    #[builder(default = "Vec::new()")]
78    pub args: Vec<String>,
79    #[builder(default = "std::collections::HashMap::new()")]
80    pub env: std::collections::HashMap<String, String>,
81    #[builder(default = "std::collections::HashMap::new()")]
82    pub input_contexts: std::collections::HashMap<String, String>,
83    #[builder(default = "Vec::new()")]
84    pub output_contexts: Vec<(String, u64)>,
85    #[builder(default = "1000")]
86    pub cpus: u64,
87    #[builder(default = "0")]
88    pub gpus: u64,
89    #[builder(default = "ByteSize::new(1024, ByteUnit::Megabyte)")]
90    pub memory: ByteSize,
91    #[builder(default = "3600")]
92    pub time: u64,
93    #[builder(default = "true")]
94    pub store_stdout: bool,
95    #[builder(default = "true")]
96    pub store_stderr: bool,
97    #[builder(default = "std::collections::HashMap::new()")]
98    pub labels: std::collections::HashMap<String, String>,
99    #[builder(default = "Vec::new()")]
100    pub tags: Vec<String>,
101}
102
103impl MsgCreateTaskBuilder {
104    pub fn into_message(&self) -> Result<gevulot::MsgCreateTask> {
105        let msg = self
106            .build()
107            .map_err(|e| Error::EncodeError(e.to_string()))?;
108        Ok(gevulot::MsgCreateTask {
109            creator: msg.creator,
110            image: msg.image,
111            command: msg.command,
112            args: msg.args,
113            env: msg
114                .env
115                .into_iter()
116                .map(|(k, v)| TaskEnv { name: k, value: v })
117                .collect(),
118            input_contexts: msg
119                .input_contexts
120                .into_iter()
121                .map(|(k, v)| InputContext {
122                    source: k,
123                    target: v,
124                })
125                .collect(),
126            output_contexts: msg
127                .output_contexts
128                .into_iter()
129                .map(|(source, retention_period)| OutputContext {
130                    source,
131                    retention_period,
132                })
133                .collect(),
134            cpus: msg.cpus,
135            gpus: msg.gpus,
136            memory: msg.memory.to_bytes(),
137            time: msg.time,
138            store_stdout: msg.store_stdout,
139            store_stderr: msg.store_stderr,
140            tags: msg.tags,
141            labels: msg
142                .labels
143                .into_iter()
144                .map(|(k, v)| Label { key: k, value: v })
145                .collect(),
146        })
147    }
148}
149
150#[derive(Builder)]
151pub struct MsgCreatePin {
152    pub creator: String,
153    pub cid: Option<String>,
154    pub bytes: ByteSize,
155    pub name: String,
156    pub redundancy: u64,
157    pub time: u64,
158    pub description: String,
159    pub fallback_urls: Vec<String>,
160    pub tags: Vec<String>,
161    pub labels: Vec<Label>,
162}
163
164impl MsgCreatePinBuilder {
165    pub fn into_message(&self) -> Result<gevulot::MsgCreatePin> {
166        let msg = self
167            .build()
168            .map_err(|e| Error::EncodeError(e.to_string()))?;
169        Ok(gevulot::MsgCreatePin {
170            creator: msg.creator,
171            cid: msg.cid.unwrap_or_default(),
172            bytes: msg.bytes.to_bytes(),
173            name: msg.name,
174            redundancy: msg.redundancy,
175            time: msg.time,
176            description: msg.description,
177            fallback_urls: msg.fallback_urls,
178            tags: msg.tags,
179            labels: msg.labels,
180        })
181    }
182}
183
184#[derive(Builder)]
185pub struct MsgDeletePin {
186    pub creator: String,
187    pub cid: String,
188    pub id: String,
189}
190
191impl MsgDeletePinBuilder {
192    pub fn into_message(&self) -> Result<gevulot::MsgDeletePin> {
193        let msg = self
194            .build()
195            .map_err(|e| Error::EncodeError(e.to_string()))?;
196        Ok(gevulot::MsgDeletePin {
197            creator: msg.creator,
198            cid: msg.cid,
199            id: msg.id,
200        })
201    }
202}
203
204#[derive(Builder)]
205pub struct MsgCreateWorker {
206    pub creator: String,
207    pub name: String,
208    pub description: String,
209    pub cpus: u64,
210    pub gpus: u64,
211    pub memory: ByteSize,
212    pub disk: ByteSize,
213    pub labels: Vec<Label>,
214    pub tags: Vec<String>,
215}
216
217impl MsgCreateWorkerBuilder {
218    pub fn into_message(&self) -> Result<gevulot::MsgCreateWorker> {
219        let msg = self
220            .build()
221            .map_err(|e| Error::EncodeError(e.to_string()))?;
222        Ok(gevulot::MsgCreateWorker {
223            creator: msg.creator,
224            name: msg.name,
225            description: msg.description,
226            cpus: msg.cpus,
227            gpus: msg.gpus,
228            memory: msg.memory.to_bytes(),
229            disk: msg.disk.to_bytes(),
230            labels: msg.labels,
231            tags: msg.tags,
232        })
233    }
234}
235
236#[derive(Builder)]
237pub struct MsgDeleteWorker {
238    pub creator: String,
239    pub id: String,
240}
241
242impl MsgDeleteWorkerBuilder {
243    pub fn into_message(&self) -> Result<gevulot::MsgDeleteWorker> {
244        let msg = self
245            .build()
246            .map_err(|e| Error::EncodeError(e.to_string()))?;
247        Ok(gevulot::MsgDeleteWorker {
248            creator: msg.creator,
249            id: msg.id,
250        })
251    }
252}
253
254#[derive(Builder)]
255pub struct MsgAckPin {
256    pub creator: String,
257    pub cid: String,
258    pub id: String,
259    pub worker_id: String,
260    pub success: bool,
261    pub error: Option<String>,
262}
263
264impl MsgAckPinBuilder {
265    pub fn into_message(&self) -> Result<gevulot::MsgAckPin> {
266        let msg = self
267            .build()
268            .map_err(|e| Error::EncodeError(e.to_string()))?;
269        Ok(gevulot::MsgAckPin {
270            creator: msg.creator,
271            cid: msg.cid,
272            id: msg.id,
273            worker_id: msg.worker_id,
274            success: msg.success,
275            error: msg.error.unwrap_or_default(),
276        })
277    }
278}
279
280#[derive(Builder)]
281pub struct MsgAnnounceWorkerExit {
282    pub creator: String,
283    pub worker_id: String,
284}
285
286impl MsgAnnounceWorkerExitBuilder {
287    pub fn into_message(&self) -> Result<gevulot::MsgAnnounceWorkerExit> {
288        let msg = self
289            .build()
290            .map_err(|e| Error::EncodeError(e.to_string()))?;
291        Ok(gevulot::MsgAnnounceWorkerExit {
292            creator: msg.creator,
293            worker_id: msg.worker_id,
294        })
295    }
296}
297
298#[derive(Builder)]
299pub struct MsgAcceptTask {
300    pub creator: String,
301    pub task_id: String,
302    pub worker_id: String,
303}
304
305impl MsgAcceptTaskBuilder {
306    pub fn into_message(&self) -> Result<gevulot::MsgAcceptTask> {
307        let msg = self
308            .build()
309            .map_err(|e| Error::EncodeError(e.to_string()))?;
310        Ok(gevulot::MsgAcceptTask {
311            creator: msg.creator,
312            task_id: msg.task_id,
313            worker_id: msg.worker_id,
314        })
315    }
316}
317
318#[derive(Builder)]
319pub struct MsgDeclineTask {
320    pub creator: String,
321    pub task_id: String,
322    pub worker_id: String,
323    pub error: Option<String>,
324}
325
326impl MsgDeclineTaskBuilder {
327    pub fn into_message(&self) -> Result<gevulot::MsgDeclineTask> {
328        let msg = self
329            .build()
330            .map_err(|e| Error::EncodeError(e.to_string()))?;
331        Ok(gevulot::MsgDeclineTask {
332            creator: msg.creator,
333            task_id: msg.task_id,
334            worker_id: msg.worker_id,
335            error: msg.error.unwrap_or_default(),
336        })
337    }
338}
339
340#[derive(Builder)]
341pub struct MsgFinishTask {
342    pub creator: String,
343    pub task_id: String,
344    pub exit_code: i32,
345    pub stdout: Option<String>,
346    pub stderr: Option<String>,
347    pub output_contexts: Option<Vec<String>>,
348    pub error: Option<String>,
349}
350
351impl MsgFinishTaskBuilder {
352    pub fn into_message(&self) -> Result<gevulot::MsgFinishTask> {
353        let msg = self
354            .build()
355            .map_err(|e| Error::EncodeError(e.to_string()))?;
356        Ok(gevulot::MsgFinishTask {
357            creator: msg.creator,
358            task_id: msg.task_id,
359            exit_code: msg.exit_code,
360            stdout: msg.stdout.unwrap_or_default(),
361            stderr: msg.stderr.unwrap_or_default(),
362            output_contexts: msg.output_contexts.unwrap_or_default(),
363            error: msg.error.unwrap_or_default(),
364        })
365    }
366}
367
368#[derive(Builder)]
369pub struct MsgSudoDeletePin {
370    pub authority: String,
371    pub cid: String,
372}
373
374impl MsgSudoDeletePinBuilder {
375    pub fn into_message(&self) -> Result<gevulot::MsgSudoDeletePin> {
376        let msg = self
377            .build()
378            .map_err(|e| Error::EncodeError(e.to_string()))?;
379        Ok(gevulot::MsgSudoDeletePin {
380            authority: msg.authority,
381            cid: msg.cid,
382        })
383    }
384}
385
386#[derive(Builder)]
387pub struct MsgSudoDeleteWorker {
388    pub authority: String,
389    pub id: String,
390}
391
392impl MsgSudoDeleteWorkerBuilder {
393    pub fn into_message(&self) -> Result<gevulot::MsgSudoDeleteWorker> {
394        let msg = self
395            .build()
396            .map_err(|e| Error::EncodeError(e.to_string()))?;
397        Ok(gevulot::MsgSudoDeleteWorker {
398            authority: msg.authority,
399            id: msg.id,
400        })
401    }
402}
403
404#[derive(Builder)]
405pub struct MsgSudoDeleteTask {
406    pub authority: String,
407    pub id: String,
408}
409
410impl MsgSudoDeleteTaskBuilder {
411    pub fn into_message(&self) -> Result<gevulot::MsgSudoDeleteTask> {
412        let msg = self
413            .build()
414            .map_err(|e| Error::EncodeError(e.to_string()))?;
415        Ok(gevulot::MsgSudoDeleteTask {
416            authority: msg.authority,
417            id: msg.id,
418        })
419    }
420}
421
422#[derive(Builder)]
423pub struct MsgSudoFreezeAccount {
424    pub authority: String,
425    pub account: String,
426}
427
428impl MsgSudoFreezeAccountBuilder {
429    pub fn into_message(&self) -> Result<gevulot::MsgSudoFreezeAccount> {
430        let msg = self
431            .build()
432            .map_err(|e| Error::EncodeError(e.to_string()))?;
433        Ok(gevulot::MsgSudoFreezeAccount {
434            authority: msg.authority,
435            account: msg.account,
436        })
437    }
438}
439
440#[derive(Builder)]
441pub struct MsgRescheduleTask {
442    pub creator: String,
443    pub task_id: String,
444}
445
446impl MsgRescheduleTaskBuilder {
447    pub fn into_message(&self) -> Result<gevulot::MsgRescheduleTask> {
448        let msg = self
449            .build()
450            .map_err(|e| Error::EncodeError(e.to_string()))?;
451        Ok(gevulot::MsgRescheduleTask {
452            creator: msg.creator,
453            id: msg.task_id,
454        })
455    }
456}
457
458#[derive(Builder)]
459pub struct MsgDeleteTask {
460    pub creator: String,
461    pub id: String,
462}
463
464impl MsgDeleteTaskBuilder {
465    pub fn into_message(&self) -> Result<gevulot::MsgDeleteTask> {
466        let msg = self
467            .build()
468            .map_err(|e| Error::EncodeError(e.to_string()))?;
469        Ok(gevulot::MsgDeleteTask {
470            creator: msg.creator,
471            id: msg.id,
472        })
473    }
474}