I'm using the below code (from here) to handle graceful shutdown:
use tokio::signal;
pub async fn shutdown_signal() {
let ctrl_c = async {
signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
};
#[cfg(unix)]
let terminate = async {
signal::unix::signal(signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
}
tracing::info!("signal received, starting graceful shutdown...");
}
Is there a way to avoid the usage of tokio::select! macro?
You can use futures's select():
use tokio::signal;
pub async fn shutdown_signal() {
let ctrl_c = async {
signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
};
#[cfg(unix)]
let terminate = async {
signal::unix::signal(signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::pin!(ctrl_c, terminate);
futures::future::select(ctrl_c, terminate).await;
tracing::info!("signal received, starting graceful shutdown...");
}
Related
I've spent a few hours trying to figure this out and I'm pretty done. I found the question with a similar name, but that looks like something was blocking synchronously which was messing with tokio. That very well may be the issue here, but I have absolutely no idea what is causing it.
Here is a heavily stripped down version of my project which hopefully gets the issue across.
use std::io;
use futures_util::{
SinkExt,
stream::{SplitSink, SplitStream},
StreamExt,
};
use tokio::{
net::TcpStream,
sync::mpsc::{channel, Receiver, Sender},
};
use tokio_tungstenite::{
connect_async,
MaybeTlsStream,
tungstenite::Message,
WebSocketStream,
};
#[tokio::main]
async fn main() {
connect_to_server("wss://a_valid_domain.com".to_string()).await;
}
async fn read_line() -> String {
loop {
let mut str = String::new();
io::stdin().read_line(&mut str).unwrap();
str = str.trim().to_string();
if !str.is_empty() {
return str;
}
}
}
async fn connect_to_server(url: String) {
let (ws_stream, _) = connect_async(url).await.unwrap();
let (write, read) = ws_stream.split();
let (tx, rx) = channel::<ChannelMessage>(100);
tokio::spawn(channel_thread(write, rx));
tokio::spawn(handle_std_input(tx.clone()));
read_messages(read, tx).await;
}
#[derive(Debug)]
enum ChannelMessage {
Text(String),
Close,
}
// PROBLEMATIC FUNCTION
async fn channel_thread(
mut write: SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
mut rx: Receiver<ChannelMessage>,
) {
while let Some(msg) = rx.recv().await {
println!("{:?}", msg); // This only fires when buffer is full
match msg {
ChannelMessage::Text(text) => write.send(Message::Text(text)).await.unwrap(),
ChannelMessage::Close => {
write.close().await.unwrap();
rx.close();
return;
}
}
}
}
async fn read_messages(
mut read: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
tx: Sender<ChannelMessage>,
) {
while let Some(msg) = read.next().await {
let msg = match msg {
Ok(m) => m,
Err(_) => continue
};
match msg {
Message::Text(m) => println!("{}", m),
Message::Close(_) => break,
_ => {}
}
}
if !tx.is_closed() {
let _ = tx.send(ChannelMessage::Close).await;
}
}
async fn handle_std_input(tx: Sender<ChannelMessage>) {
loop {
let str = read_line().await;
if tx.is_closed() {
break;
}
tx.send(ChannelMessage::Text(str)).await.unwrap();
}
}
As you can see, what I'm trying to do is:
Connect to a websocket
Print outgoing messages from the websocket
Forward any input from stdin to the websocket
Also a custom heartbeat solution which was trimmed out
The issue lies in the channel_thread() function. I move the websocket writer into this function as well as the channel receiver. The issue is, it only loops over the sent objects when the buffer is full.
I've spent a lot of time trying to solve this, any help is greatly appreciated.
Here, you make a blocking synchronous call in an async context:
async fn read_line() -> String {
loop {
let mut str = String::new();
io::stdin().read_line(&mut str).unwrap();
// ^^^^^^^^^^^^^^^^^^^
// This is sync+blocking
str = str.trim().to_string();
if !str.is_empty() {
return str;
}
}
}
You never ever make blocking synchronous calls in an async context, because that prevents the entire thread from running other async tasks. Your channel receiver task is likely also assigned to this thread, so it's having to wait until all the blocking calls are done and whatever invokes this function yields back to the async runtime.
Tokio has its own async version of stdin, which you should use instead.
How can I have some part of my application reading user input and also listening for a shutdown.
According to tokio-docs to do this type of things I should use blocking IO in a spawned task.
For interactive uses, it is recommended to spawn a thread dedicated to user input and use blocking IO directly in that thread.
And so I did with something like this
async fn read_input(mut rx: watch::Receiver<&str>) {
let mut line = String::new();
let stdin = io::stdin();
loop {
stdin.lock().read_line(&mut line).expect("Could not read line");
let op = line.trim_right();
if op == "EXIT" {
break;
} else if op == "send" {
// send_stuff();
}
line.clear();
}
}
the thing is, how can I check the receiver channel for a shutdown and break this loop?
If I await the code will block.
Am I approaching this with the wrong concept/architecture ?
Without managing your own thread, there has to be a way to use some non-blocking OS API on stdin and wrap it for tokio (tokio::io::Stdin 1.12 uses a blocking variant).
Otherwise if we follow the advice from the docs and spawn our own thread,
this is how it could be done:
fn start_reading_stdin_lines(
sender: tokio::sync::mpsc::Sender<String>,
runtime: tokio::runtime::Handle
) {
std::thread::spawn(move || {
let stdin = std::io::stdin();
let mut line_buf = String::new();
while let Ok(_) = stdin.read_line(&mut line_buf) {
let line = line_buf.trim_end().to_string();
line_buf.clear();
let sender2 = sender.clone();
runtime.spawn(async move {
let result = sender2.send(line).await;
if let Err(error) = result {
println!("start_reading_stdin_lines send error: {:?}", error);
}
});
}
});
}
fn start_activity_until_shutdown(watch_sender: tokio::sync::watch::Sender<bool>) {
tokio::spawn(async move {
tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
println!("exiting after a signal...");
let result = watch_sender.send(true);
if let Err(error) = result {
println!("watch_sender send error: {:?}", error);
}
});
}
async fn read_input(
mut line_receiver: tokio::sync::mpsc::Receiver<String>,
mut watch_receiver: tokio::sync::watch::Receiver<bool>
) {
loop {
tokio::select! {
Some(line) = line_receiver.recv() => {
println!("line: {}", line);
// process the input
match line.as_str() {
"exit" => {
println!("exiting manually...");
break;
},
"send" => {
println!("send_stuff");
}
unexpected_line => {
println!("unexpected command: {}", unexpected_line);
}
}
}
Ok(_) = watch_receiver.changed() => {
println!("shutdown");
break;
}
}
}
}
#[tokio::main]
async fn main() {
let (line_sender, line_receiver) = tokio::sync::mpsc::channel(1);
start_reading_stdin_lines(line_sender, tokio::runtime::Handle::current());
let (watch_sender, watch_receiver) = tokio::sync::watch::channel(false);
// this will send a shutdown signal at some point
start_activity_until_shutdown(watch_sender);
read_input(line_receiver, watch_receiver).await;
}
Potential improvements:
if you are ok with tokio_stream wrappers, this could be combined more elegantly with start_reading_stdin_lines producing a stream of lines and mapping them to typed commands. read_input could be then based on StreamMap instead of select!.
enabling experimental stdin_forwarders feature makes reading lines easier with a for loop over lines()
I have code that sends file modification events over a tokio channel. The messages originate in a dedicated thread (I think), and end up in a tokio thread. It works fine if the received is in the main tokio task but if I move it to a spawned task then for some reason it rx.next() immediately fails. Here's the code:
use futures::{
channel::mpsc::{channel, Receiver},
SinkExt, StreamExt,
};
use notify::{Event, RecommendedWatcher, RecursiveMode, Watcher};
use std::path::Path;
fn async_watcher() -> notify::Result<(RecommendedWatcher, Receiver<notify::Result<Event>>)> {
let (mut tx, rx) = channel(1);
let watcher = RecommendedWatcher::new(move |res| {
futures::executor::block_on(async {
tx.send(res).await.unwrap();
})
})?;
Ok((watcher, rx))
}
#[tokio::main]
async fn main() {
let path = std::env::args()
.nth(1)
.expect("Argument 1 needs to be a path");
println!("watching {}", path);
// watch_with_task(&path).await.unwrap();
watch_without_task(&path).await;
}
fn watch_with_task(path: &str) -> tokio::task::JoinHandle<()> {
let (mut watcher, mut rx) = async_watcher().unwrap();
watcher
.watch(Path::new(path), RecursiveMode::NonRecursive)
.unwrap();
tokio::spawn(async move {
eprintln!("Watch task spawned");
while let Some(res) = rx.next().await {
match res {
Ok(event) => eprintln!("File changed: {:?}", event),
Err(e) => eprintln!("Watch error: {:?}", e),
}
}
eprintln!("Watch task finished");
})
}
async fn watch_without_task(path: &str) {
let (mut watcher, mut rx) = async_watcher().unwrap();
watcher
.watch(Path::new(path), RecursiveMode::NonRecursive)
.unwrap();
while let Some(res) = rx.next().await {
match res {
Ok(event) => eprintln!("File changed: {:?}", event),
Err(e) => eprintln!("Watch error: {:?}", e),
}
}
}
Dependencies:
[dependencies]
futures = "0.3"
notify = "5.0.0-pre.13"
tokio = { version = "1.6", features = ["full"] }
With watch_without_task I get this output - it works:
watching /path/to/file
File changed: Event { kind: Modify(Metadata(Any)), paths: ["/path/to/file"], attr:tracker: None, attr:flag: None, attr:info: None, attr:source: None }
File changed: Event { kind: Modify(Data(Content)), paths: ["/path/to/file"], attr:tracker: None, attr:flag: None, attr:info: None, attr:source: None }
With watch_with_task I get this:
watching /path/to/file
Watch task spawned
Watch task finished
It exits immediately. Why?
Oooo actually I was wrong. The issue is that when using the task the watcher gets dropped immediately, whereas without the task it is kept alive by the while loop.
I'm trying to read from stdin in an async task, spawned with tokio::spawn. The
executor is crated as
let mut executor = tokio::runtime::current_thread::Runtime::new().unwrap();
the main task is then run with executor.task(...), which spawns other tasks
with tokio::spawn().
fn main then calls executor.run().unwrap(); to wait for all tasks to finish.
The problem is when I do
let mut stdin = tokio::io::stdin();
let mut read_buf: [u8; 1024] = [0; 1024];
...
stdin.read(&mut read_buf).await
I get "blocking annotated I/O must be called from the context of the Tokio runtime" error.
Dependencies:
futures-preview = { version = "0.3.0-alpha.18", features = ["async-await", "nightly"] }
tokio = "0.2.0-alpha.2"
tokio-net = "0.2.0-alpha.2"
tokio-sync = "0.2.0-alpha.2"
Full code:
extern crate futures;
extern crate tokio;
extern crate tokio_net;
extern crate tokio_sync;
use std::io::Write;
use std::net::SocketAddr;
use tokio::io::AsyncReadExt;
use tokio::net::tcp::split::{TcpStreamReadHalf, TcpStreamWriteHalf};
use tokio::net::TcpListener;
use tokio_sync::oneshot;
use futures::select;
use futures::future::FutureExt;
#[derive(Debug)]
enum AppErr {
CantBindAddr(std::io::Error),
CantAccept(std::io::Error),
}
fn main() {
let mut executor = tokio::runtime::current_thread::Runtime::new().unwrap();
executor.spawn(async {
match server_task().await {
Ok(()) => {}
Err(err) => {
println!("Error: {:?}", err);
}
}
});
executor.run().unwrap(); // ignores RunError
}
async fn server_task() -> Result<(), AppErr> {
let addr: SocketAddr = "127.0.0.1:8080".parse().unwrap();
let mut listener = TcpListener::bind(&addr).map_err(AppErr::CantBindAddr)?;
loop {
print!("Waiting for incoming connection...");
let _ = std::io::stdout().flush();
let (socket, _) = listener.accept().await.map_err(AppErr::CantAccept)?;
println!("{:?} connected.", socket);
let (read, write) = socket.split();
let (abort_in_task_snd, abort_in_task_rcv) = oneshot::channel();
let (abort_out_task_snd, abort_out_task_rcv) = oneshot::channel();
tokio::spawn(handle_incoming(read, abort_in_task_rcv, abort_out_task_snd));
tokio::spawn(handle_outgoing(
write,
abort_out_task_rcv,
abort_in_task_snd,
));
}
}
async fn handle_incoming(
mut conn: TcpStreamReadHalf,
abort_in: oneshot::Receiver<()>,
abort_out: oneshot::Sender<()>,
) {
println!("handle_incoming");
let mut read_buf: [u8; 1024] = [0; 1024];
let mut abort_in_fused = abort_in.fuse();
loop {
select! {
abort_ret = abort_in_fused => {
// TODO match abort_ret {..}
println!("abort signalled, handle_incoming returning");
return;
},
bytes = conn.read(&mut read_buf).fuse() => {
match bytes {
Err(io_err) => {
println!("io error when reading input stream: {:?}", io_err);
return;
}
Ok(bytes) => {
println!("read {} bytes: {:?}", bytes, &read_buf[0..bytes]);
}
}
}
}
}
}
async fn handle_outgoing(
conn: TcpStreamWriteHalf,
abort_in: oneshot::Receiver<()>,
abort_out: oneshot::Sender<()>,
) {
println!("handle_outgoing");
let mut stdin = tokio::io::stdin();
let mut read_buf: [u8; 1024] = [0; 1024];
let mut abort_in_fused = abort_in.fuse();
loop {
select! {
abort_ret = abort_in_fused => {
println!("abort signalled, handle_outgoing returning");
return;
}
input = stdin.read(&mut read_buf).fuse() => {
match input {
Err(io_err) => {
println!("io error when reading stdin: {:?}", io_err);
return;
}
Ok(bytes) => {
println!("handle_outgoing read {} bytes", bytes);
// TODO
}
}
},
}
}
}
Questions:
Am I doing task spawning right? Can I safely do tokio::spawn in the main
task passed to executor.spawn()?
What's wrong with the way I read stdin in this program?
Thanks
Tokio stdin blocks enclosing thread from the pool of executor because it is annotated with blocking from tokio-executor. From the reference :
When the blocking function enters, it hands off the responsibility
of processing the current work queue to another thread.
Your code is not working because the executor that you have used is multiplexing tasks in a single thread(tokio::runtime::current_thread::Runtime::new()). Thus there will remain no other thread to execute other tasks for executor.
If you properly configure your runtime(thread pool with multiple threads) your code will work fine :
fn main() {
let rt = tokio::runtime::Runtime::new().unwrap();
let mut executor = rt.executor();
executor.spawn(async {
match server_task().await {
Ok(()) => {}
Err(err) => {
println!("Error: {:?}", err);
}
}
});
rt.shutdown_on_idle();
}
See also: How can I stop reading from a tokio::io::lines stream?
I am dabbling in tokio-core and can figure out how to spawn an event loop. However there are two things i am not sure of - how to gracefully exit the event loop and how to exit a stream running inside an event loop. For e.g consider this simple piece of code which spawns two listeners into the event loop and waits for another thread to indicate an exit condition:
extern crate tokio_core;
extern crate futures;
use tokio_core::reactor::Core;
use futures::sync::mpsc::unbounded;
use tokio_core::net::TcpListener;
use std::net::SocketAddr;
use std::str::FromStr;
use futures::{Stream, Future};
use std::thread;
use std::time::Duration;
use std::sync::mpsc::channel;
fn main() {
let (get_tx, get_rx) = channel();
let j = thread::spawn(move || {
let mut core = Core::new().unwrap();
let (tx, rx) = unbounded();
get_tx.send(tx).unwrap(); // <<<<<<<<<<<<<<< (1)
// Listener-0
{
let l = TcpListener::bind(&SocketAddr::from_str("127.0.0.1:44444").unwrap(),
&core.handle())
.unwrap();
let fe = l.incoming()
.for_each(|(_sock, peer)| {
println!("Accepted from {}", peer);
Ok(())
})
.map_err(|e| println!("----- {:?}", e));
core.handle().spawn(fe);
}
// Listener1
{
let l = TcpListener::bind(&SocketAddr::from_str("127.0.0.1:55555").unwrap(),
&core.handle())
.unwrap();
let fe = l.incoming()
.for_each(|(_sock, peer)| {
println!("Accepted from {}", peer);
Ok(())
})
.map_err(|e| println!("----- {:?}", e));
core.handle().spawn(fe);
}
let work = rx.for_each(|v| {
if v {
// (3) I want to shut down listener-0 above the release the resources
Ok(())
} else {
Err(()) // <<<<<<<<<<<<<<< (2)
}
});
let _ = core.run(work);
println!("Exiting event loop thread");
});
let tx = get_rx.recv().unwrap();
thread::sleep(Duration::from_secs(2));
println!("Want to terminate listener-0"); // <<<<<< (3)
tx.send(true).unwrap();
thread::sleep(Duration::from_secs(2));
println!("Want to exit event loop");
tx.send(false).unwrap();
j.join().unwrap();
}
So say after the sleep in the main thread i want a clean exit of the event loop thread. Currently I send something to the event loop to make it exit and thus releasing the thread.
However both, (1) and (2) feel hacky - i am forcing an error as an exit condition. My questions are:
1) Am I doing it right ? If not then what is the correct way to gracefully exit the event loop thread.
2) I don't event know how to do (3) - i.e. indicate a condition externally to shutdown listener-0 and free all it's resources. How do i achieve this ?
The event loop (core) is not being turned any more (e.g. by run()) or is forgotten (drop()ed). There is no synchronous exit. core.run() returns and stops turning the loop when the Future passed to it completes.
A Stream completes by yielding None (marked with (3) in the code below).
When e.g. a TCP connection is closed the Stream representing it completes and the other way around.
extern crate tokio_core;
extern crate futures;
use tokio_core::reactor::Core;
use futures::sync::mpsc::unbounded;
use tokio_core::net::TcpListener;
use std::net::SocketAddr;
use std::str::FromStr;
use futures::{Async, Stream, Future, Poll};
use std::thread;
use std::time::Duration;
struct CompletionPact<S, C>
where S: Stream,
C: Stream,
{
stream: S,
completer: C,
}
fn stream_completion_pact<S, C>(s: S, c: C) -> CompletionPact<S, C>
where S: Stream,
C: Stream,
{
CompletionPact {
stream: s,
completer: c,
}
}
impl<S, C> Stream for CompletionPact<S, C>
where S: Stream,
C: Stream,
{
type Item = S::Item;
type Error = S::Error;
fn poll(&mut self) -> Poll<Option<S::Item>, S::Error> {
match self.completer.poll() {
Ok(Async::Ready(None)) |
Err(_) |
Ok(Async::Ready(Some(_))) => {
// We are done, forget us
Ok(Async::Ready(None)) // <<<<<< (3)
},
Ok(Async::NotReady) => {
self.stream.poll()
},
}
}
}
fn main() {
// unbounded() is the equivalent of a Stream made from a channel()
// directly create it in this thread instead of receiving a Sender
let (tx, rx) = unbounded::<()>();
// A second one to cause forgetting the listener
let (l0tx, l0rx) = unbounded::<()>();
let j = thread::spawn(move || {
let mut core = Core::new().unwrap();
// Listener-0
{
let l = TcpListener::bind(
&SocketAddr::from_str("127.0.0.1:44444").unwrap(),
&core.handle())
.unwrap();
// wrap the Stream of incoming connections (which usually doesn't
// complete) into a Stream that completes when the
// other side is drop()ed or sent on
let fe = stream_completion_pact(l.incoming(), l0rx)
.for_each(|(_sock, peer)| {
println!("Accepted from {}", peer);
Ok(())
})
.map_err(|e| println!("----- {:?}", e));
core.handle().spawn(fe);
}
// Listener1
{
let l = TcpListener::bind(
&SocketAddr::from_str("127.0.0.1:55555").unwrap(),
&core.handle())
.unwrap();
let fe = l.incoming()
.for_each(|(_sock, peer)| {
println!("Accepted from {}", peer);
Ok(())
})
.map_err(|e| println!("----- {:?}", e));
core.handle().spawn(fe);
}
let _ = core.run(rx.into_future());
println!("Exiting event loop thread");
});
thread::sleep(Duration::from_secs(2));
println!("Want to terminate listener-0");
// A drop() will result in the rx side Stream being completed,
// which is indicated by Ok(Async::Ready(None)).
// Our wrapper behaves the same when something is received.
// When the event loop encounters a
// Stream that is complete it forgets about it. Which propagates to a
// drop() that close()es the file descriptor, which closes the port if
// nothing else uses it.
l0tx.send(()).unwrap(); // alternatively: drop(l0tx);
// Note that this is async and is only the signal
// that starts the forgetting.
thread::sleep(Duration::from_secs(2));
println!("Want to exit event loop");
// Same concept. The reception or drop() will cause Stream completion.
// A completed Future will cause run() to return.
tx.send(()).unwrap();
j.join().unwrap();
}
I implemented graceful shutdown via a oneshot channel.
The trick was to use both a oneshot channel to cancel the tcp listener, and use a select! on the two futures. Note I'm using tokio 0.2 and futures 0.3 in the example below.
use futures::channel::oneshot;
use futures::{FutureExt, StreamExt};
use std::thread;
use tokio::net::TcpListener;
pub struct ServerHandle {
// This is the thread in which the server will block
thread: thread::JoinHandle<()>,
// This switch can be used to trigger shutdown of the server.
kill_switch: oneshot::Sender<()>,
}
impl ServerHandle {
pub fn stop(self) {
self.kill_switch.send(()).unwrap();
self.thread.join().unwrap();
}
}
pub fn run_server() -> ServerHandle {
let (kill_switch, kill_switch_receiver) = oneshot::channel::<()>();
let thread = thread::spawn(move || {
info!("Server thread begun!!!");
let mut runtime = tokio::runtime::Builder::new()
.basic_scheduler()
.enable_all()
.thread_name("Tokio-server-thread")
.build()
.unwrap();
runtime.block_on(async {
server_prog(kill_switch_receiver).await.unwrap();
});
info!("Server finished!!!");
});
ServerHandle {
thread,
kill_switch,
}
}
async fn server_prog(kill_switch_receiver: oneshot::Receiver<()>) -> std::io::Result<()> {
let addr = "127.0.0.1:12345";
let addr: std::net::SocketAddr = addr.parse().unwrap();
let mut listener = TcpListener::bind(&addr).await?;
let mut kill_switch_receiver = kill_switch_receiver.fuse();
let mut incoming = listener.incoming().fuse();
loop {
futures::select! {
x = kill_switch_receiver => {
break;
},
optional_new_client = incoming.next() => {
if let Some(new_client) = optional_new_client {
let peer_socket = new_client?;
info!("Client connected!");
let peer = process_client(peer_socket, db.clone());
peers.lock().unwrap().push(peer);
} else {
info!("No more incoming connections.");
break;
}
},
};
}
Ok(())
}
Hopes this helps others (or future me ;)).
My code lives here:
https://github.com/windelbouwman/lognplot/blob/master/lognplot/src/server/server.rs