Non-Blocking Execution of Concurrent Tasks in Tokio with Custom Thread Pool

[Hojun-Son]

I am trying to run two tasks using Tokio. Since I need to limit the number of threads, I initialize the runtime with a Builder. I use the block_on method to run both tasks. These two tasks are channel_loop and a loop that handles I/O requests. These tasks should not block each other.

In the I/O loop, data is sent through a priority channel, and the channel_loop task receives this data. Inside channel_loop, I create additional tasks using tokio::spawn(process_data()). My goal is to run multiple process_data() tasks asynchronously. Here is my code:

fn main() -> Result<(), Box<dyn std::error::Error>> {
    //...//

    let num_threads = 4;

    let runtime = Builder::new_multi_thread()
        .worker_threads(num_threads)
        .enable_all()
        .build()
        .unwrap();

    let _ = runtime.block_on(async { 
        tokio::spawn(channel_loop( 
            rx.clone(),
            result_queue.clone(),
            loopback,
            prompt_args.clone(),
            tx.clone(),
        ));

        loop {
            // This is an I/O loop
        }
           
    });
    Ok(())
}

pub async fn channel_loop(
    rx: Receiver<DataForChannel, u8>,
    result_queue: Arc<Mutex<VecDeque<FinalResult>>>,
    loopback: SocketAddr,
    prompt_args: PromptArgs,
    tx: Sender<DataForChannel, u8>,
) -> Result<()> {
    let doc_sum_vec: Arc<Mutex<Vec<DocChunk>>> = Arc::new(Mutex::new(Vec::new()));

    let loopback_socket = match TcpStream::connect(loopback) {
        Ok(stream) => stream,
        Err(err) => {
            error!("failed to connect loopback socket! : {:?}", err);
            return Ok(());
        }
    };
    let loopback_socket: Arc<Mutex<TcpStream>> = Arc::new(Mutex::new(loopback_socket));
    loop {
        match rx.try_recv() {
            core::result::Result::Ok(data) => {
                let doc_sum_vec = doc_sum_vec.clone();
                let result_queue = result_queue.clone();
                let loopback_socket = loopback_socket.clone();
                let prompt_args = prompt_args.clone();
                let tx = tx.clone();

                tokio::spawn(async move {
                    if let Err(e) = process_data(
                        data,
                        doc_sum_vec,
                        result_queue,
                        loopback_socket,
                        prompt_args,
                        tx,
                    )
                    .await
                    {
                        error!("Error processing data: {:?}", e);
                    }
                });
            }
            Err(async_priority_channel::TryRecvError::Empty) => {
                // debug!("TryRecvError::Empty!");
                tokio::time::sleep(std::time::Duration::from_millis(1)).await;
                continue;
            }
            Err(async_priority_channel::TryRecvError::Closed) => {
                // debug!("TryRecvError::Closed!");
                continue;
            }
        }
    }
}

How can I ensure that the channel_loop and the loop handling I/O requests do not block each other while running multiple process_data() tasks asynchronously?

[Darksonn]

You have not posted complete code, but I notice some issues immediately:

1. The call to TcpStream::connect has no .await. This means that you are not using an async TcpStream. That is a mistake.
2. You're calling try_recv in a loop. This is also always wrong. Use the async recv method instead.
3. When you get a closed error from the channel, you should exit. Do not use continue. Closed errors do not go away - channels are closed forever.

Consider reading Async: What is blocking?

[Hojun-Son]

Thank you for your help.

Is there any issue with using the block_on method to create the I/O loop and the channel_loop task? I'm wondering if they will not block each other.

[Darksonn]

Your use of block_on is okay. However, your code will probably have better performance if you instead move the I/O loop inside a tokio::spawned task.