Loop over Rust collection (borrowed immutably) then mutate it afterward

Question

I'm having trouble while trying to manipulate a HashMap in Rust.

use std::os::fd::AsRawFd;
use std::collections::HashMap;
use std::collections::HashSet;
use tokio::net::TcpListener;
use tokio::sync::mpsc;

#[tokio::main]
async fn main() -> std::io::Result<()>{
    let (tx, mut rx) = mpsc::channel(32);
    let listener = TcpListener::bind("127.0.0.1:1234").await?;
    
    tokio::spawn(async move {
        while let Ok((stream, _)) = listener.accept().await {
            tx.send(stream).await;
        }
    });

    let mut should_close = HashSet::new();
    let mut clients = HashMap::new();
    let mut buf = [0u8; 1024];
    loop {
        if let Ok(stream) = rx.try_recv() {
            // clients.entry(stream.as_raw_fd()).or_insert(stream);
            clients.insert(stream.as_raw_fd(), stream);
        }
        should_close.clear();
        
        for (fd, stream) in clients.iter() {
            match stream.try_read(&mut buf) {
                Ok(0) => {
                    println!("Client {fd} disconnected");
                    should_close.insert(fd);
                },
                Ok(_) => {
                    let s = String::from_utf8(buf.to_vec()).unwrap();
                    println!("Client {fd} says: {}", s);
                    stream.try_write(format!("Thank you for saying: {s}").as_bytes());
                },
                Err(e) if e.kind() == tokio::io::ErrorKind::WouldBlock => continue,
                Err(e) => {
                    println!("Client {fd} error: {e}");
                    should_close.insert(fd);
                }
            }
        }
        for fd in &should_close {
            // clients.remove_entry(fd);
            clients.remove(fd);
        }
        
    }

Playground

Since this is a TCP server, I don't want to spawn one task for each connected client and allocate a dedicated buffer for it, so I'm keeping a list of clients being stored in variable clients. In side the infinite loop, if rx receives a new client from the listener, this client is added to clients. Then I loop over this map to read the clients' messages in a non-blocking fashion, keeping track of which clients should be untracked in the should_close variable. Once done looping over clients I loop over should_close to remove elements from clients.

The compiler won't let me borrow mutably and immutably inter-changeably. Here is the error produced:

error[E0502]: cannot borrow `clients` as mutable because it is also borrowed as immutable
  --> src/main.rs:21:13
   |
21 |             clients.insert(stream.as_raw_fd(), stream);
   |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
22 |         }
23 |         should_close.clear();
   |         -------------------- immutable borrow later used here
24 |         
25 |         for (fd, stream) in clients.iter() {
   |                             -------------- immutable borrow occurs here

Wrapping the map inside a RefCell or even Rc<RefCell> doesn't work:

error[E0716]: temporary value dropped while borrowed
  --> src/main.rs:25:29
   |
25 |         for (fd, stream) in clients.borrow().iter() {
   |                             ^^^^^^^^^^^^^^^^ creates a temporary value which is freed while still in use
...
42 |         }
   |         - temporary value is freed at the end of this statement
43 |         for fd in &should_close {
   |                   ------------- borrow later used here
   |
   = note: consider using a `let` binding to create a longer lived value

The compiler's suggestion to create a let outside of the loop doesn't work, because borrowing mutable inside the loop to mutate the map would cause panic.

What is the correct way to do this?

Answer 1

You problems happen because your hash set should_close store references to the keys of clients. Try to store descriptors by value:

let mut clients: HashSet<RawFd> = HashMap::new();
...
for (&fd, stream) in clients.iter() {
   ...
}

Should work.

In such cases, it is often insightful to write variables types fully when declare them. Those inferred variable types can confuse.