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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
use chrono;
use std::borrow::Borrow;
use std::collections::HashMap;
use std::os::unix::net::UnixStream;
use std::path::PathBuf;
use std::{env, fs};
use super::*;
pub const STAGING_DIR: &str = "staging";
pub const TASKS_DIR: &str = "tasks";
/// Alias type for a executor task function with arbitrary number of inputs
/// and outputs.
pub type TaskFn = Fn(&mut Context, &[DataInstance], &mut [Output]) -> TaskResult<()> + Send + Sync;
//pub trait TaskFn: Fn(&mut Context, &[DataInstance], &mut [Output]) -> TaskResult<()> + Send + Sync {}
/// The executor event loop and the set of registered tasks.
pub struct Executor {
/// An identifier for the local governor
executor_id: ExecutorId,
/// Any name given to the executor type (denoting the group of provided tasks)
executor_type: String,
/// Path to the socket (should be absolute)
socket_path: PathBuf,
/// Registered task functions
tasks: HashMap<String, Box<TaskFn>>,
/// Executor working directory
working_dir: PathBuf,
/// Executor staging subdirectory
staging_dir: PathBuf,
/// Executor subdirectory with individual tasks
tasks_dir: PathBuf,
/// Prevent running `Executor::run` repeatedly
was_run: bool,
/// If true, failed task directories (but not outputs) are kept in "tasks/"
keep_failed_tasks: bool,
}
impl Executor {
/// Create a executor based on env variables `RAIN_EXECUTOR_ID` and `RAIN_EXECUTOR_SOCKET`
/// and working dir in the current directory. See also `Executor::with_params`.
///
/// Panics when either env variable is missing or invalid.
pub fn new(executor_type: &str) -> Self {
let id = env::var("RAIN_EXECUTOR_ID")
.expect("Env variable RAIN_EXECUTOR_ID required")
.parse()
.expect("Error parsing RAIN_EXECUTOR_ID");
let socket: PathBuf = env::var_os("RAIN_EXECUTOR_SOCKET")
.expect("Env variable RAIN_EXECUTOR_SOCKET required")
.into();
let workdir = env::current_dir().unwrap();
Executor::with_params(executor_type, id, &socket, &workdir)
}
/// Creates a Sbgovernor with the given attributes. Note that the attributes are only
/// recorded at this point and no initialization is performed.
pub fn with_params(
executor_type: &str,
executor_id: ExecutorId,
socket_path: &Path,
working_dir: &Path,
) -> Self {
Executor {
executor_type: executor_type.into(),
executor_id,
socket_path: socket_path.into(),
tasks: HashMap::new(),
staging_dir: working_dir.join(STAGING_DIR),
tasks_dir: working_dir.join(TASKS_DIR),
working_dir: working_dir.into(),
was_run: false,
keep_failed_tasks: false,
}
}
/// Register task function.
///
/// The provided function must accept a list of inpts and outputs, expanding them manually.
/// For functions accepting individual inputs and outputs, see the macro `register_task!`.
/// The accepted function type is equivalent to `'static + TaskFn`.
///
/// Panics when a task with the same name has been registered previously.
pub fn register_task<S, F>(&mut self, task_name: S, task_fun: F)
where
S: Into<String>,
F: 'static
+ Fn(&mut Context, &[DataInstance], &mut [Output]) -> TaskResult<()>
+ Send
+ Sync,
{
let key: String = task_name.into();
if self.tasks.contains_key(&key) {
panic!("can't add task named {:?}: already present", &key);
}
self.tasks.insert(key, Box::new(task_fun));
}
/// Run the executor loop, connecting to the governor, registering and handling requests
/// until the connection is closed. May be only called once.
///
/// Panics on any error outside of the task functions. See `README.md` for rationale.
pub fn run(&mut self) {
if self.was_run {
panic!("Executor::run may only be ran once");
}
self.was_run = true;
// Prepare the directories
if self.staging_dir.exists() || self.tasks_dir.exists() {
panic!(
"Executor must be ran in a clean directory (workdir: {:?})",
self.working_dir
);
}
fs::create_dir(&self.staging_dir).expect("error creating staging dir");
fs::create_dir(&self.tasks_dir).expect("error creating tasks dir");
// Connect to socket
info!(
"Connecting to governor at socket {:?} with ID {}",
self.socket_path, self.executor_id
);
if !self.socket_path.exists() {
panic!("Socket file not found at {:?}", self.socket_path);
}
// Change directory to prevent too long socket pathnames
env::set_current_dir(self.socket_path.parent().unwrap())
.expect("error chdir to socket dir");
let mut sock = UnixStream::connect(&self.socket_path.file_name().unwrap())
.expect("error opening socket");
env::set_current_dir(&self.working_dir).expect("error chdir to working dir");
// Register and run the task loop, catching any errors
let res = (|| {
let regmsg = ExecutorToGovernorMessage::Register(self.register());
sock.write_msg(®msg)?;
loop {
match sock.read_msg()? {
GovernorToExecutorMessage::Call(call_msg) => {
let reply = self.handle_call(call_msg);
sock.write_msg(&ExecutorToGovernorMessage::Result(reply))?;
}
GovernorToExecutorMessage::DropCached(drop_msg) => {
if !drop_msg.objects.is_empty() {
panic!("received nonempty dropCached request with no cached objects");
}
}
}
}
})();
match res {
Err(Error(ErrorKind::Io(ref e), _))
if (e.kind() == io::ErrorKind::ConnectionAborted)
|| (e.kind() == io::ErrorKind::UnexpectedEof) =>
{
info!("Connection closed, shutting down");
}
other => other.expect("error in the message loop"),
}
}
/// Create a register message to the governor.
fn register(&mut self) -> RegisterMsg {
RegisterMsg {
protocol: MSG_PROTOCOL.into(),
executor_id: self.executor_id,
executor_type: self.executor_type.clone(),
}
}
/// Handle one call msg: decode, run the task function, cleanup finished task
/// (and already staged files on failure), create reply message.
///
/// Panics on any IO error. TaskErrors are handled and returned to the governor.
fn handle_call(&mut self, call_msg: CallMsg) -> ResultMsg {
let task_name = format!(
"{}-task-{}_{}",
chrono::Local::now().format("%Y%m%d-%H%M%S"),
call_msg.spec.id.get_session_id(),
call_msg.spec.id.get_id()
);
let task_dir = self.tasks_dir.join(task_name);
let (task_exec, task_method): (String, String) = {
let mut m_split = call_msg.spec.task_type.splitn(2, "/");
(
m_split.next().unwrap().into(),
m_split
.next()
.expect("Call method must be in the form \"executor_type/method\"")
.into(),
)
};
let mut context = Context::for_call_msg(call_msg, &self.staging_dir, &task_dir);
if task_exec != self.executor_type {
context.fail(format!(
"Mismatch of executor type in call: {:?} vs {:?}",
task_exec, self.executor_type
))
}
match self.tasks.get(&task_method) {
None => context.fail(format!(
"Task {:?} not found in executor {:?}",
task_method, self.executor_type
)),
Some(f) => {
fs::create_dir(&task_dir).expect("error creating task dir");
// Call the method function with context
let res = context.call_with_context(f.borrow());
// Check and handle in-task errors
env::set_current_dir(&self.working_dir).expect("error on chdir to work dir");
if let Err(ref e) = res {
debug!("Method {:?} in {:?} failed: {}", task_method, task_dir, e);
context.fail(format!(
"error returned from call to {:?} (in {:?}):\n{}",
task_method, task_dir, e
));
// Clean already staged/open outputs
for mut o in context.outputs.iter_mut() {
o.cleanup_failed_task();
}
if !self.keep_failed_tasks {
// cleanup of the task working dir
fs::remove_dir_all(task_dir)
.expect("error removing the task direcotry after failure");
}
} else {
debug!("Method {:?} finished successfully", task_method);
// cleanup of the task working dir
fs::remove_dir_all(task_dir)
.expect("error removing the finished task direcotry");
}
}
}
context.into_result_msg()
}
}