task.rs

use crate::mm::{
    MemorySet,
    PhysPageNum,
    KERNEL_SPACE, 
    VirtAddr,
    translated_refmut,
};
//use crate::timer::get_time_ms;
use crate::trap::{TrapContext, trap_handler};
use crate::config::{TRAP_CONTEXT};
use super::TaskContext;
use super::{PidHandle, pid_alloc, KernelStack};
use alloc::sync::{Weak, Arc};
use alloc::vec;
use alloc::vec::Vec;
use alloc::string::{String, ToString};
use spin::{Mutex, MutexGuard};
use crate::fs::{File, ProcInfo, Stdin, Stdout};
#[derive(Copy, Clone)]
pub struct SyscallCount{
    pub syscall_dup: (usize, usize),
    pub syscall_open: (usize, usize),
    pub syscall_close: (usize, usize),
    pub syscall_pipe: (usize, usize),
    pub syscall_read: (usize, usize),
    pub syscall_write: (usize, usize),
    pub syscall_exit: (usize, usize),
    pub syscall_yield: (usize, usize),
    pub syscall_get_time: (usize, usize),
    pub syscall_getpid: (usize, usize),
    pub syscall_fork: (usize, usize),
    pub syscall_exec: (usize, usize),
    pub syscall_waitpid: (usize, usize),
}

pub struct TaskControlBlock {
    // immutable
    pub pid: PidHandle,
    pub kernel_stack: KernelStack,
    // mutable
    inner: Mutex<TaskControlBlockInner>,
    pub ppid: isize,
}

pub struct TaskControlBlockInner {
    pub trap_cx_ppn: PhysPageNum,
    pub base_size: usize,
    pub task_cx_ptr: usize,
    pub task_status: TaskStatus,
    pub name: String,
    pub memory_set: MemorySet,
    pub parent: Option<Weak<TaskControlBlock>>,
    pub children: Vec<Arc<TaskControlBlock>>,
    pub exit_code: i32,
    pub fd_table: Vec<Option<Arc<dyn File + Send + Sync>>>,
    pub start_time: usize,
    pub cpu_time: usize,
    pub syscall: Arc<Mutex<SyscallCount>>,
}

impl TaskControlBlockInner {
    pub fn get_task_cx_ptr2(&self) -> *const usize {
        &self.task_cx_ptr as *const usize
    }
    pub fn get_trap_cx(&self) -> &'static mut TrapContext {
        self.trap_cx_ppn.get_mut()
    }
    pub fn get_user_token(&self) -> usize {
        self.memory_set.token()
    }
    pub fn get_status(&self) -> TaskStatus {
        self.task_status
    }
    pub fn is_zombie(&self) -> bool {
        self.get_status() == TaskStatus::Zombie
    }
    pub fn alloc_fd(&mut self) -> usize {
        if let Some(fd) = (0..self.fd_table.len())
            .find(|fd| self.fd_table[*fd].is_none()) {
            fd
        } else {
            self.fd_table.push(None);
            self.fd_table.len() - 1
        }
    }
    /*pub fn stamp(&mut self){
        self.start_run_stamp = get_time_ms();
    }
    pub fn cpu_time_update(&mut self){
        self.cpu_time += get_time_ms()-self.start_run_stamp;
    }*/
}

impl TaskControlBlock {
    pub fn get_info(&self)-> ProcInfo{
        let pid = self.getpid();
        let status = self.inner.lock().get_status();
        let name = self.get_name();
        let ppid = self.get_ppid();
        let cpu_time = self.get_cpu_time();
        let syscall_cnt = self.get_syscall_cnt();
        ProcInfo{
            pid,
            name,
            status,
            ppid,
            cpu_time,
            syscall_cnt,
        }
    }

    pub fn acquire_inner_lock(&self) -> MutexGuard<TaskControlBlockInner> {
        self.inner.lock()
    }

    pub fn set_name(&self, name: String){
        self.acquire_inner_lock().name = name;
    }

    pub fn new(elf_data: &[u8]) -> Self {
        // memory_set with elf program headers/trampoline/trap context/user stack
        let (memory_set, user_sp, entry_point) = MemorySet::from_elf(elf_data);
        let trap_cx_ppn = memory_set
            .translate(VirtAddr::from(TRAP_CONTEXT).into())
            .unwrap()
            .ppn();
        // alloc a pid and a kernel stack in kernel space
        let pid_handle = pid_alloc();
        let kernel_stack = KernelStack::new(&pid_handle);
        let kernel_stack_top = kernel_stack.get_top();
        let name = String::new();
        // push a task context which goes to trap_return to the top of kernel stack
        let task_cx_ptr = kernel_stack.push_on_top(TaskContext::goto_trap_return());
        let task_control_block = Self {
            pid: pid_handle,
            kernel_stack,
            ppid: -1,
            inner: Mutex::new(TaskControlBlockInner {
                trap_cx_ppn,
                base_size: user_sp,
                task_cx_ptr: task_cx_ptr as usize,
                task_status: TaskStatus::Ready,
                name,
                memory_set,
                parent: None,
                children: Vec::new(),
                exit_code: 0,
                fd_table: vec![
                    // 0 -> stdin
                    Some(Arc::new(Stdin)),
                    // 1 -> stdout
                    Some(Arc::new(Stdout)),
                    // 2 -> stderr
                    Some(Arc::new(Stdout)),
                    // ? -> proc_info
                    // 注意调用号严格来说不是自己分配的，打开一个新文件默认分配能分配的最小调用号
                    // Some(Arc::new(ProcInfoList)),
                ],
                start_time: 0,
                cpu_time: 0,
                syscall: Arc::new(Mutex::new(SyscallCount{
                    syscall_dup: (0, 0),
                    syscall_open: (0, 0),
                    syscall_close: (0, 0),
                    syscall_pipe: (0, 0),
                    syscall_read: (0, 0),
                    syscall_write: (0, 0),
                    syscall_exit: (0, 0),
                    syscall_yield: (0, 0),
                    syscall_get_time: (0, 0),
                    syscall_getpid: (0, 0),
                    syscall_fork: (0, 0),
                    syscall_exec: (0, 0),
                    syscall_waitpid: (0, 0),
                }))
                //start_run_stamp:0,
            }),
        };
        // prepare TrapContext in user space
        let trap_cx = task_control_block.acquire_inner_lock().get_trap_cx();
        *trap_cx = TrapContext::app_init_context(
            entry_point,
            user_sp,
            KERNEL_SPACE.lock().token(),
            kernel_stack_top,
            trap_handler as usize,
        );
        task_control_block
    }
    pub fn exec(&self, elf_data: &[u8], args: Vec<String>) {
        // memory_set with elf program headers/trampoline/trap context/user stack
        let (memory_set, mut user_sp, entry_point) = MemorySet::from_elf(elf_data);
        let trap_cx_ppn = memory_set
            .translate(VirtAddr::from(TRAP_CONTEXT).into())
            .unwrap()
            .ppn();
        // push arguments on user stack
        user_sp -= (args.len() + 1) * core::mem::size_of::<usize>();
        let argv_base = user_sp;
        let mut argv: Vec<_> = (0..=args.len())
            .map(|arg| {
                translated_refmut(
                    memory_set.token(),
                    (argv_base + arg * core::mem::size_of::<usize>()) as *mut usize
                )
            })
            .collect();
        *argv[args.len()] = 0;
        for i in 0..args.len() {
            user_sp -= args[i].len() + 1;
            *argv[i] = user_sp;
            let mut p = user_sp;
            for c in args[i].as_bytes() {
                *translated_refmut(memory_set.token(), p as *mut u8) = *c;
                p += 1;
            }
            *translated_refmut(memory_set.token(), p as *mut u8) = 0;
        }
        // make the user_sp aligned to 8B for k210 platform
        user_sp -= user_sp % core::mem::size_of::<usize>();

        // **** hold current PCB lock
        let mut inner = self.acquire_inner_lock();
        // substitute memory_set
        inner.memory_set = memory_set;
        // update trap_cx ppn
        inner.trap_cx_ppn = trap_cx_ppn;
        // initialize trap_cx
        let mut trap_cx = TrapContext::app_init_context(
            entry_point,
            user_sp,
            KERNEL_SPACE.lock().token(),
            self.kernel_stack.get_top(),
            trap_handler as usize,
        );
        trap_cx.x[10] = args.len();
        trap_cx.x[11] = argv_base;
        *inner.get_trap_cx() = trap_cx;
        // **** release current PCB lock
    }
    pub fn fork(self: &Arc<TaskControlBlock>) -> Arc<TaskControlBlock> {
        // ---- hold parent PCB lock
        let mut parent_inner = self.acquire_inner_lock();
        // copy user space(include trap context)
        let memory_set = MemorySet::from_existed_user(
            &parent_inner.memory_set
        );
        let trap_cx_ppn = memory_set
            .translate(VirtAddr::from(TRAP_CONTEXT).into())
            .unwrap()
            .ppn();
        // alloc a pid and a kernel stack in kernel space
        let pid_handle = pid_alloc();
        let kernel_stack = KernelStack::new(&pid_handle);
        let kernel_stack_top = kernel_stack.get_top();
        // push a goto_trap_return task_cx on the top of kernel stack
        let task_cx_ptr = kernel_stack.push_on_top(TaskContext::goto_trap_return());
        // copy fd table
        let mut new_fd_table: Vec<Option<Arc<dyn File + Send + Sync>>> = Vec::new();
        for fd in parent_inner.fd_table.iter() {
            if let Some(file) = fd {
                new_fd_table.push(Some(file.clone()));
            } else {
                new_fd_table.push(None);
            }
        }
        let name = parent_inner.name.clone();
        let ppid = self.pid.0 as isize;
        let task_control_block = Arc::new(TaskControlBlock {
            pid: pid_handle,
            kernel_stack,
            ppid,
            inner: Mutex::new(TaskControlBlockInner {
                trap_cx_ppn,
                base_size: parent_inner.base_size,
                task_cx_ptr: task_cx_ptr as usize,
                task_status: TaskStatus::Ready,
                name,
                memory_set,
                parent: Some(Arc::downgrade(self)),
                children: Vec::new(),
                exit_code: 0,
                fd_table: new_fd_table,
                start_time: 0,
                cpu_time: 0,
                syscall: Arc::new(Mutex::new(SyscallCount{
                    syscall_dup: (0, 0),
                    syscall_open: (0, 0),
                    syscall_close: (0, 0),
                    syscall_pipe: (0, 0),
                    syscall_read: (0, 0),
                    syscall_write: (0, 0),
                    syscall_exit: (0, 0),
                    syscall_yield: (0, 0),
                    syscall_get_time: (0, 0),
                    syscall_getpid: (0, 0),
                    syscall_fork: (0, 0),
                    syscall_exec: (0, 0),
                    syscall_waitpid: (0, 0),
                }))
                // czy
                //cpu_time:0,
                // czy
                //start_run_stamp:0,
            }),
        });
        // add child
        parent_inner.children.push(task_control_block.clone());
        // modify kernel_sp in trap_cx
        // **** acquire child PCB lock
        let trap_cx = task_control_block.acquire_inner_lock().get_trap_cx();
        // **** release child PCB lock
        trap_cx.kernel_sp = kernel_stack_top;
        // return
        task_control_block
        // ---- release parent PCB lock
    }
    pub fn getpid(&self) -> usize {
        self.pid.0
    }

    pub fn get_ppid(&self) -> isize{
        self.ppid
    }

    pub fn get_cpu_time(&self) -> usize{
        self.acquire_inner_lock().cpu_time
    }

    pub fn get_name(&self) -> String{
        let temp1 = self.acquire_inner_lock();
        temp1.name.clone()
    }

    pub fn get_syscall_cnt(&self) -> Arc<Mutex<SyscallCount>>{
        Arc::clone(&self.acquire_inner_lock().syscall) 
    }


}

#[derive(Copy, Clone, PartialEq)]
pub enum TaskStatus {
    Ready,
    Running,
    Zombie,
}