use std::collections::HashMap;
use std::io::{BufRead, Write};
use std::process::Stdio;
use std::sync::Arc;
use crate::channel::{bounded, Receiver, Sender};
use crate::core::cluster::TaskResourceInfo;
use crate::core::env::{StreamApp, StreamExecutionEnvironment};
use crate::core::runtime::ClusterDescriptor;
use crate::deployment::TResourceManager;
use crate::runtime::context::Context;
use crate::runtime::ManagerType;
use crate::utils;
pub(crate) struct YarnResourceManager {
context: Arc<Context>,
cluster_descriptor: Option<ClusterDescriptor>,
yarn_command: Option<YarnCliCommand>,
}
impl YarnResourceManager {
pub fn new(context: Arc<Context>) -> Self {
YarnResourceManager {
context,
cluster_descriptor: None,
yarn_command: None,
}
}
}
impl TResourceManager for YarnResourceManager {
fn prepare(&mut self, context: &Context, job_descriptor: &ClusterDescriptor) {
self.cluster_descriptor = Some(job_descriptor.clone());
self.yarn_command = Some(YarnCliCommand::new(&context, job_descriptor));
}
fn worker_allocate<S>(
&self,
_stream_app: &S,
_stream_env: &StreamExecutionEnvironment,
) -> anyhow::Result<Vec<TaskResourceInfo>>
where
S: StreamApp + 'static,
{
let cluster_descriptor = self.cluster_descriptor.as_ref().unwrap();
let mut task_args = Vec::new();
for task_manager_descriptor in &cluster_descriptor.worker_managers {
let mut args = HashMap::new();
args.insert(
"cluster_mode".to_string(),
self.context.cluster_mode.to_string(),
);
args.insert("manager_type".to_string(), ManagerType::Worker.to_string());
args.insert(
"application_id".to_string(),
self.context.application_id.clone(),
);
args.insert(
"task_manager_id".to_string(),
task_manager_descriptor.task_manager_id.clone(),
);
args.insert(
"coordinator_address".to_string(),
cluster_descriptor.coordinator_manager.web_address.clone(),
);
task_args.push(args);
}
self.yarn_command.as_ref().unwrap().allocate(task_args)
}
fn stop_workers(&self, task_ids: Vec<TaskResourceInfo>) -> anyhow::Result<()> {
self.yarn_command.as_ref().unwrap().stop(task_ids)
}
}
#[derive(Serialize, Deserialize)]
struct Data<T> {
cmd: String,
cmd_id: String,
data: T,
}
impl<T> Data<T> {
fn new(cmd: String, cmd_id: String, data: T) -> Self {
Data { cmd, cmd_id, data }
}
}
type Command = Data<Vec<HashMap<String, String>>>;
type Response = Data<Vec<TaskResourceInfo>>;
type StopCommand = Data<Vec<TaskResourceInfo>>;
const COMMAND_PREFIX: &'static str = "/*rlink-rs_yarn*/";
fn parse_command(command_line: &str) -> Option<&str> {
command_line
.find(COMMAND_PREFIX)
.map(|pos| &command_line[pos + COMMAND_PREFIX.len()..command_line.len()])
}
fn parse_line(command_line: &std::io::Result<String>) -> Option<&str> {
command_line
.as_ref()
.map(|x| parse_command(x.as_str()))
.unwrap_or(None)
}
struct YarnCliCommand {
cmd_sender: Sender<String>,
ret_receiver: Receiver<String>,
}
impl YarnCliCommand {
pub fn new(context: &Context, cluster_descriptor: &ClusterDescriptor) -> Self {
let coordinator_manager = &cluster_descriptor.coordinator_manager;
let child = std::process::Command::new("java")
.arg("-Xmx256M")
.arg(context.yarn_manager_main_class.as_str())
.arg("--coordinator_address")
.arg(coordinator_manager.web_address.as_str())
.arg("--worker_process_path")
.arg(context.worker_process_path.as_str())
.arg("--memory_mb")
.arg(coordinator_manager.memory_mb.to_string())
.arg("--v_cores")
.arg(coordinator_manager.v_cores.to_string())
.arg("--exclusion_nodes")
.arg(context.exclusion_nodes.as_str())
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.unwrap();
let (cmd_sender, cmd_receiver) = bounded::<String>(2);
let (ret_sender, ret_receiver) = bounded::<String>(200);
match child.stdin {
Some(mut stdin) => {
std::thread::spawn(move || {
let cmd_receiver = cmd_receiver;
loop {
match cmd_receiver.recv() {
Ok(cmd) => {
let c = format!("{}\n", cmd);
info!("cmd_receiver cmd={}", c);
match stdin.write(c.as_bytes()) {
Ok(_) => {}
Err(e) => error!("stdin write error. {}", e),
}
}
Err(e) => {
panic!("stdin recv error. {}", e);
}
}
}
});
}
None => error!("stdin not found"),
};
match child.stderr {
Some(stderr) => {
let ret_sender = ret_sender.clone();
std::thread::spawn(move || {
std::io::BufReader::new(stderr).lines().for_each(|txt| {
error!("command line: {:?}", txt);
parse_line(&txt)
.map(|command| ret_sender.send(command.to_string()).unwrap());
});
});
}
_ => {}
};
match child.stdout {
Some(stdout) => {
std::thread::spawn(move || {
std::io::BufReader::new(stdout).lines().for_each(|txt| {
info!("command line: {:?}", txt);
parse_line(&txt)
.map(|command| ret_sender.send(command.to_string()).unwrap());
});
});
}
_ => {}
};
YarnCliCommand {
cmd_sender,
ret_receiver,
}
}
pub fn allocate(
&self,
task_args: Vec<HashMap<String, String>>,
) -> anyhow::Result<Vec<TaskResourceInfo>> {
let cmd_id = utils::generator::gen_with_ts();
let command = Command::new("allocate".to_string(), cmd_id.to_string(), task_args);
let command_json = serde_json::to_string(&command).unwrap();
let cmd_str = format!("{}\n", command_json);
info!("send command: {}", cmd_str);
match self.cmd_sender.send(cmd_str) {
Ok(_) => {}
Err(e) => {
error!("send command error. {}", e);
}
}
let txt = self.ret_receiver.recv()?;
let response = serde_json::from_slice::<Response>(txt.as_bytes())?;
if !response.cmd_id.eq(cmd_id.as_str()) {
Err(anyhow::Error::msg("`cmd_id` is inconsistency"))
} else {
Ok(response.data)
}
}
pub fn stop(&self, task_infos: Vec<TaskResourceInfo>) -> anyhow::Result<()> {
let cmd_id = utils::generator::gen_with_ts();
let command = StopCommand::new("stop".to_string(), cmd_id.to_string(), task_infos);
let command_json = serde_json::to_string(&command).unwrap();
let cmd_str = format!("{}\n", command_json);
info!("send command: {}", cmd_str);
match self.cmd_sender.send(cmd_str) {
Ok(_) => {}
Err(e) => {
error!("send command error. {}", e);
}
}
let txt = self.ret_receiver.recv()?;
let response = serde_json::from_slice::<Response>(txt.as_bytes())?;
if !response.cmd_id.eq(cmd_id.as_str()) {
Err(anyhow::Error::msg("`cmd_id` is inconsistency"))
} else {
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use crate::deployment::yarn::Data;
use crate::utils;
#[test]
pub fn command_json_test() {
let mut map = HashMap::new();
map.insert("a".to_string(), "b".to_string());
let mut task_args: Vec<HashMap<String, String>> = Vec::new();
task_args.push(map);
let cmd_id = utils::generator::gen_with_ts();
let command = Data::new("allocate".to_string(), cmd_id, task_args);
let command_json = serde_json::to_string(&command).unwrap();
let cmd_str = format!("{}\n", command_json);
println!("{}", cmd_str)
}
}