1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

/* -------------------------------------------------------------------------- *\
 *        Apache 2.0 License Copyright © 2022-2023 The Aurae Authors          *
 *                                                                            *
 *                +--------------------------------------------+              *
 *                |   █████╗ ██╗   ██╗██████╗  █████╗ ███████╗ |              *
 *                |  ██╔══██╗██║   ██║██╔══██╗██╔══██╗██╔════╝ |              *
 *                |  ███████║██║   ██║██████╔╝███████║█████╗   |              *
 *                |  ██╔══██║██║   ██║██╔══██╗██╔══██║██╔══╝   |              *
 *                |  ██║  ██║╚██████╔╝██║  ██║██║  ██║███████╗ |              *
 *                |  ╚═╝  ╚═╝ ╚═════╝ ╚═╝  ╚═╝╚═╝  ╚═╝╚══════╝ |              *
 *                +--------------------------------------------+              *
 *                                                                            *
 *                         Distributed Systems Runtime                        *
 *                                                                            *
 * -------------------------------------------------------------------------- *
 *                                                                            *
 *   Licensed under the Apache License, Version 2.0 (the "License");          *
 *   you may not use this file except in compliance with the License.         *
 *   You may obtain a copy of the License at                                  *
 *                                                                            *
 *       http://www.apache.org/licenses/LICENSE-2.0                           *
 *                                                                            *
 *   Unless required by applicable law or agreed to in writing, software      *
 *   distributed under the License is distributed on an "AS IS" BASIS,        *
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. *
 *   See the License for the specific language governing permissions and      *
 *   limitations under the License.                                           *
 *                                                                            *
\* -------------------------------------------------------------------------- */

use super::{
    Executable, ExecutableName, ExecutableSpec, ExecutablesError, Result,
};
use std::{collections::HashMap, process::ExitStatus};

type Cache = HashMap<ExecutableName, Executable>;

/// An in-memory store for the list of executables created with Aurae.
#[derive(Debug, Default)]
pub struct Executables {
    cache: Cache,
}

impl Executables {
    pub fn start<T: Into<ExecutableSpec>>(
        &mut self,
        executable_spec: T,
    ) -> Result<&Executable> {
        let executable_spec = executable_spec.into();

        // TODO: replace with try_insert when it becomes stable
        // Check if there was already an executable with the same name.
        if self.cache.contains_key(&executable_spec.name) {
            return Err(ExecutablesError::ExecutableExists {
                executable_name: executable_spec.name,
            });
        }

        let executable_name = executable_spec.name.clone();
        // `or_insert` will always insert as we've already assured ourselves that the key does not exist.
        let executable = self
            .cache
            .entry(executable_name.clone())
            .or_insert_with(|| Executable::new(executable_spec));

        // TODO: if we fail to start, the exe remains in the cache and start cannot be called again
        // solving ^^ was a borrow checker fight and I (future-highway) lost this round.
        executable.start().map_err(|e| {
            ExecutablesError::FailedToStartExecutable {
                executable_name,
                source: e,
            }
        })?;

        Ok(executable)
    }

    pub fn get(&self, executable_name: &ExecutableName) -> Result<&Executable> {
        let Some(executable) = self.cache.get(executable_name) else {
            return Err(ExecutablesError::ExecutableNotFound { executable_name: executable_name.clone() });
        };
        Ok(executable)
    }

    pub async fn stop(
        &mut self,
        executable_name: &ExecutableName,
    ) -> Result<ExitStatus> {
        let Some(executable) = self.cache.get_mut(executable_name) else {
            return Err(ExecutablesError::ExecutableNotFound { executable_name: executable_name.clone() });
        };

        let exit_status = executable.kill().await.map_err(|e| {
            ExecutablesError::FailedToStopExecutable {
                executable_name: executable_name.clone(),
                source: e,
            }
        })?;

        let Some(exit_status) = exit_status else {
            // Exes that never started return None
            let executable = self.cache.remove(executable_name).expect("exe in cache");
            return Err(ExecutablesError::ExecutableNotFound {
                executable_name: executable.name,
            });
        };

        let _ = self.cache.remove(executable_name).ok_or_else(|| {
            // get_mut would have already thrown this error, so we should never reach here
            ExecutablesError::ExecutableNotFound {
                executable_name: executable_name.clone(),
            }
        })?;

        Ok(exit_status)
    }

    /// Stops all executables concurrently
    pub async fn broadcast_stop(&mut self) {
        let mut names = vec![];
        for exe in self.cache.values_mut() {
            let _ = exe.kill().await;
            names.push(exe.name.clone())
        }

        for name in names {
            let _ = self.cache.remove(&name);
        }
    }
}