LocalSet pool to handle multiple tcp connection concurrently

[yongman]

struct LocalSpawner {
   send: mpsc::UnboundedSender<Handler>,
}

impl LocalSpawner {
    pub fn new() -> Self {
        let (send, mut recv) = mpsc::unbounded_channel::<Handler>();
        
        std::thread::spawn(move || {
            let local = LocalSet::new();

            let rt = Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();

            local.spawn_local(async move {
                while let Some(mut handler) = recv.recv().await {
                        // Process the connection. If an error is encountered, log it.
                        // handle.run will be loop runing inside until connection close
                        if let Err(err) = handler.run().await {
                            println!("Connection Error: {:?}", &err);
                            error!(cause = ?err, "connection error");
                        }
                }
                // If the while loop returns, then all the LocalSpawner
                // objects have have been dropped.
            });

            rt.block_on(local);
        });

        Self {
            send,
        }
    }

    pub fn spawn(&self, handler: Handler) {
        self.send.send(handler).expect("Thread with LocalSet has shut down.");
    }
}

As the code example in tokio docs, the handler is the tcp connection accepted by main thread, I want to handle thousands of connections at the same time.
When I run the code, it will block when handling one connection, can not accept new connection until the previous one disconnected.
I can use one LocalSet instance with run_until to loop run accept and handle connection, but I want to use multiple LocalSet workers and balance the connections to them to use multiple cpu core.
I want to know the tokio usage pattern to implement this, have any suggetions?

[Darksonn]

You should use tokio::task::spawn_local inside your loop whenever you receive a new connection.

[yongman]

Thanks for your reply.

        loop {
            let handler = accept new connections...
            // Use localset to spawn task, because of mlua async function is not `Send`
            task::spawn_local(async move {
                // Process the connection. If an error is encountered, log it.
                if let Err(err) = handler.run().await {
                    println!("Connection Error: {:?}", &err);
                    error!(cause = ?err, "connection error");
                }
            });
        }

It works in this way, but the performance is bad, it can use just one cpu core. Are there any ways to balance connections to multiple LocalSet?

[Darksonn]

Tokio-util has an utility for balancing tasks on multiple LocalSets, which you can find here.

[yongman]

I have noticed this tokio utility and tried it, but I can not spawn_pinned tasks and await it inner the loop which will block the main thread.
How can I await the JoinHandle and accept loop at the same time?

[Darksonn]

You don't need to await the JoinHandle. The task will run even if you don't.

[yongman]

I have tried it and worked as expected.
Thanks a lot. ❤️