How can I implement a timeout for closing a TCP connection using TcpStream in Rust?

I have a function that connects to a TCP endpoint, sends a message, and then waits for two messages. After the function completes, I want to close the connection. However, I would like to add a timeout to the entire function so that if the message exchange doesn't finish within a specified time, the connection will be disconnected. What is the ideal way to achieve this in Rust with the TcpStream?

Currently, I am using the timeout function from tokio https://docs.rs/tokio/latest/tokio/time/fn.timeout.html. However, it doesn't appear to be the optimal choice for my use case.

use std::time::Duration;

use tokio::{net::TcpStream, time::timeout};

#[tokio::main]
async fn main() {
    connect().await;
}

async fn connect() {
    let conn = TcpStream::connect("127.0.0.1:7000").await.unwrap();

    // drop the connection within 2 seconds
    let _ = timeout(Duration::from_secs(2), do_something(conn)).await;
}
async fn do_something(conn: TcpStream) {
    // do something
}

Solution

For blocking streams, a timeout is usually set for a single connect, read or write operation. Non-blocking steams don't have an analogue at the syscall level. Using tokio::time::timeout() for each connect/read/write should give you roughly the same behaviour as timeouts for the corresponding blocking operations. You can also use tokio_io_timeout::TimeoutStream to configure timeouts for all read and write operations.