use std::{
collections::{BTreeMap, HashSet},
fs::File,
io::{Read, Write},
os::unix::io::{FromRawFd, IntoRawFd},
os::unix::net::UnixStream,
time::{Duration, Instant},
};
use anyhow::{bail, Context};
use async_io::Async;
use futures::{channel::mpsc::*, SinkExt, StreamExt};
use nom::{Err, HexDisplay, Offset};
use sozu_command_lib::{
buffer::fixed::Buffer,
command::{
CommandRequest, CommandRequestOrder, CommandResponse, CommandResponseContent,
CommandStatus, FrontendFilters, ListedFrontends, RunState, WorkerInfo, PROTOCOL_VERSION,
},
config::Config,
logging,
parser::parse_several_commands,
proxy::{
AggregatedMetricsData, MetricsConfiguration, ProxyRequest, ProxyRequestOrder,
ProxyResponseContent, ProxyResponseStatus, Query, QueryAnswer, QueryClusterType,
},
scm_socket::Listeners,
state::get_cluster_ids_by_domain,
};
use sozu::metrics::METRICS;
use crate::{
command::{CommandMessage, CommandServer, RequestIdentifier, Response, Success, Worker},
upgrade::fork_main_into_new_main,
worker::start_worker,
};
impl CommandServer {
pub async fn handle_client_request(
&mut self,
client_id: String,
request: CommandRequest,
) -> anyhow::Result<Success> {
trace!("Received order {:?}", request);
let request_identifier = RequestIdentifier {
client: client_id.to_owned(),
request: request.id.to_owned(),
};
let cloned_identifier = request_identifier.clone();
let result: anyhow::Result<Option<Success>> = match request.order {
CommandRequestOrder::SaveState { path } => self.save_state(&path).await,
CommandRequestOrder::DumpState => self.dump_state().await,
CommandRequestOrder::ListWorkers => self.list_workers().await,
CommandRequestOrder::ListFrontends(filters) => self.list_frontends(filters).await,
CommandRequestOrder::LoadState { path } => {
self.load_state(
Some(request_identifier.client),
request_identifier.request,
&path,
)
.await
}
CommandRequestOrder::LaunchWorker(tag) => {
self.launch_worker(request_identifier, &tag).await
}
CommandRequestOrder::UpgradeMain => self.upgrade_main(request_identifier).await,
CommandRequestOrder::UpgradeWorker(worker_id) => {
self.upgrade_worker(request_identifier, worker_id).await
}
CommandRequestOrder::Proxy(proxy_request_order) => match *proxy_request_order {
ProxyRequestOrder::ConfigureMetrics(config) => {
self.configure_metrics(request_identifier, config).await
}
ProxyRequestOrder::Query(query) => self.query(request_identifier, query).await,
ProxyRequestOrder::Logging(logging_filter) => {
self.set_logging_level(logging_filter)
}
order => {
self.worker_order(request_identifier, order, request.worker_id)
.await
}
},
CommandRequestOrder::SubscribeEvents => {
self.event_subscribers.insert(client_id.clone());
Ok(Some(Success::SubscribeEvent(client_id.clone())))
}
CommandRequestOrder::ReloadConfiguration { path } => {
self.reload_configuration(request_identifier, path).await
}
CommandRequestOrder::Status => self.status(request_identifier).await,
};
match result {
Ok(Some(success)) => {
info!("{}", success);
return_success(self.command_tx.clone(), cloned_identifier, success).await;
}
Err(anyhow_error) => {
let formatted = format!("{:#}", anyhow_error);
error!("{:#}", formatted);
return_error(self.command_tx.clone(), cloned_identifier, formatted).await;
}
Ok(None) => {
}
}
Ok(Success::HandledClientRequest)
}
pub async fn save_state(&mut self, path: &str) -> anyhow::Result<Option<Success>> {
let mut file = File::create(&path)
.with_context(|| format!("could not open file at path: {}", &path))?;
let counter = self
.save_state_to_file(&mut file)
.with_context(|| "failed writing state to file")?;
info!("wrote {} commands to {}", counter, path);
Ok(Some(Success::SaveState(counter, path.into())))
}
pub fn save_state_to_file(&mut self, file: &mut File) -> anyhow::Result<usize> {
let mut counter = 0usize;
let orders = self.state.generate_orders();
let result: anyhow::Result<usize> = (move || {
for command in orders {
let message = CommandRequest::new(
format!("SAVE-{}", counter),
CommandRequestOrder::Proxy(Box::new(command)),
None,
);
file.write_all(
&serde_json::to_string(&message)
.map(|s| s.into_bytes())
.unwrap_or_default(),
)
.with_context(|| {
format!(
"Could not add this instruction line to the saved state file: {:?}",
message
)
})?;
file.write_all(&b"\n\0"[..])
.with_context(|| "Could not add new line to the saved state file")?;
if counter % 1000 == 0 {
info!("writing command {}", counter);
file.sync_all()
.with_context(|| "Failed to sync the saved state file")?;
}
counter += 1;
}
file.sync_all()
.with_context(|| "Failed to sync the saved state file")?;
Ok(counter)
})();
result.with_context(|| "Could not write the state onto the state file")
}
pub async fn dump_state(&mut self) -> anyhow::Result<Option<Success>> {
let state = self.state.clone();
Ok(Some(Success::DumpState(CommandResponseContent::State(
Box::new(state),
))))
}
pub async fn load_state(
&mut self,
client_id: Option<String>,
request_id: String,
path: &str,
) -> anyhow::Result<Option<Success>> {
let mut file =
File::open(&path).with_context(|| format!("Cannot open file at path {}", path))?;
let mut buffer = Buffer::with_capacity(200000);
info!("starting to load state from {}", path);
let mut message_counter = 0usize;
let mut diff_counter = 0usize;
let (load_state_tx, mut load_state_rx) = futures::channel::mpsc::channel(10000);
loop {
let previous = buffer.available_data();
match file.read(buffer.space()) {
Ok(sz) => buffer.fill(sz),
Err(e) => {
bail!("Error reading the saved state file: {}", e);
}
};
if buffer.available_data() == 0 {
debug!("Empty buffer");
break;
}
let mut offset = 0usize;
match parse_several_commands::<CommandRequest>(buffer.data()) {
Ok((i, requests)) => {
if !i.is_empty() {
debug!("could not parse {} bytes", i.len());
if previous == buffer.available_data() {
bail!("error consuming load state message");
}
}
offset = buffer.data().offset(i);
if requests.iter().any(|o| {
if o.version > PROTOCOL_VERSION {
error!("configuration protocol version mismatch: Sōzu handles up to version {}, the message uses version {}", PROTOCOL_VERSION, o.version);
true
} else {
false
}
}) {
break;
}
for request in requests {
if let CommandRequestOrder::Proxy(order) = request.order {
message_counter += 1;
if self.state.handle_order(&order).is_ok() {
diff_counter += 1;
let mut found = false;
let id = format!("LOAD-STATE-{}-{}", request_id, diff_counter);
for ref mut worker in self.workers.iter_mut().filter(|worker| {
worker.run_state != RunState::Stopping
&& worker.run_state != RunState::Stopped
}) {
let worker_message_id = format!("{}-{}", id, worker.id);
worker.send(worker_message_id.clone(), *order.clone()).await;
self.in_flight
.insert(worker_message_id, (load_state_tx.clone(), 1));
found = true;
}
if !found {
bail!("no worker found");
}
}
}
}
}
Err(Err::Incomplete(_)) => {
if buffer.available_data() == buffer.capacity() {
error!(
"message too big, stopping parsing:\n{}",
buffer.data().to_hex(16)
);
break;
}
}
Err(parse_error) => {
bail!("saved state parse error: {:?}", parse_error);
}
}
buffer.consume(offset);
}
info!(
"stopped loading data from file, remaining: {} bytes, saw {} messages, generated {} diff messages",
buffer.available_data(), message_counter, diff_counter
);
if diff_counter > 0 {
info!(
"state loaded from {}, will start sending {} messages to workers",
path, diff_counter
);
let command_tx = self.command_tx.to_owned();
let path = path.to_owned();
smol::spawn(async move {
let mut ok = 0usize;
let mut error = 0usize;
while let Some((proxy_response, _)) = load_state_rx.next().await {
match proxy_response.status {
ProxyResponseStatus::Ok => {
ok += 1;
}
ProxyResponseStatus::Processing => {}
ProxyResponseStatus::Error(message) => {
error!("{}", message);
error += 1;
}
};
debug!("ok:{}, error: {}", ok, error);
}
let request_identifier = match client_id {
Some(client_id) => RequestIdentifier::new(client_id, request_id),
None => {
match error {
0 => info!("loading state: {} ok messages, 0 errors", ok),
_ => error!("loading state: {} ok messages, {} errors", ok, error),
}
return;
}
};
match error {
0 => {
return_success(
command_tx,
request_identifier,
Success::LoadState(path.to_string(), ok, error),
)
.await;
}
_ => {
return_error(
command_tx,
request_identifier,
format!(
"Loading state failed, ok: {}, error: {}, path: {}",
ok, error, path
),
)
.await;
}
}
})
.detach();
} else {
info!("no messages sent to workers: local state already had those messages");
if let Some(client_id) = client_id {
return_success(
self.command_tx.clone(),
RequestIdentifier::new(client_id, request_id),
Success::LoadState(path.to_string(), 0, 0),
)
.await;
}
}
self.backends_count = self.state.count_backends();
self.frontends_count = self.state.count_frontends();
gauge!("configuration.clusters", self.state.clusters.len());
gauge!("configuration.backends", self.backends_count);
gauge!("configuration.frontends", self.frontends_count);
Ok(None)
}
pub async fn list_frontends(
&mut self,
filters: FrontendFilters,
) -> anyhow::Result<Option<Success>> {
info!(
"Received a request to list frontends, along these filters: {:?}",
filters
);
let list_all = !filters.http && !filters.https && !filters.tcp;
let mut listed_frontends = ListedFrontends::default();
if filters.http || list_all {
for http_frontend in self.state.http_fronts.iter().filter(|f| {
if let Some(domain) = &filters.domain {
f.1.hostname.contains(domain)
} else {
true
}
}) {
listed_frontends
.http_frontends
.push(http_frontend.1.to_owned());
}
}
if filters.https || list_all {
for https_frontend in self.state.https_fronts.iter().filter(|f| {
if let Some(domain) = &filters.domain {
f.1.hostname.contains(domain)
} else {
true
}
}) {
listed_frontends
.https_frontends
.push(https_frontend.1.to_owned());
}
}
if (filters.tcp || list_all) && filters.domain.is_none() {
for tcp_frontend in self.state.tcp_fronts.values().flat_map(|v| v.iter()) {
listed_frontends.tcp_frontends.push(tcp_frontend.to_owned())
}
}
Ok(Some(Success::ListFrontends(
CommandResponseContent::FrontendList(listed_frontends),
)))
}
pub async fn list_workers(&mut self) -> anyhow::Result<Option<Success>> {
let workers: Vec<WorkerInfo> = self
.workers
.iter()
.map(|worker| WorkerInfo {
id: worker.id,
pid: worker.pid,
run_state: worker.run_state,
})
.collect();
debug!("workers: {:#?}", workers);
Ok(Some(Success::ListWorkers(CommandResponseContent::Workers(
workers,
))))
}
pub async fn launch_worker(
&mut self,
request_identifier: RequestIdentifier,
_tag: &str,
) -> anyhow::Result<Option<Success>> {
let mut worker = start_worker(
self.next_worker_id,
&self.config,
self.executable_path.clone(),
&self.state,
None,
)
.with_context(|| format!("Failed at creating worker {}", self.next_worker_id))?;
return_processing(
self.command_tx.clone(),
request_identifier.clone(),
"Sending configuration orders to the new worker...",
)
.await;
info!("created new worker: {}", worker.id);
self.next_worker_id += 1;
let sock = worker
.worker_channel
.take()
.expect("No channel on the worker being launched")
.sock;
let (worker_tx, worker_rx) = channel(10000);
worker.sender = Some(worker_tx);
let stream = Async::new(unsafe {
let fd = sock.into_raw_fd();
UnixStream::from_raw_fd(fd)
})?;
let id = worker.id;
let command_tx = self.command_tx.clone();
smol::spawn(async move {
super::worker_loop(id, stream, command_tx, worker_rx).await;
})
.detach();
info!(
"sending listeners: to the new worker: {:?}",
worker.scm_socket.send_listeners(&Listeners {
http: Vec::new(),
tls: Vec::new(),
tcp: Vec::new(),
})
);
let activate_orders = self.state.generate_activate_orders();
for (count, order) in activate_orders.into_iter().enumerate() {
worker
.send(format!("{}-ACTIVATE-{}", id, count), order)
.await;
}
self.workers.push(worker);
return_success(
self.command_tx.clone(),
request_identifier,
Success::WorkerLaunched(id),
)
.await;
Ok(None)
}
pub async fn upgrade_main(
&mut self,
request_identifier: RequestIdentifier,
) -> anyhow::Result<Option<Success>> {
self.disable_cloexec_before_upgrade()?;
return_processing(
self.command_tx.clone(),
request_identifier,
"The proxy is processing the upgrade command.",
)
.await;
let upgrade_data = self.generate_upgrade_data();
let (new_main_pid, mut fork_confirmation_channel) =
fork_main_into_new_main(self.executable_path.clone(), upgrade_data)
.with_context(|| "Could not start a new main process")?;
if let Err(e) = fork_confirmation_channel.blocking() {
error!(
"Could not block the fork confirmation channel: {}. This is not normal, you may need to restart sozu",
e
);
}
let received_ok_from_new_process = fork_confirmation_channel.read_message();
debug!("upgrade channel sent {:?}", received_ok_from_new_process);
if let Err(e) = self
.accept_cancel
.take() .expect("No channel on the main process")
.send(())
{
error!("could not close the accept loop: {:?}", e);
}
if !received_ok_from_new_process
.with_context(|| "Did not receive fork confirmation from new worker")?
{
bail!("forking the new worker failed")
}
info!("wrote final message, closing");
Ok(Some(Success::UpgradeMain(new_main_pid)))
}
pub async fn upgrade_worker(
&mut self,
request_identifier: RequestIdentifier,
id: u32,
) -> anyhow::Result<Option<Success>> {
info!(
"client[{}] msg {} wants to upgrade worker {}",
request_identifier.client, request_identifier.request, id
);
if !self.workers.iter().any(|worker| {
worker.id == id
&& worker.run_state != RunState::Stopping
&& worker.run_state != RunState::Stopped
}) {
bail!(format!(
"The worker {} does not exist, or is stopped / stopping.",
&id
));
}
let next_id = self.next_worker_id;
let mut new_worker = start_worker(
next_id,
&self.config,
self.executable_path.clone(),
&self.state,
None,
)
.with_context(|| "failed at creating worker")?;
return_processing(
self.command_tx.clone(),
request_identifier.clone(),
"Sending configuration orders to the worker",
)
.await;
info!("created new worker: {}", next_id);
self.next_worker_id += 1;
let sock = new_worker
.worker_channel
.take()
.with_context(|| "No channel on new worker".to_string())?
.sock;
let (worker_tx, worker_rx) = channel(10000);
new_worker.sender = Some(worker_tx);
new_worker
.sender
.as_mut()
.with_context(|| "No sender on new worker".to_string())?
.send(ProxyRequest {
id: format!("UPGRADE-{}-STATUS", id),
order: ProxyRequestOrder::Status,
})
.await
.with_context(|| {
format!(
"could not send status message to worker {:?}",
new_worker.id,
)
})?;
let mut listeners = None;
{
let old_worker: &mut Worker = self
.workers
.iter_mut()
.find(|worker| worker.id == id)
.unwrap();
let (sockets_return_tx, mut sockets_return_rx) = futures::channel::mpsc::channel(3);
let id = format!("{}-return-sockets", request_identifier.client);
self.in_flight.insert(id.clone(), (sockets_return_tx, 1));
old_worker
.send(id.clone(), ProxyRequestOrder::ReturnListenSockets)
.await;
info!("sent ReturnListenSockets to old worker");
let cloned_command_tx = self.command_tx.clone();
let cloned_req_id = request_identifier.clone();
smol::spawn(async move {
while let Some((proxy_response, _)) = sockets_return_rx.next().await {
match proxy_response.status {
ProxyResponseStatus::Ok => {
info!("returnsockets OK");
break;
}
ProxyResponseStatus::Processing => {
info!("returnsockets processing");
}
ProxyResponseStatus::Error(message) => {
return_error(cloned_command_tx, cloned_req_id, message).await;
break;
}
};
}
})
.detach();
let mut counter = 0usize;
loop {
info!("waiting for listen sockets from the old worker");
if let Err(e) = old_worker.scm_socket.set_blocking(true) {
error!("Could not set the old worker socket to blocking: {}", e);
};
match old_worker.scm_socket.receive_listeners() {
Ok(l) => {
listeners = Some(l);
break;
}
Err(error) => {
error!(
"Could not receive listerners from scm socket with file descriptor {}:\n{:?}",
old_worker.scm_socket.fd, error
);
counter += 1;
if counter == 50 {
break;
}
std::thread::sleep(Duration::from_millis(100));
}
}
}
info!("got the listen sockets from the old worker");
old_worker.run_state = RunState::Stopping;
let (softstop_tx, mut softstop_rx) = futures::channel::mpsc::channel(10);
let softstop_id = format!("{}-softstop", request_identifier.client);
self.in_flight.insert(softstop_id.clone(), (softstop_tx, 1));
old_worker
.send(softstop_id.clone(), ProxyRequestOrder::SoftStop)
.await;
let mut command_tx = self.command_tx.clone();
let cloned_request_identifier = request_identifier.clone();
let worker_id = old_worker.id;
smol::spawn(async move {
while let Some((proxy_response, _)) = softstop_rx.next().await {
match proxy_response.status {
ProxyResponseStatus::Ok => {
info!("softstop OK"); if let Err(e) = command_tx
.send(CommandMessage::WorkerClose { worker_id })
.await
{
error!(
"could not send worker close message to {}: {:?}",
worker_id, e
);
}
break;
}
ProxyResponseStatus::Processing => {
info!("softstop processing");
}
ProxyResponseStatus::Error(message) => {
info!("softstop error: {:?}", message);
break;
}
};
}
return_processing(
command_tx.clone(),
cloned_request_identifier,
"Processing softstop responses from the workers...",
)
.await;
})
.detach();
}
match listeners {
Some(l) => {
info!(
"sending listeners: to the new worker: {:?}",
new_worker.scm_socket.send_listeners(&l)
);
l.close();
}
None => error!("could not get the list of listeners from the previous worker"),
};
let stream = Async::new(unsafe {
let fd = sock.into_raw_fd();
UnixStream::from_raw_fd(fd)
})?;
let id = new_worker.id;
let command_tx = self.command_tx.clone();
smol::spawn(async move {
super::worker_loop(id, stream, command_tx, worker_rx).await;
})
.detach();
let activate_orders = self.state.generate_activate_orders();
for (count, order) in activate_orders.into_iter().enumerate() {
new_worker
.send(
format!("{}-ACTIVATE-{}", request_identifier.client, count),
order,
)
.await;
}
info!("sent config messages to the new worker");
self.workers.push(new_worker);
info!("finished upgrade");
Ok(Some(Success::UpgradeWorker(id)))
}
pub async fn reload_configuration(
&mut self,
request_identifier: RequestIdentifier,
config_path: Option<String>,
) -> anyhow::Result<Option<Success>> {
let path = config_path.as_deref().unwrap_or(&self.config.config_path);
let new_config = Config::load_from_path(path)
.with_context(|| format!("cannot load configuration from '{}'", path))?;
let mut diff_counter = 0usize;
let (load_state_tx, mut load_state_rx) = futures::channel::mpsc::channel(10000);
return_processing(
self.command_tx.clone(),
request_identifier.clone(),
"Reloading configuration, sending config messages to workers...",
)
.await;
for message in new_config.generate_config_messages() {
if let CommandRequestOrder::Proxy(order) = message.order {
if self.state.handle_order(&order).is_ok() {
diff_counter += 1;
let mut found = false;
let id = format!(
"LOAD-STATE-{}-{}",
&request_identifier.request, diff_counter
);
for ref mut worker in self.workers.iter_mut().filter(|worker| {
worker.run_state != RunState::Stopping
&& worker.run_state != RunState::Stopped
}) {
let worker_message_id = format!("{}-{}", id, worker.id);
worker.send(worker_message_id.clone(), *order.clone()).await;
self.in_flight
.insert(worker_message_id, (load_state_tx.clone(), 1));
found = true;
}
if !found {
error!("no worker found");
}
}
}
}
let command_tx = self.command_tx.clone();
let cloned_identifier = request_identifier.clone();
if diff_counter > 0 {
info!(
"state loaded from {}, will start sending {} messages to workers",
new_config.config_path, diff_counter
);
smol::spawn(async move {
let mut ok = 0usize;
let mut error = 0usize;
while let Some((proxy_response, _)) = load_state_rx.next().await {
match proxy_response.status {
ProxyResponseStatus::Ok => {
ok += 1;
}
ProxyResponseStatus::Processing => {}
ProxyResponseStatus::Error(message) => {
error!("{}", message);
error += 1;
}
};
debug!("ok:{}, error: {}", ok, error);
}
if error == 0 {
return_success(
command_tx,
cloned_identifier,
Success::ReloadConfiguration(ok, error),
)
.await;
} else {
return_error(
command_tx,
cloned_identifier,
format!(
"Reloading configuration failed. ok: {} messages, error: {}",
ok, error
),
)
.await;
}
})
.detach();
} else {
info!("no messages sent to workers: local state already had those messages");
}
self.backends_count = self.state.count_backends();
self.frontends_count = self.state.count_frontends();
gauge!("configuration.clusters", self.state.clusters.len());
gauge!("configuration.backends", self.backends_count);
gauge!("configuration.frontends", self.frontends_count);
self.config = new_config;
Ok(None)
}
pub async fn status(
&mut self,
request_identifier: RequestIdentifier,
) -> anyhow::Result<Option<Success>> {
info!("Requesting the status of all workers.");
let (status_tx, mut status_rx) = futures::channel::mpsc::channel(self.workers.len() * 2);
let mut worker_info_map: BTreeMap<String, WorkerInfo> = BTreeMap::new();
let prefix = format!("{}-status-", request_identifier.client);
return_processing(
self.command_tx.clone(),
request_identifier.clone(),
"Sending status requests to workers...",
)
.await;
let mut count = 0usize;
for ref mut worker in self.workers.iter_mut() {
info!("Worker {} is {}", worker.id, worker.run_state);
let worker_request_id = format!("{}{}", prefix, worker.id);
if worker.run_state == RunState::Running {
info!("Summoning status of worker {}", worker.id);
worker
.send(worker_request_id.clone(), ProxyRequestOrder::Status)
.await;
count += 1;
self.in_flight
.insert(worker_request_id.clone(), (status_tx.clone(), 1));
}
worker_info_map.insert(worker_request_id, worker.info());
}
let command_tx = self.command_tx.clone();
let thread_request_identifier = request_identifier.clone();
let now = Instant::now();
smol::spawn(async move {
let mut i = 0;
while let Some((proxy_response, _)) = status_rx.next().await {
info!(
"received response with id {}: {:?}",
proxy_response.id, proxy_response
);
let new_run_state = match proxy_response.status {
ProxyResponseStatus::Ok => RunState::Running,
ProxyResponseStatus::Processing => continue,
ProxyResponseStatus::Error(_) => RunState::NotAnswering,
};
worker_info_map
.entry(proxy_response.id)
.and_modify(|worker_info| worker_info.run_state = new_run_state);
i += 1;
if i == count || now.elapsed() > Duration::from_secs(10) {
break;
}
}
let worker_info_vec: Vec<WorkerInfo> = worker_info_map
.iter()
.map(|(_, worker_info)| worker_info.to_owned())
.collect();
return_success(
command_tx,
thread_request_identifier,
Success::Status(CommandResponseContent::Status(worker_info_vec)),
)
.await;
})
.detach();
Ok(None)
}
pub async fn configure_metrics(
&mut self,
request_identifier: RequestIdentifier,
config: MetricsConfiguration,
) -> anyhow::Result<Option<Success>> {
let (metrics_tx, mut metrics_rx) = futures::channel::mpsc::channel(self.workers.len() * 2);
let mut count = 0usize;
for ref mut worker in self
.workers
.iter_mut()
.filter(|worker| worker.run_state != RunState::Stopped)
{
let req_id = format!("{}-metrics-{}", request_identifier.client, worker.id);
worker
.send(
req_id.clone(),
ProxyRequestOrder::ConfigureMetrics(config.clone()),
)
.await;
count += 1;
self.in_flight.insert(req_id, (metrics_tx.clone(), 1));
}
let prefix = format!("{}-metrics-", request_identifier.client);
let command_tx = self.command_tx.clone();
let thread_request_identifier = request_identifier.clone();
smol::spawn(async move {
let mut responses = Vec::new();
let mut i = 0;
while let Some((proxy_response, _)) = metrics_rx.next().await {
match proxy_response.status {
ProxyResponseStatus::Ok => {
let tag = proxy_response.id.trim_start_matches(&prefix).to_string();
responses.push((tag, proxy_response));
}
ProxyResponseStatus::Processing => {
continue;
}
ProxyResponseStatus::Error(_) => {
let tag = proxy_response.id.trim_start_matches(&prefix).to_string();
responses.push((tag, proxy_response));
}
};
i += 1;
if i == count {
break;
}
}
let mut messages = vec![];
let mut has_error = false;
for response in responses.iter() {
match response.1.status {
ProxyResponseStatus::Error(ref e) => {
messages.push(format!("{}: {}", response.0, e));
has_error = true;
}
_ => messages.push(format!("{}: OK", response.0)),
}
}
if has_error {
return_error(command_tx, thread_request_identifier, messages.join(", ")).await;
} else {
return_success(
command_tx,
thread_request_identifier,
Success::Metrics(config),
)
.await;
}
})
.detach();
Ok(None)
}
pub async fn query(
&mut self,
request_identifier: RequestIdentifier,
query: Query,
) -> anyhow::Result<Option<Success>> {
debug!("Received this query: {:?}", query);
let (query_tx, mut query_rx) = futures::channel::mpsc::channel(self.workers.len() * 2);
let mut count = 0usize;
for ref mut worker in self
.workers
.iter_mut()
.filter(|worker| worker.run_state != RunState::Stopped)
{
let req_id = format!("{}-query-{}", request_identifier.client, worker.id);
worker
.send(req_id.clone(), ProxyRequestOrder::Query(query.clone()))
.await;
count += 1;
self.in_flight.insert(req_id, (query_tx.clone(), 1));
}
return_processing(
self.command_tx.clone(),
request_identifier.clone(),
"Query was sent to the workers...",
)
.await;
let mut main_query_answer = None;
match &query {
Query::ClustersHashes => {
main_query_answer = Some(QueryAnswer::ClustersHashes(self.state.hash_state()));
}
Query::Clusters(query_type) => {
main_query_answer = Some(QueryAnswer::Clusters(match query_type {
QueryClusterType::ClusterId(cluster_id) => {
vec![self.state.cluster_state(cluster_id)]
}
QueryClusterType::Domain(domain) => {
let cluster_ids = get_cluster_ids_by_domain(
&self.state,
domain.hostname.clone(),
domain.path.clone(),
);
cluster_ids
.iter()
.map(|cluster_id| self.state.cluster_state(cluster_id))
.collect()
}
}));
}
Query::Certificates(_) => {}
Query::Metrics(_) => {}
};
let command_tx = self.command_tx.clone();
let cloned_identifier = request_identifier.clone();
let main_metrics =
METRICS.with(|metrics| (*metrics.borrow_mut()).dump_local_proxy_metrics());
smol::spawn(async move {
let mut responses = Vec::new();
let mut i = 0;
while let Some((proxy_response, worker_id)) = query_rx.next().await {
match proxy_response.status {
ProxyResponseStatus::Ok => {
responses.push((worker_id, proxy_response));
}
ProxyResponseStatus::Processing => {
info!("metrics processing");
continue;
}
ProxyResponseStatus::Error(_) => {
responses.push((worker_id, proxy_response));
}
};
i += 1;
if i == count {
break;
}
}
let mut proxy_responses_map: BTreeMap<String, QueryAnswer> = responses
.into_iter()
.filter_map(|(worker_id, proxy_response)| {
if let Some(ProxyResponseContent::Query(d)) = proxy_response.content {
Some((worker_id.to_string(), d))
} else {
None
}
})
.collect();
let success = match &query {
&Query::ClustersHashes | &Query::Clusters(_) => {
let main = main_query_answer.unwrap(); proxy_responses_map.insert(String::from("main"), main);
Success::Query(CommandResponseContent::Query(proxy_responses_map))
}
&Query::Certificates(_) => {
info!(
"certificates query answer received: {:?}",
proxy_responses_map
);
Success::Query(CommandResponseContent::Query(proxy_responses_map))
}
Query::Metrics(options) => {
debug!("metrics query answer received: {:?}", proxy_responses_map);
if options.list {
Success::Query(CommandResponseContent::Query(proxy_responses_map))
} else {
Success::Query(CommandResponseContent::Metrics(AggregatedMetricsData {
main: main_metrics,
workers: proxy_responses_map,
}))
}
}
};
return_success(command_tx, cloned_identifier, success).await;
})
.detach();
Ok(None)
}
pub fn set_logging_level(&mut self, logging_filter: String) -> anyhow::Result<Option<Success>> {
debug!("Changing main process log level to {}", logging_filter);
logging::LOGGER.with(|l| {
let directives = logging::parse_logging_spec(&logging_filter);
l.borrow_mut().set_directives(directives);
});
::std::env::set_var("RUST_LOG", &logging_filter);
debug!("Logging level now: {}", ::std::env::var("RUST_LOG")?);
Ok(Some(Success::Logging(logging_filter)))
}
pub async fn worker_order(
&mut self,
request_identifier: RequestIdentifier,
order: ProxyRequestOrder,
worker_id: Option<u32>,
) -> anyhow::Result<Option<Success>> {
if let &ProxyRequestOrder::AddCertificate(_) = &order {
debug!("workerconfig client order AddCertificate()");
} else {
debug!("workerconfig client order {:?}", order);
}
self.state
.handle_order(&order)
.with_context(|| "Could not execute order on the state")?;
if self.config.automatic_state_save
& (order != ProxyRequestOrder::SoftStop || order != ProxyRequestOrder::HardStop)
{
if let Some(path) = self.config.saved_state.clone() {
return_processing(
self.command_tx.clone(),
request_identifier.clone(),
"Saving state to file",
)
.await;
let mut file = File::create(&path)
.with_context(|| "Could not create file to automatically save the state")?;
self.save_state_to_file(&mut file)
.with_context(|| format!("could not save state automatically to {}", path))?;
}
}
return_processing(
self.command_tx.clone(),
request_identifier.clone(),
match worker_id {
Some(id) => format!("Sending the order to worker {}", id),
None => "Sending the order to all workers".to_owned(),
},
)
.await;
let (worker_order_tx, mut worker_order_rx) =
futures::channel::mpsc::channel(self.workers.len() * 2);
let mut found = false;
let mut stopping_workers = HashSet::new();
let mut worker_count = 0usize;
for ref mut worker in self.workers.iter_mut().filter(|worker| {
worker.run_state != RunState::Stopping && worker.run_state != RunState::Stopped
}) {
if let Some(id) = worker_id {
if id != worker.id {
continue;
}
}
let should_stop_worker =
order == ProxyRequestOrder::SoftStop || order == ProxyRequestOrder::HardStop;
if should_stop_worker {
worker.run_state = RunState::Stopping;
stopping_workers.insert(worker.id);
}
let req_id = format!("{}-worker-{}", request_identifier.client, worker.id);
worker.send(req_id.clone(), order.clone()).await;
self.in_flight.insert(req_id, (worker_order_tx.clone(), 1));
found = true;
worker_count += 1;
}
let should_stop_main = (order == ProxyRequestOrder::SoftStop
|| order == ProxyRequestOrder::HardStop)
&& worker_id.is_none();
let mut command_tx = self.command_tx.clone();
let thread_request_identifier = request_identifier.clone();
smol::spawn(async move {
let mut responses = Vec::new();
let mut response_count = 0usize;
while let Some((proxy_response, worker_id)) = worker_order_rx.next().await {
match proxy_response.status {
ProxyResponseStatus::Ok => {
responses.push((worker_id, proxy_response));
if stopping_workers.contains(&worker_id) {
if let Err(e) = command_tx
.send(CommandMessage::WorkerClose { worker_id })
.await
{
error!(
"could not send worker close message to {}: {:?}",
worker_id, e
);
}
}
}
ProxyResponseStatus::Processing => {
info!("Order is processing");
continue;
}
ProxyResponseStatus::Error(_) => {
responses.push((worker_id, proxy_response));
}
};
response_count += 1;
if response_count == worker_count {
break;
}
}
if should_stop_main {
if let Err(e) = command_tx.send(CommandMessage::MasterStop).await {
error!("could not send main stop message: {:?}", e);
}
}
let mut messages = vec![];
let mut has_error = false;
for response in responses.iter() {
match response.1.status {
ProxyResponseStatus::Error(ref e) => {
messages.push(format!("{}: {}", response.0, e));
has_error = true;
}
_ => messages.push(format!("{}: OK", response.0)),
}
}
if has_error {
return_error(command_tx, thread_request_identifier, messages.join(", ")).await;
} else {
return_success(
command_tx,
thread_request_identifier,
Success::WorkerOrder(worker_id),
)
.await;
}
})
.detach();
if !found {
bail!("no worker found");
}
match order {
ProxyRequestOrder::AddBackend(_) | ProxyRequestOrder::RemoveBackend(_) => {
self.backends_count = self.state.count_backends()
}
ProxyRequestOrder::AddHttpFrontend(_)
| ProxyRequestOrder::AddHttpsFrontend(_)
| ProxyRequestOrder::AddTcpFrontend(_)
| ProxyRequestOrder::RemoveHttpFrontend(_)
| ProxyRequestOrder::RemoveHttpsFrontend(_)
| ProxyRequestOrder::RemoveTcpFrontend(_) => {
self.frontends_count = self.state.count_frontends()
}
_ => {}
};
gauge!("configuration.clusters", self.state.clusters.len());
gauge!("configuration.backends", self.backends_count);
gauge!("configuration.frontends", self.frontends_count);
Ok(None)
}
pub async fn notify_advancement_to_client(
&mut self,
request_identifier: RequestIdentifier,
response: Response,
) -> anyhow::Result<Success> {
let RequestIdentifier {
client: client_id,
request: request_id,
} = request_identifier.to_owned();
let command_response = match response {
Response::Ok(success) => {
let success_message = success.to_string();
let command_response_data = match success {
Success::DumpState(crd)
| Success::ListFrontends(crd)
| Success::ListWorkers(crd)
| Success::Query(crd)
| Success::Status(crd) => Some(crd),
_ => None,
};
CommandResponse::new(
request_id.clone(),
CommandStatus::Ok,
success_message,
command_response_data,
)
}
Response::Processing(processing_message) => CommandResponse::new(
request_id.clone(),
CommandStatus::Processing,
processing_message,
None,
),
Response::Error(error_message) => CommandResponse::new(
request_id.clone(),
CommandStatus::Error,
error_message,
None,
),
};
trace!(
"Sending response to request {} of client {}: {:?}",
request_id,
client_id,
command_response
);
match self.clients.get_mut(&client_id) {
Some(client_tx) => {
trace!("sending from main process to client loop");
client_tx.send(command_response).await.with_context(|| {
format!(
"Could not notify client {} about request {}",
client_id, request_identifier.request,
)
})?;
}
None => bail!(format!("Could not find client {}", client_id)),
}
Ok(Success::NotifiedClient(client_id))
}
}
async fn return_error<T>(
mut command_tx: Sender<CommandMessage>,
request_identifier: RequestIdentifier,
error_message: T,
) where
T: ToString,
{
let error_command_message = CommandMessage::Advancement {
request_identifier,
response: Response::Error(error_message.to_string()),
};
trace!("return_error: sending event to the command server");
if let Err(e) = command_tx.send(error_command_message).await {
error!("Error while return error to the command server: {}", e)
}
}
async fn return_processing<T>(
mut command_tx: Sender<CommandMessage>,
request_identifier: RequestIdentifier,
processing_message: T,
) where
T: ToString,
{
let processing_command_message = CommandMessage::Advancement {
request_identifier,
response: Response::Processing(processing_message.to_string()),
};
trace!("return_processing: sending event to the command server");
if let Err(e) = command_tx.send(processing_command_message).await {
error!(
"Error while returning processing to the command server: {}",
e
)
}
}
async fn return_success(
mut command_tx: Sender<CommandMessage>,
request_identifier: RequestIdentifier,
success: Success,
) {
let success_command_message = CommandMessage::Advancement {
request_identifier,
response: Response::Ok(success),
};
trace!("return_success: sending event to the command server");
if let Err(e) = command_tx.send(success_command_message).await {
error!("Error while returning success to the command server: {}", e)
}
}