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 MsgUpdateWorker {
238    pub creator: String,
239    pub id: String,
240    pub name: String,
241    pub description: String,
242    pub cpus: u64,
243    pub gpus: u64,
244    pub memory: ByteSize,
245    pub disk: ByteSize,
246    pub labels: Vec<Label>,
247    pub tags: Vec<String>,
248}
249
250impl MsgUpdateWorkerBuilder {
251    pub fn into_message(&self) -> Result<gevulot::MsgUpdateWorker> {
252        let msg = self
253            .build()
254            .map_err(|e| Error::EncodeError(e.to_string()))?;
255        Ok(gevulot::MsgUpdateWorker {
256            creator: msg.creator,
257            id: msg.id,
258            name: msg.name,
259            description: msg.description,
260            cpus: msg.cpus,
261            gpus: msg.gpus,
262            memory: msg.memory.to_bytes(),
263            disk: msg.disk.to_bytes(),
264            labels: msg.labels,
265            tags: msg.tags,
266        })
267    }
268}
269
270#[derive(Builder)]
271pub struct MsgDeleteWorker {
272    pub creator: String,
273    pub id: String,
274}
275
276impl MsgDeleteWorkerBuilder {
277    pub fn into_message(&self) -> Result<gevulot::MsgDeleteWorker> {
278        let msg = self
279            .build()
280            .map_err(|e| Error::EncodeError(e.to_string()))?;
281        Ok(gevulot::MsgDeleteWorker {
282            creator: msg.creator,
283            id: msg.id,
284        })
285    }
286}
287
288#[derive(Builder)]
289pub struct MsgAckPin {
290    pub creator: String,
291    pub cid: String,
292    pub id: String,
293    pub worker_id: String,
294    pub success: bool,
295    pub error: Option<String>,
296}
297
298impl MsgAckPinBuilder {
299    pub fn into_message(&self) -> Result<gevulot::MsgAckPin> {
300        let msg = self
301            .build()
302            .map_err(|e| Error::EncodeError(e.to_string()))?;
303        Ok(gevulot::MsgAckPin {
304            creator: msg.creator,
305            cid: msg.cid,
306            id: msg.id,
307            worker_id: msg.worker_id,
308            success: msg.success,
309            error: msg.error.unwrap_or_default(),
310        })
311    }
312}
313
314#[derive(Builder)]
315pub struct MsgAnnounceWorkerExit {
316    pub creator: String,
317    pub worker_id: String,
318}
319
320impl MsgAnnounceWorkerExitBuilder {
321    pub fn into_message(&self) -> Result<gevulot::MsgAnnounceWorkerExit> {
322        let msg = self
323            .build()
324            .map_err(|e| Error::EncodeError(e.to_string()))?;
325        Ok(gevulot::MsgAnnounceWorkerExit {
326            creator: msg.creator,
327            worker_id: msg.worker_id,
328        })
329    }
330}
331
332#[derive(Builder)]
333pub struct MsgAcceptTask {
334    pub creator: String,
335    pub task_id: String,
336    pub worker_id: String,
337}
338
339impl MsgAcceptTaskBuilder {
340    pub fn into_message(&self) -> Result<gevulot::MsgAcceptTask> {
341        let msg = self
342            .build()
343            .map_err(|e| Error::EncodeError(e.to_string()))?;
344        Ok(gevulot::MsgAcceptTask {
345            creator: msg.creator,
346            task_id: msg.task_id,
347            worker_id: msg.worker_id,
348        })
349    }
350}
351
352#[derive(Builder)]
353pub struct MsgDeclineTask {
354    pub creator: String,
355    pub task_id: String,
356    pub worker_id: String,
357    pub error: Option<String>,
358}
359
360impl MsgDeclineTaskBuilder {
361    pub fn into_message(&self) -> Result<gevulot::MsgDeclineTask> {
362        let msg = self
363            .build()
364            .map_err(|e| Error::EncodeError(e.to_string()))?;
365        Ok(gevulot::MsgDeclineTask {
366            creator: msg.creator,
367            task_id: msg.task_id,
368            worker_id: msg.worker_id,
369            error: msg.error.unwrap_or_default(),
370        })
371    }
372}
373
374#[derive(Builder)]
375pub struct MsgFinishTask {
376    pub creator: String,
377    pub task_id: String,
378    pub exit_code: i32,
379    pub stdout: Option<String>,
380    pub stderr: Option<String>,
381    pub output_contexts: Option<Vec<String>>,
382    pub error: Option<String>,
383}
384
385impl MsgFinishTaskBuilder {
386    pub fn into_message(&self) -> Result<gevulot::MsgFinishTask> {
387        let msg = self
388            .build()
389            .map_err(|e| Error::EncodeError(e.to_string()))?;
390        Ok(gevulot::MsgFinishTask {
391            creator: msg.creator,
392            task_id: msg.task_id,
393            exit_code: msg.exit_code,
394            stdout: msg.stdout.unwrap_or_default(),
395            stderr: msg.stderr.unwrap_or_default(),
396            output_contexts: msg.output_contexts.unwrap_or_default(),
397            error: msg.error.unwrap_or_default(),
398        })
399    }
400}
401
402#[derive(Builder)]
403pub struct MsgSudoDeletePin {
404    pub authority: String,
405    pub cid: String,
406}
407
408impl MsgSudoDeletePinBuilder {
409    pub fn into_message(&self) -> Result<gevulot::MsgSudoDeletePin> {
410        let msg = self
411            .build()
412            .map_err(|e| Error::EncodeError(e.to_string()))?;
413        Ok(gevulot::MsgSudoDeletePin {
414            authority: msg.authority,
415            cid: msg.cid,
416        })
417    }
418}
419
420#[derive(Builder)]
421pub struct MsgSudoDeleteWorker {
422    pub authority: String,
423    pub id: String,
424}
425
426impl MsgSudoDeleteWorkerBuilder {
427    pub fn into_message(&self) -> Result<gevulot::MsgSudoDeleteWorker> {
428        let msg = self
429            .build()
430            .map_err(|e| Error::EncodeError(e.to_string()))?;
431        Ok(gevulot::MsgSudoDeleteWorker {
432            authority: msg.authority,
433            id: msg.id,
434        })
435    }
436}
437
438#[derive(Builder)]
439pub struct MsgSudoDeleteTask {
440    pub authority: String,
441    pub id: String,
442}
443
444impl MsgSudoDeleteTaskBuilder {
445    pub fn into_message(&self) -> Result<gevulot::MsgSudoDeleteTask> {
446        let msg = self
447            .build()
448            .map_err(|e| Error::EncodeError(e.to_string()))?;
449        Ok(gevulot::MsgSudoDeleteTask {
450            authority: msg.authority,
451            id: msg.id,
452        })
453    }
454}
455
456#[derive(Builder)]
457pub struct MsgSudoFreezeAccount {
458    pub authority: String,
459    pub account: String,
460}
461
462impl MsgSudoFreezeAccountBuilder {
463    pub fn into_message(&self) -> Result<gevulot::MsgSudoFreezeAccount> {
464        let msg = self
465            .build()
466            .map_err(|e| Error::EncodeError(e.to_string()))?;
467        Ok(gevulot::MsgSudoFreezeAccount {
468            authority: msg.authority,
469            account: msg.account,
470        })
471    }
472}
473
474#[derive(Builder)]
475pub struct MsgRescheduleTask {
476    pub creator: String,
477    pub task_id: String,
478}
479
480impl MsgRescheduleTaskBuilder {
481    pub fn into_message(&self) -> Result<gevulot::MsgRescheduleTask> {
482        let msg = self
483            .build()
484            .map_err(|e| Error::EncodeError(e.to_string()))?;
485        Ok(gevulot::MsgRescheduleTask {
486            creator: msg.creator,
487            id: msg.task_id,
488        })
489    }
490}
491
492#[derive(Builder)]
493pub struct MsgDeleteTask {
494    pub creator: String,
495    pub id: String,
496}
497
498impl MsgDeleteTaskBuilder {
499    pub fn into_message(&self) -> Result<gevulot::MsgDeleteTask> {
500        let msg = self
501            .build()
502            .map_err(|e| Error::EncodeError(e.to_string()))?;
503        Ok(gevulot::MsgDeleteTask {
504            creator: msg.creator,
505            id: msg.id,
506        })
507    }
508}