How to share tokio::net::TcpStream act on it concurrently?

Question

I have a requirement to send and receive normal data on the same TcpStream, while sending heartbeat data at regular intervals. In the current implementation, I used Arc>, but it compiled with errors:

use anyhow::Result;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

#[tokio::main]
async fn main() -> Result<()> {
    let stream = TcpStream::connect("127.0.0.1:8888").await.unwrap();
    let stream = Arc::new(Mutex::new(stream));

    let common_stream = stream.clone();
    let handler1 = tokio::spawn(async {
        loop {
            let mut stream = common_stream.lock().unwrap();
            let mut buf = [0u8; 10];
            stream.read_exact(&mut buf).await.unwrap();
            buf.reverse();
            stream.write(&buf).await.unwrap();
        }
    });

    let heartbeat_stream = stream.clone();
    let handler2 = tokio::spawn(async {
        loop {
            let mut stream = heartbeat_stream.lock().unwrap();
            stream.write_u8(1).await.unwrap();

            thread::sleep(Duration::from_millis(200));
        }
    });

    handler1.await?;
    handler2.await?;

    Ok(())
}

error: future cannot be sent between threads safely
   --> src\main.rs:14:20
    |
14  |     let handler1 = tokio::spawn(async {
    |                    ^^^^^^^^^^^^ future created by async block is not `Send`
    |
    = help: within `impl Future`, the trait `Send` is not implemented for `std::sync::MutexGuard<'_, tokio::net::TcpStream>`
note: future is not `Send` as this value is used across an await
   --> src\main.rs:20:31
    |
16  |             let mut stream = common_stream.lock().unwrap();
    |                 ---------- has type `std::sync::MutexGuard<'_, tokio::net::TcpStream>` which is not `Send`
...
20  |             stream.write(&buf).await.unwrap();
    |                               ^^^^^^ await occurs here, with `mut stream` maybe used later
21  |         }
    |         - `mut stream` is later dropped here
note: required by a bound in `tokio::spawn`
   --> .cargo
egistry\src\mirrors.tuna.tsinghua.edu.cn-df7c3c540f42cdbd	okio-1.17.0\src	ask\spawn.rs:127:21
    |
127 |         T: Future + Send + 'static,
    |                     ^^^^ required by this bound in `tokio::spawn`

error: future cannot be sent between threads safely
   --> src\main.rs:25:20
    |
25  |     let handler2 = tokio::spawn(async {
    |                    ^^^^^^^^^^^^ future created by async block is not `Send`
    |
    = help: within `impl Future`, the trait `Send` is not implemented for `std::sync::MutexGuard<'_, tokio::net::TcpStream>`
note: future is not `Send` as this value is used across an await
   --> src\main.rs:28:31
    |
27  |             let mut stream = heartbeat_stream.lock().unwrap();
    |                 ---------- has type `std::sync::MutexGuard<'_, tokio::net::TcpStream>` which is not `Send`
28  |             stream.write_u8(1).await.unwrap();
    |                               ^^^^^^ await occurs here, with `mut stream` maybe used later
...
31  |         }
    |         - `mut stream` is later dropped here
note: required by a bound in `tokio::spawn`
   --> .cargo
egistry\src\mirrors.tuna.tsinghua.edu.cn-df7c3c540f42cdbd	okio-1.17.0\src	ask\spawn.rs:127:21
    |
127 |         T: Future + Send + 'static,
    |                     ^^^^ required by this bound in `tokio::spawn`

How can these errors be fixed or is there another way to achieve the same goal?

user_3420151 · Answer

Here is a solution that splits the stream to two parts for reading and writing plus does in a loop:

waiting for heartbeat events and sends a byte to write half of stream when this happens
waits data from read half (10 bytes), reverses it and writes again to write half

Also this does not spawn threads and does everything nicely in current one without locks.

use anyhow::Result;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

#[tokio::main]
async fn main() -> Result<()> {
    let mut stream = TcpStream::connect("127.0.0.1:8888").await?;
    let (mut read, mut write) = stream.split();
    let mut heartbeat_interval = tokio::time::interval(Duration::from_millis(200));
    let mut buf = [0u8; 10];

    loop {
        tokio::select! {
            _ = heartbeat_interval.tick() => {
                write.write(&[1]).await?;
            }

            result = read.read_exact(&mut buf) => {
                let _bytes_read = result?;
                buf.reverse();
                write.write(&buf).await?;
            }
        }
    }
}

user_1695172 · Answer

Several errors in your code, although the idea behind it is almost good. You should use any available tool in async as possible. Some of the needed/desired changes:

Use tokio::time::sleep because it is async, otherwise the call is blocking
Use an async version of mutex (the one from futures crate for example)
Use some kind of generic error handling (anyhow would help)

use futures::lock::Mutex;
use anyhow::Error;
use tokio::time::sleep;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let stream = TcpStream::connect("127.0.0.1:8888").await.unwrap();
    let stream = Arc::new(Mutex::new(stream));

    let common_stream = stream.clone();
    let handler1 = tokio::spawn(async move {
        loop {
            let mut stream = common_stream.lock().await;
            let mut buf = [0u8; 10];
            stream.read_exact(&mut buf).await.unwrap();
            buf.reverse();
            stream.write(&buf).await.unwrap();
        }
    });

    let heartbeat_stream = stream.clone();
    let handler2 = tokio::spawn(async move {
        loop {
            let mut stream = heartbeat_stream.lock().await;
            stream.write_u8(1).await.unwrap();

            sleep(Duration::from_millis(200)).await;
        }
    });

    handler1.await?;
    handler2.await?;

    Ok(())
}

Playground

How to share tokio::net::TcpStream act on it concurrently?

2 Answers