Agent to teardown sessions on signals
// license that can be found in the LICENSE file.
use env::{MANAGER_URL, PEM};
-use futures::{FutureExt, TryFutureExt};
+use futures::{
+ future::{select, AbortHandle, Abortable},
+ FutureExt, TryFutureExt,
+};
use iml_agent::{
agent_error::Result,
daemon_plugins, env,
http_comms::{agent_client::AgentClient, crypto_client, session},
poller, reader,
};
+use tokio::signal::unix::{signal, SignalKind};
#[tokio::main]
async fn main() -> Result<()> {
@@ -31,15 +35,30 @@ async fn main() -> Result<()> {
let registry_keys: Vec<iml_wire_types::PluginName> = registry.keys().cloned().collect();
let sessions = session::Sessions::new(&registry_keys);
- tokio::spawn(
+ let (reader, reader_reg) = AbortHandle::new_pair();
+ tokio::spawn(Abortable::new(
reader::create_reader(sessions.clone(), agent_client.clone(), registry)
.map_err(|e| {
tracing::error!("{}", e);
})
.map(drop),
- );
+ reader_reg,
+ ));
+
+ let (poller, poller_reg) = AbortHandle::new_pair();
+ tokio::spawn(Abortable::new(
+ poller::create_poller(agent_client, sessions.clone()),
+ poller_reg,
+ ));
+
+ let mut sigterm = signal(SignalKind::terminate()).expect("Could not listen to SIGTERM");
+ let mut sigint = signal(SignalKind::interrupt()).expect("Could not listen to SIGINT");
+ select(sigterm.recv().boxed(), sigint.recv().boxed()).await;
- poller::create_poller(agent_client, sessions).await;
+ tracing::info!("Terminating on signal...");
+ poller.abort();
+ reader.abort();
+ sessions.terminate_all_sessions().await?;
Ok(())
}
