eval_stack/
judge.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
use std::{
    fs, future::Future, io::Read, os::unix::process::ExitStatusExt, path::PathBuf, task::Poll,
    time::Duration,
};

use anyhow::Result;

use crate::utils::get_memory_usage;

#[derive(Debug)]
pub enum JudgeStatus {
    Accepted,
    WrongAnswer,
    TimeLimitExceeded,
    MemoryLimitExceeded,
    RuntimeError,
    CompileError { message: String },
    SystemError { code: i32 },
    SegmentFault,
}

#[derive(Debug)]
pub struct JudgeResult {
    pub status: JudgeStatus,
    pub time_used: Duration,
    pub memory_used: u64,
}

pub struct Judge {
    pub child: std::process::Child,
    pub id: u32,
    pub time_limit: Duration,
    pub memory_limit: u64,
    pub instant: tokio::time::Instant,
    pub memory_used: u64,
    pub time_used: Duration,
    pub stdout_file: PathBuf,
    pub expected_output_file: PathBuf,
}

impl Future for Judge {
    type Output = Result<JudgeResult>;

    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        match self.child.try_wait()? {
            Some(status) => {
                self.time_used = self.instant.elapsed();
                drop(self.child.stdin.take());
                drop(self.child.stdout.take());
                if status.success() {
                    let mut stdout = fs::File::open(&self.stdout_file)?;
                    let mut expected_out = fs::File::open(&self.expected_output_file)?;

                    let mut output = String::new();
                    let mut expected_output = String::new();
                    stdout.read_to_string(&mut output)?;
                    expected_out.read_to_string(&mut expected_output)?;

                    if output.trim_end_matches(|c: char| c.is_whitespace() || c == '\n')
                        == expected_output
                            .trim_end_matches(|c: char| c.is_whitespace() || c == '\n')
                    {
                        Poll::Ready(Ok(JudgeResult {
                            status: JudgeStatus::Accepted,
                            time_used: self.time_used,
                            memory_used: self.memory_used,
                        }))
                    } else {
                        Poll::Ready(Ok(JudgeResult {
                            status: JudgeStatus::WrongAnswer,
                            time_used: self.time_used,
                            memory_used: self.memory_used,
                        }))
                    }
                } else {
                    match status.signal() {
                        Some(libc::SIGSEGV) | Some(libc::SIGBUS) | Some(libc::SIGILL) => {
                            Poll::Ready(Ok(JudgeResult {
                                status: JudgeStatus::SegmentFault,
                                time_used: self.time_used,
                                memory_used: self.memory_used,
                            }))
                        }
                        Some(code) => Poll::Ready(Ok(JudgeResult {
                            status: JudgeStatus::SystemError { code },
                            time_used: self.time_used,
                            memory_used: self.memory_used,
                        })),
                        None => Poll::Ready(Ok(JudgeResult {
                            status: JudgeStatus::RuntimeError,
                            time_used: self.time_used,
                            memory_used: self.memory_used,
                        })),
                    }
                }
            }
            None => {
                if let Some(memory_used) = get_memory_usage(self.id) {
                    self.memory_used = memory_used.max(self.memory_used);
                };
                if self.memory_used > self.memory_limit {
                    self.child.kill()?;
                    self.time_used = self.instant.elapsed();
                    return Poll::Ready(Ok(JudgeResult {
                        status: JudgeStatus::MemoryLimitExceeded,
                        time_used: self.time_used,
                        memory_used: self.memory_used,
                    }));
                }
                if self.instant.elapsed() > self.time_limit {
                    self.child.kill()?;
                    self.time_used = self.instant.elapsed();
                    return Poll::Ready(Ok(JudgeResult {
                        status: JudgeStatus::TimeLimitExceeded,
                        time_used: self.time_used,
                        memory_used: self.memory_used,
                    }));
                }
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }
}