Rust实现Protohackers Prime Time多客户端测试超时求助

我正在学习Rust，尝试完成Protohackers的Prime Time挑战，单客户端测试已通过，但多客户端测试出现60秒超时失败，测试日志如下：

[Mon Nov 28 14:21:10 2022 UTC] [0simple.test] NOTE:check starts
[Mon Nov 28 14:21:10 2022 UTC] [0simple.test] NOTE:connected to 213.89.253.71 port 5006
[Mon Nov 28 14:21:13 2022 UTC] [0simple.test] NOTE:handled all 25 requests correctly
[Mon Nov 28 14:21:13 2022 UTC] [0simple.test] PASS
[Mon Nov 28 14:21:14 2022 UTC] [1multiclient.test] NOTE:check starts
[Mon Nov 28 14:21:15 2022 UTC] [1multiclient.test] NOTE:connected to 213.89.253.71 port 5006
[Mon Nov 28 14:21:15 2022 UTC] [1multiclient.test] NOTE:connected to 213.89.253.71 port 5006
[Mon Nov 28 14:21:15 2022 UTC] [1multiclient.test] NOTE:connected to 213.89.253.71 port 5006
[Mon Nov 28 14:21:15 2022 UTC] [1multiclient.test] NOTE:connected to 213.89.253.71 port 5006
[Mon Nov 28 14:21:15 2022 UTC] [1multiclient.test] NOTE:connected to 213.89.253.71 port 5006
[Mon Nov 28 14:22:14 2022 UTC] [1multiclient.test] FAIL:alarm timeout after 60 seconds

我的实现代码：

use std::{net::TcpListener, io::{Result, Read}, thread};
use serde_json::json;
use serde::{Deserialize};
use std::io::Write;

#[derive(Deserialize, Debug)]
struct Message {
    method: String,
    number: i64
}

impl Message {
    fn is_ok(&self) -> bool {
        if self.method == "isPrime"{
            return true;
        }
        return false; 
    }

    fn is_prime(&self) -> bool {
        if self.number <= 1{
            return false; 
        }
        for i in  2..self.number {
            if self.number % i ==0 {
                let res = self.number % i; 
                println!("{res}"); 
                return false; 
            }
        }
        return true; 
    }

    fn get_variables(&self){
        println!("Method is: {0}, Number is: {1} ", self.method, self.number);
    }
    
}

fn main() -> Result<()>{
    let addr= "192.168.0.9:5006"; 
    let listener= TcpListener::bind(addr)?; 
    for stream in listener.incoming(){
        let _ :thread::JoinHandle<Result<()>> = thread::spawn(move || {
            loop {
                let mut m = serde_json::Deserializer::from_reader(stream.as_ref().unwrap());
                let msg = Message::deserialize(&mut m);

                match msg{
                    Ok(msg)=>{
                        if msg.is_ok(){
                            let ret_msg = json!({"method":"isPrime","prime":msg.is_prime()});
                            serde_json::to_writer(stream.as_ref().unwrap(), &ret_msg).unwrap();
                            stream.as_ref().expect("Error").write(b"\n").unwrap(); 
                        }
                     }, 
                    _ => {
                        break Ok(()); 
                    }
                }
            }
        });
    }
    Ok(())
}

问题分析与修复方案 #

1. 质数判断效率极低（核心超时原因）

原is_prime函数遍历从2到目标数的所有整数，对于大数字（比如接近i64上限的数）会花费数十秒甚至更久，多客户端场景下只要有一个请求卡住，就会触发测试超时。

优化后的质数判断逻辑：

• 直接排除小于等于1的数、偶数（除了2）
• 仅遍历到目标数的平方根，且只检查奇数，大幅减少循环次数

fn is_prime(&self) -> bool {
    match self.number {
        n if n <= 1 => false,
        2 => true,
        n if n % 2 == 0 => false,
        _ => {
            let sqrt_n = (self.number as f64).sqrt() as i64;
            (3..=sqrt_n).step_by(2).all(|i| self.number % i != 0)
        }
    }
}

2. Stream可变借用错误

serde_json::Deserializer::from_reader需要一个实现Read的可变引用，原代码中stream.as_ref().unwrap()得到的是不可变引用，无法正确读取数据流，多客户端下会引发读取异常。

3. JSON反序列化方式错误

每次循环创建新的Deserializer，无法正确处理同一个TCP连接中的多个JSON请求。应该复用Deserializer，通过连续反序列化来逐个读取请求。

4. 未刷新输出缓冲区

写入响应后没有调用flush()，可能导致响应滞留在缓冲区，客户端无法及时收到数据，引发超时。

修正后的完整代码 #

use std::{net::TcpListener, io::{Result, Read, Write}, thread};
use serde_json::json;
use serde::{Deserialize};

#[derive(Deserialize, Debug)]
struct Message {
    method: String,
    number: i64
}

impl Message {
    fn is_prime(&self) -> bool {
        match self.number {
            n if n <= 1 => false,
            2 => true,
            n if n % 2 == 0 => false,
            _ => {
                let sqrt_n = (self.number as f64).sqrt() as i64;
                (3..=sqrt_n).step_by(2).all(|i| self.number % i != 0)
            }
        }
    }
}

fn main() -> Result<()>{
    let addr = "192.168.0.9:5006"; 
    let listener = TcpListener::bind(addr)?; 
    
    for stream_result in listener.incoming() {
        let _ = thread::spawn(move || -> Result<()> {
            let mut stream = stream_result?;
            let mut deserializer = serde_json::Deserializer::from_reader(&mut stream);
            
            loop {
                match Message::deserialize(&mut deserializer) {
                    Ok(msg) => {
                        if msg.method == "isPrime" {
                            let prime = msg.is_prime();
                            let ret_msg = json!({"method": "isPrime", "prime": prime});
                            serde_json::to_writer(&mut stream, &ret_msg)?;
                            stream.write_all(b"\n")?;
                            stream.flush()?;
                        }
                    }
                    Err(_) => {
                        // 连接断开或解析失败时退出循环
                        break Ok(());
                    }
                }
            }
        });
    }
    
    Ok(())
}

内容的提问来源于stack exchange，提问作者HeyGoogle