esp_wifi/preempt_builtin/
mod.rs

1#[cfg_attr(target_arch = "riscv32", path = "preempt_riscv.rs")]
2#[cfg_attr(target_arch = "xtensa", path = "preempt_xtensa.rs")]
3mod arch_specific;
4pub mod timer;
5
6use core::{ffi::c_void, mem::MaybeUninit};
7
8use allocator_api2::boxed::Box;
9use arch_specific::*;
10pub(crate) use timer::setup_timer;
11use timer::{disable_multitasking, setup_multitasking};
12
13use crate::{
14    compat::malloc::InternalMemory,
15    hal::{sync::Locked, trapframe::TrapFrame},
16    preempt::Scheduler,
17    preempt_builtin::timer::disable_timebase,
18};
19
20struct Context {
21    trap_frame: TrapFrame,
22    pub thread_semaphore: u32,
23    pub next: *mut Context,
24    pub _allocated_stack: Box<[MaybeUninit<u8>], InternalMemory>,
25}
26
27impl Context {
28    pub(crate) fn new(
29        task_fn: extern "C" fn(*mut c_void),
30        param: *mut c_void,
31        task_stack_size: usize,
32    ) -> Self {
33        trace!("task_create {:?} {:?} {}", task_fn, param, task_stack_size);
34
35        let mut stack = Box::<[u8], _>::new_uninit_slice_in(task_stack_size, InternalMemory);
36
37        let stack_top = unsafe { stack.as_mut_ptr().add(task_stack_size).cast() };
38
39        Context {
40            trap_frame: new_task_context(task_fn, param, stack_top),
41            thread_semaphore: 0,
42            next: core::ptr::null_mut(),
43            _allocated_stack: stack,
44        }
45    }
46}
47
48struct SchedulerState {
49    /// Pointer to the current task.
50    ///
51    /// Tasks are stored in a circular linked list. CTX_NOW points to the
52    /// current task.
53    current_task: *mut Context,
54
55    /// Pointer to the task that is scheduled for deletion.
56    to_delete: *mut Context,
57}
58
59impl SchedulerState {
60    const fn new() -> Self {
61        Self {
62            current_task: core::ptr::null_mut(),
63            to_delete: core::ptr::null_mut(),
64        }
65    }
66
67    fn delete_task(&mut self, task: *mut Context) {
68        let mut current_task = self.current_task;
69        // Save the first pointer so we can prevent an accidental infinite loop.
70        let initial = current_task;
71        loop {
72            // We don't have the previous pointer, so we need to walk forward in the circle
73            // even if we need to delete the first task.
74
75            // If the next task is the one we want to delete, we need to remove it from the
76            // list, then drop it.
77            let next_task = unsafe { (*current_task).next };
78            if core::ptr::eq(next_task, task) {
79                unsafe {
80                    (*current_task).next = (*next_task).next;
81
82                    core::ptr::drop_in_place(task);
83                    break;
84                }
85            }
86
87            // If the next task is the first task, we can stop. If we needed to delete the
88            // first task, we have already handled it in the above case. If we needed to
89            // delete another task, it has already been deleted in a previous iteration.
90            if core::ptr::eq(next_task, initial) {
91                break;
92            }
93
94            // Move to the next task.
95            current_task = next_task;
96        }
97    }
98
99    fn switch_task(&mut self, trap_frame: &mut TrapFrame) {
100        save_task_context(unsafe { &mut *self.current_task }, trap_frame);
101
102        if !self.to_delete.is_null() {
103            let task_to_delete = core::mem::replace(&mut self.to_delete, core::ptr::null_mut());
104            self.delete_task(task_to_delete);
105        }
106
107        unsafe { self.current_task = (*self.current_task).next };
108
109        restore_task_context(unsafe { &mut *self.current_task }, trap_frame);
110    }
111
112    fn schedule_task_deletion(&mut self, task: *mut Context) -> bool {
113        self.to_delete = task;
114        core::ptr::eq(task, self.current_task)
115    }
116}
117
118static SCHEDULER_STATE: Locked<SchedulerState> = Locked::new(SchedulerState::new());
119
120struct BuiltinScheduler {}
121
122crate::scheduler_impl!(static SCHEDULER: BuiltinScheduler = BuiltinScheduler {});
123
124impl Scheduler for BuiltinScheduler {
125    fn enable(&self) {
126        // allocate the main task
127        allocate_main_task();
128        setup_multitasking();
129    }
130
131    fn disable(&self) {
132        disable_timebase();
133        disable_multitasking();
134        delete_all_tasks();
135    }
136
137    fn yield_task(&self) {
138        timer::yield_task()
139    }
140
141    fn task_create(
142        &self,
143        task: extern "C" fn(*mut c_void),
144        param: *mut c_void,
145        task_stack_size: usize,
146    ) -> *mut c_void {
147        let task = Box::new_in(Context::new(task, param, task_stack_size), InternalMemory);
148        let task_ptr = Box::into_raw(task);
149
150        SCHEDULER_STATE.with(|state| unsafe {
151            let current_task = state.current_task;
152            debug_assert!(
153                !current_task.is_null(),
154                "Tried to allocate a task before allocating the main task"
155            );
156            // Insert the new task at the next position.
157            let next = (*current_task).next;
158            (*task_ptr).next = next;
159            (*current_task).next = task_ptr;
160        });
161
162        task_ptr as *mut c_void
163    }
164
165    fn current_task(&self) -> *mut c_void {
166        current_task() as *mut c_void
167    }
168
169    fn schedule_task_deletion(&self, task_handle: *mut c_void) {
170        schedule_task_deletion(task_handle as *mut Context)
171    }
172
173    fn current_task_thread_semaphore(&self) -> *mut crate::binary::c_types::c_void {
174        unsafe {
175            &mut ((*current_task()).thread_semaphore) as *mut _
176                as *mut crate::binary::c_types::c_void
177        }
178    }
179}
180
181fn allocate_main_task() {
182    // This context will be filled out by the first context switch.
183    let context = Box::new_in(
184        Context {
185            trap_frame: TrapFrame::default(),
186            thread_semaphore: 0,
187            next: core::ptr::null_mut(),
188            _allocated_stack: Box::<[u8], _>::new_uninit_slice_in(0, InternalMemory),
189        },
190        InternalMemory,
191    );
192
193    let context_ptr = Box::into_raw(context);
194    unsafe {
195        // The first task loops back to itself.
196        (*context_ptr).next = context_ptr;
197    }
198
199    SCHEDULER_STATE.with(|state| {
200        debug_assert!(
201            state.current_task.is_null(),
202            "Tried to allocate main task multiple times"
203        );
204        state.current_task = context_ptr;
205    })
206}
207
208fn delete_all_tasks() {
209    let first_task = SCHEDULER_STATE.with(|state| {
210        // Remove all tasks from the list. We will drop them outside of the critical
211        // section.
212        core::mem::replace(&mut state.current_task, core::ptr::null_mut())
213    });
214
215    if first_task.is_null() {
216        return;
217    }
218
219    let mut task_to_delete = first_task;
220
221    loop {
222        let next_task = unsafe {
223            // SAFETY: Tasks are in a circular linked list. We are guaranteed that the next
224            // task is a valid pointer, or the first task that may already have been
225            // deleted. In the second case, we will not move on to the next
226            // iteration, so the loop will not try to free a task twice.
227            let next_task = (*task_to_delete).next;
228            core::ptr::drop_in_place(task_to_delete);
229            next_task
230        };
231
232        if core::ptr::eq(next_task, first_task) {
233            break;
234        }
235
236        task_to_delete = next_task;
237    }
238}
239
240fn current_task() -> *mut Context {
241    SCHEDULER_STATE.with(|state| state.current_task)
242}
243
244fn schedule_task_deletion(task: *mut Context) {
245    let deleting_current = SCHEDULER_STATE.with(|state| state.schedule_task_deletion(task));
246
247    // Tasks are deleted during context switches, so we need to yield if we are
248    // deleting the current task.
249    if deleting_current {
250        loop {
251            timer::yield_task();
252        }
253    }
254}
255
256pub(crate) fn task_switch(trap_frame: &mut TrapFrame) {
257    SCHEDULER_STATE.with(|state| state.switch_task(trap_frame));
258}