1use super::{Future, Pin};
2use std::cell::RefCell;
3use std::os::fd::{AsRawFd, FromRawFd, RawFd};
4use std::sync::OnceLock;
5use std::sync::atomic::{
6 AtomicBool, AtomicI32, AtomicPtr, AtomicU32, AtomicUsize, Ordering, fence,
7};
8use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
9
10use crate::io::trace::io_trace;
13#[allow(unused_imports)]
14use crate::io::trace::trace_now_us;
15
16static WORKER_ROUND_ROBIN: AtomicUsize = AtomicUsize::new(0);
17
18use crate::lockfree::{BufferPool, SpscQueue, TreiberStack};
29
30#[repr(align(64))]
34struct LocalAllocator {
35 thread_idx: usize,
36 local_chunks: Vec<u32>,
37}
38
39thread_local! {
40 static LOCAL_ALLOCATOR: RefCell<Option<LocalAllocator>> = const { RefCell::new(None) };
41 static THREAD_ID: usize = {
42 static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
43 NEXT_ID.fetch_add(1, Ordering::Relaxed)
44 };
45}
46
47#[doc(hidden)]
48#[must_use]
49pub fn get_local_thread_id() -> usize {
50 THREAD_ID.with(|id| *id)
51}
52
53static THREAD_RETURNED_STACKS: OnceLock<Box<[TreiberStack]>> = OnceLock::new();
54static GLOBAL_BUFFER_POOL: OnceLock<BufferPool> = OnceLock::new();
55static CHUNK_OWNERS: OnceLock<Box<[AtomicU32]>> = OnceLock::new();
63
64fn get_or_init_local_allocator() -> Option<usize> {
65 LOCAL_ALLOCATOR.with(|cell| {
66 let mut borrow = cell.borrow_mut();
67 if borrow.is_none() {
68 let idx = get_local_thread_id();
69 if idx < 512 {
70 *borrow = Some(LocalAllocator {
71 thread_idx: idx,
72 local_chunks: Vec::new(),
73 });
74 }
75 }
76 borrow.as_ref().map(|alloc| alloc.thread_idx)
77 })
78}
79
80#[doc(hidden)]
81pub fn allocate_buffer() -> Option<u32> {
82 let t_idx_opt = get_or_init_local_allocator();
83 if let Some(t_idx) = t_idx_opt {
84 LOCAL_ALLOCATOR.with(|cell| {
85 let mut borrow = cell.borrow_mut();
86 let alloc = borrow.as_mut().unwrap();
87
88 if let Some(idx) = alloc.local_chunks.pop() {
90 return Some(idx);
91 }
92 if let Some(stacks) = THREAD_RETURNED_STACKS.get()
94 && let Some(stack) = stacks.get(t_idx)
95 {
96 while let Some(idx) = stack.pop() {
97 alloc.local_chunks.push(idx);
98 }
99 if let Some(idx) = alloc.local_chunks.pop() {
100 return Some(idx);
101 }
102 }
103 if let Some(pool) = GLOBAL_BUFFER_POOL.get()
105 && let Some(idx) = pool.acquire()
106 {
107 if let Some(owners) = CHUNK_OWNERS.get() {
108 owners[idx as usize].store(t_idx as u32, Ordering::Release);
109 }
110 return Some(idx);
111 }
112 None
113 })
114 } else if let Some(pool) = GLOBAL_BUFFER_POOL.get() {
115 if let Some(idx) = pool.acquire() {
116 if let Some(owners) = CHUNK_OWNERS.get() {
117 owners[idx as usize].store(u32::MAX, Ordering::Release);
118 }
119 return Some(idx);
120 }
121 None
122 } else {
123 None
124 }
125}
126
127#[doc(hidden)]
128pub fn free_buffer(idx: u32) {
129 if let Some(pool) = GLOBAL_BUFFER_POOL.get() {
130 let owner = CHUNK_OWNERS.get().map_or(u32::MAX, |owners| {
131 owners[idx as usize].load(Ordering::Acquire)
132 });
133 if owner == u32::MAX {
134 pool.release(idx);
135 return;
136 }
137
138 let current_thread_idx = get_or_init_local_allocator();
139 if Some(owner as usize) == current_thread_idx {
140 LOCAL_ALLOCATOR.with(|cell| {
141 if let Some(alloc) = cell.borrow_mut().as_mut() {
142 alloc.local_chunks.push(idx);
143 }
144 });
145 } else if let Some(stacks) = THREAD_RETURNED_STACKS.get() {
146 if let Some(stack) = stacks.get(owner as usize) {
147 stack.push(idx);
148 } else {
149 pool.release(idx);
150 }
151 } else {
152 pool.release(idx);
153 }
154 }
155}
156
157#[doc(hidden)]
158#[repr(align(64))]
159pub struct BufferSlice {
160 pub buf_idx: u32,
161 pub read_pos: usize,
162 pub write_pos: usize,
163}
164
165impl BufferSlice {
166 #[must_use]
167 pub const fn new(buf_idx: u32, len: usize) -> Self {
168 Self {
169 buf_idx,
170 read_pos: 0,
171 write_pos: len,
172 }
173 }
174
175 #[inline]
182 pub fn data(&self) -> *mut u8 {
183 GLOBAL_BUFFER_POOL.get().unwrap().get_ptr(self.buf_idx)
184 }
185
186 #[inline]
187 #[must_use]
188 pub const fn remaining(&self) -> usize {
189 self.write_pos.saturating_sub(self.read_pos)
190 }
191}
192
193impl Drop for BufferSlice {
194 fn drop(&mut self) {
195 free_buffer(self.buf_idx);
196 }
197}
198
199#[derive(Clone, Copy, Debug, PartialEq, Eq)]
212pub enum OpCode {
213 Read,
215 Write,
217 Accept,
219 Connect,
221 SendTo,
225 RecvFrom,
229}
230
231#[derive(Clone, Copy)]
238pub enum IoRequest {
239 Read {
242 fd: u32,
244 direct_fd_idx: u32,
247 buf_ptr: *mut u8,
249 len: usize,
251 offset: i64,
253 slot_idx: usize,
255 },
256 Write {
258 fd: u32,
260 direct_fd_idx: u32,
263 buf_ptr: *const u8,
265 len: usize,
267 offset: i64,
269 slot_idx: usize,
271 },
272 Accept {
274 fd: u32,
276 direct_fd_idx: u32,
279 slot_idx: usize,
281 },
282 Connect {
284 fd: u32,
286 direct_fd_idx: u32,
289 addr: libc::sockaddr_storage,
291 addr_len: libc::socklen_t,
293 slot_idx: usize,
295 },
296 SendTo {
298 fd: u32,
300 direct_fd_idx: u32,
303 msg_ptr: *mut libc::msghdr,
306 slot_idx: usize,
308 },
309 RecvFrom {
311 fd: u32,
313 direct_fd_idx: u32,
316 msg_ptr: *mut libc::msghdr,
319 slot_idx: usize,
321 },
322 RegisterFile {
325 fd: RawFd,
327 slot_idx: usize,
329 },
330 UnregisterFile {
332 direct_fd_idx: u32,
334 slot_idx: usize,
336 },
337}
338
339#[repr(align(64))]
351struct WakerSlot {
352 waker_data: AtomicPtr<()>,
354 waker_vtable: AtomicPtr<RawWakerVTable>,
357 waker_lock: AtomicBool,
358 result: AtomicI32,
359 completed: AtomicBool,
360 dropped: AtomicBool,
361 origin_fd: AtomicU32,
365}
366
367#[repr(align(64))]
368struct WaitSlot {
369 waker_data: AtomicPtr<()>,
370 waker_vtable: AtomicPtr<RawWakerVTable>,
371}
372
373impl WakerSlot {
374 #[inline(always)]
375 fn lock_waker(&self) {
376 while self
377 .waker_lock
378 .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
379 .is_err()
380 {
381 core::hint::spin_loop();
382 }
383 }
384
385 #[inline(always)]
386 fn unlock_waker(&self) {
387 self.waker_lock.store(false, Ordering::Release);
388 }
389}
390
391#[inline(always)]
392fn wake_next_waiting_fiber(state: &WorkerState) {
393 if let Some(wait_idx) = state.waiting_queue.pop() {
394 let wait_slot = &state.wait_slots[wait_idx as usize];
395 let data = wait_slot
396 .waker_data
397 .swap(std::ptr::null_mut(), Ordering::Relaxed);
398 let vtable = wait_slot
399 .waker_vtable
400 .swap(std::ptr::null_mut(), Ordering::Relaxed);
401 state.free_wait_slots.push(wait_idx);
402
403 if !data.is_null() && !vtable.is_null() {
404 let raw = RawWaker::new(data.cast_const(), unsafe { &*vtable });
405 let w = unsafe { Waker::from_raw(raw) };
406 w.wake();
407 }
408 }
409}
410
411pub struct WorkerState {
418 #[cfg(target_os = "linux")]
419 ring: std::cell::UnsafeCell<io_uring::IoUring>,
420 #[cfg(not(target_os = "linux"))]
421 poll: std::cell::UnsafeCell<mio::Poll>,
422
423 queues: Box<[SpscQueue<IoRequest>]>,
424 slots: Box<[WakerSlot]>,
425 free_slots: TreiberStack,
426
427 wait_slots: Box<[WaitSlot]>,
428 free_wait_slots: TreiberStack,
429 waiting_queue: TreiberStack,
430 is_sleeping: AtomicBool,
431
432 cancel_queue: TreiberStack,
439
440 #[cfg(target_os = "linux")]
441 wake_eventfd: RawFd,
442 #[cfg(target_os = "linux")]
443 sqpoll_enabled: bool,
444 #[cfg(not(target_os = "linux"))]
445 waker: std::sync::Arc<mio::Waker>,
446
447 direct_fd_free: TreiberStack,
448}
449
450#[allow(clippy::non_send_fields_in_send_ty)]
461unsafe impl Send for WorkerState {}
462unsafe impl Sync for WorkerState {}
463
464struct GlobalConfig {
473 workers: usize,
474 pin_cpus: Vec<usize>,
475}
476
477static GLOBAL_CONFIG: OnceLock<GlobalConfig> = OnceLock::new();
478static WORKERS: OnceLock<Box<[WorkerState]>> = OnceLock::new();
479static SHUTDOWN: AtomicBool = AtomicBool::new(false);
480
481#[cfg(target_os = "linux")]
482fn pin_thread_to_cpu(cpu_id: usize) -> Result<(), &'static str> {
483 unsafe {
484 let mut cpuset: libc::cpu_set_t = std::mem::zeroed();
485 libc::CPU_SET(cpu_id, &mut cpuset);
486 let thread = libc::pthread_self();
487 let res = libc::pthread_setaffinity_np(
488 thread,
489 std::mem::size_of::<libc::cpu_set_t>(),
490 &raw const cpuset,
491 );
492 if res == 0 {
493 Ok(())
494 } else {
495 Err("pthread_setaffinity_np failed")
496 }
497 }
498}
499
500#[cfg(not(target_os = "linux"))]
501fn pin_thread_to_cpu(_cpu_id: usize) -> Result<(), &'static str> {
502 Ok(())
503}
504
505#[cfg(target_os = "linux")]
514fn pick_direct_fd_count(desired_max: usize) -> usize {
515 unsafe {
516 let mut lim: libc::rlimit = std::mem::zeroed();
517 if libc::getrlimit(libc::RLIMIT_NOFILE, &raw mut lim) == 0 {
518 if lim.rlim_cur < lim.rlim_max {
519 let raised = libc::rlimit {
520 rlim_cur: lim.rlim_max,
521 rlim_max: lim.rlim_max,
522 };
523 let _ = libc::setrlimit(libc::RLIMIT_NOFILE, &raw const raised);
524 let _ = libc::getrlimit(libc::RLIMIT_NOFILE, &raw mut lim);
525 }
526 let headroom = 256usize;
527 let available = (lim.rlim_cur as usize).saturating_sub(headroom);
528 return available.clamp(64, desired_max);
529 }
530 }
531 64
534}
535
536#[allow(clippy::too_many_lines)]
568pub fn init_runtime(
569 workers: usize,
570 ring_depth: u32,
571 buffer_pool_size: usize,
572 chunk_size: usize,
573 pin_cpus: &[usize],
574) {
575 let config = GlobalConfig {
576 workers,
577 pin_cpus: pin_cpus.to_vec(),
578 };
579 if GLOBAL_CONFIG.set(config).is_err() {
580 return;
581 }
582
583 let pool = BufferPool::new(buffer_pool_size, chunk_size);
584 let _ = GLOBAL_BUFFER_POOL.set(pool);
585 let owners: Vec<AtomicU32> = (0..buffer_pool_size)
586 .map(|_| AtomicU32::new(u32::MAX))
587 .collect();
588 let _ = CHUNK_OWNERS.set(owners.into_boxed_slice());
589
590 let mut returned_stacks = Vec::with_capacity(512);
591 for _ in 0..512 {
592 returned_stacks.push(TreiberStack::new(0));
593 }
594 let _ = THREAD_RETURNED_STACKS.set(returned_stacks.into_boxed_slice());
595
596 let mut worker_states = Vec::with_capacity(workers);
597 for _worker_idx in 0..workers {
598 let mut queues = Vec::with_capacity(512);
599 for _ in 0..512 {
600 queues.push(SpscQueue::new(256));
601 }
602 let queues = queues.into_boxed_slice();
603
604 let mut slots = Vec::with_capacity(ring_depth as usize);
605 for _ in 0..ring_depth {
606 slots.push(WakerSlot {
607 waker_data: AtomicPtr::new(std::ptr::null_mut()),
608 waker_vtable: AtomicPtr::new(std::ptr::null_mut()),
609 waker_lock: AtomicBool::new(false),
610 result: AtomicI32::new(0),
611 completed: AtomicBool::new(false),
612 dropped: AtomicBool::new(false),
613 origin_fd: AtomicU32::new(u32::MAX),
614 });
615 }
616 let slots = slots.into_boxed_slice();
617 let free_slots = TreiberStack::new(ring_depth as usize);
618 for i in 0..ring_depth {
619 free_slots.push(i);
620 }
621 let cancel_queue = TreiberStack::new(ring_depth as usize);
622
623 let wait_slots_depth = 65536;
624 let mut wait_slots = Vec::with_capacity(wait_slots_depth);
625 for _ in 0..wait_slots_depth {
626 wait_slots.push(WaitSlot {
627 waker_data: AtomicPtr::new(std::ptr::null_mut()),
628 waker_vtable: AtomicPtr::new(std::ptr::null_mut()),
629 });
630 }
631 let wait_slots = wait_slots.into_boxed_slice();
632 let free_wait_slots = TreiberStack::new(wait_slots_depth);
633 for i in 0..wait_slots_depth {
634 free_wait_slots.push(i as u32);
635 }
636 let waiting_queue = TreiberStack::new(wait_slots_depth);
637 let is_sleeping = AtomicBool::new(false);
638
639 #[cfg(target_os = "linux")]
640 let direct_fd_count = pick_direct_fd_count(4096);
641 #[cfg(not(target_os = "linux"))]
642 let direct_fd_count = 4096usize;
643
644 let direct_fd_free = TreiberStack::new(direct_fd_count);
645 for i in 0..direct_fd_count as u32 {
646 direct_fd_free.push(i);
647 }
648
649 #[cfg(target_os = "linux")]
650 {
651 let wake_eventfd = unsafe { libc::eventfd(0, libc::EFD_CLOEXEC | libc::EFD_NONBLOCK) };
652 assert!(wake_eventfd >= 0, "Failed to create eventfd");
653
654 let (ring, sqpoll_enabled) = io_uring::IoUring::builder()
655 .setup_sqpoll(2000)
656 .build(ring_depth)
657 .map_or_else(
658 |_| {
659 (
660 io_uring::IoUring::new(ring_depth)
661 .expect("Failed to initialize io_uring fallback"),
662 false,
663 )
664 },
665 |r| (r, true),
666 );
667
668 let initial_fds = vec![-1; direct_fd_count];
669 ring.submitter()
670 .register_files(&initial_fds)
671 .expect("Failed to register direct FDs");
672
673 worker_states.push(WorkerState {
674 ring: std::cell::UnsafeCell::new(ring),
675 queues,
676 slots,
677 free_slots,
678 wait_slots,
679 free_wait_slots,
680 waiting_queue,
681 is_sleeping,
682 cancel_queue,
683 wake_eventfd,
684 sqpoll_enabled,
685 direct_fd_free,
686 });
687 }
688
689 #[cfg(not(target_os = "linux"))]
690 {
691 let poll = mio::Poll::new().expect("Failed to initialize mio Poll");
692 let waker = std::sync::Arc::new(
693 mio::Waker::new(poll.registry(), mio::Token(0))
694 .expect("Failed to create mio waker"),
695 );
696
697 worker_states.push(WorkerState {
698 poll: std::cell::UnsafeCell::new(poll),
699 queues,
700 slots,
701 free_slots,
702 wait_slots,
703 free_wait_slots,
704 waiting_queue,
705 is_sleeping,
706 cancel_queue,
707 waker,
708 direct_fd_free,
709 });
710 }
711 }
712
713 let worker_states = worker_states.into_boxed_slice();
714 let _ = WORKERS.set(worker_states);
715
716 for worker_idx in 0..workers {
717 std::thread::Builder::new()
718 .name(format!("dtact-io-worker-{worker_idx}"))
719 .spawn(move || {
720 LOCAL_ALLOCATOR.with(|cell| {
721 *cell.borrow_mut() = Some(LocalAllocator {
722 thread_idx: worker_idx,
723 local_chunks: Vec::new(),
724 });
725 });
726
727 let state = &WORKERS.get().unwrap()[worker_idx];
728
729 #[cfg(target_os = "linux")]
730 run_linux_worker_loop(worker_idx, state);
731
732 #[cfg(not(target_os = "linux"))]
733 run_mio_worker_loop(worker_idx, state);
734 })
735 .expect("Failed to spawn dtact-io worker thread");
736 }
737}
738
739pub fn init(workers: usize) {
746 init_runtime(workers, 1024, 65536, 4096, &[]);
747}
748
749pub fn shutdown_runtime() {
753 SHUTDOWN.store(true, Ordering::Release);
754 if let Some(workers) = WORKERS.get() {
755 for state in workers {
756 #[cfg(target_os = "linux")]
757 let _ = unsafe {
758 libc::write(
759 state.wake_eventfd,
760 std::ptr::from_ref::<u64>(&1u64).cast::<libc::c_void>(),
761 8,
762 )
763 };
764 #[cfg(not(target_os = "linux"))]
765 state.waker.wake();
766 }
767 }
768}
769
770#[allow(clippy::too_many_lines)]
780#[cfg(target_os = "linux")]
781fn run_linux_worker_loop(worker_idx: usize, state: &WorkerState) {
782 if let Some(config) = GLOBAL_CONFIG.get()
783 && let Some(&cpu_id) = config.pin_cpus.get(worker_idx)
784 {
785 let _ = pin_thread_to_cpu(cpu_id);
786 }
787
788 let ring = unsafe { &mut *state.ring.get() };
789 let mut eventfd_buf = 0u64;
790 let mut eventfd_submitted = false;
791 #[cfg(feature = "spin")]
796 let mut idle_streak: u32 = 0;
797
798 loop {
799 if SHUTDOWN.load(Ordering::Relaxed) {
800 break;
801 }
802
803 if !eventfd_submitted {
804 let sqe = io_uring::opcode::Read::new(
805 io_uring::types::Fd(state.wake_eventfd),
806 (&raw mut eventfd_buf).cast::<u8>(),
807 8,
808 )
809 .build()
810 .user_data(u64::MAX);
811
812 unsafe {
813 if ring.submission().push(&sqe).is_ok() {
814 eventfd_submitted = true;
815 }
816 }
817 }
818
819 let mut pushed_sqe = false;
820 for q in &state.queues {
821 while let Some(req) = q.pop() {
822 pushed_sqe = true;
823 let _ = submit_linux_request(state, &req);
824 }
825 }
826
827 while let Some(slot_idx) = state.cancel_queue.pop() {
828 pushed_sqe = true;
829 let sqe = io_uring::opcode::AsyncCancel::new(u64::from(slot_idx))
830 .build()
831 .user_data(u64::MAX - 1);
832 unsafe {
833 let _ = push_sqe(ring, &sqe);
834 }
835 }
836
837 let any_pending = state.queues.iter().any(|q| !q.is_empty());
848
849 let mut folded_wait = false;
857 if pushed_sqe || eventfd_submitted {
858 ring.submission().sync();
859 if pushed_sqe && !any_pending {
860 state.is_sleeping.store(true, Ordering::SeqCst);
861 fence(Ordering::SeqCst);
870 let missed = state.queues.iter().any(|q| !q.is_empty());
871 if missed {
872 state.is_sleeping.store(false, Ordering::SeqCst);
873 let sr = ring.submit();
874 io_trace!(
875 "[dtact-io] t={} loop submit(folded-missed) result={:?}",
876 trace_now_us(),
877 sr
878 );
879 } else {
880 let sr = ring.submit_and_wait(1);
881 state.is_sleeping.store(false, Ordering::Release);
882 io_trace!(
883 "[dtact-io] t={} loop submit_and_wait(folded) result={:?}",
884 trace_now_us(),
885 sr
886 );
887 }
888 folded_wait = true;
889 } else {
890 let should_enter = if state.sqpoll_enabled {
891 ring.submission().need_wakeup()
892 } else {
893 true
894 };
895 if should_enter {
896 let sr = ring.submit();
897 io_trace!(
898 "[dtact-io] t={} loop submit() result={:?}",
899 trace_now_us(),
900 sr
901 );
902 }
903 }
904 }
905
906 let mut has_completions = false;
907 let mut cq = ring.completion();
908 cq.sync();
909 let cq_len = cq.len();
910 if cq_len > 0 {
911 io_trace!("[dtact-io] t={} loop cq_len={}", trace_now_us(), cq_len);
912 }
913 for cqe in cq {
914 has_completions = true;
915 let user_data = cqe.user_data();
916 let res = cqe.result();
917
918 if user_data == u64::MAX {
919 eventfd_submitted = false;
920 } else if user_data == u64::MAX - 1 {
921 } else {
923 process_linux_completion(state, user_data as usize, res);
924 }
925 }
926
927 #[cfg(feature = "spin")]
928 if folded_wait || pushed_sqe || has_completions {
929 idle_streak = 0;
930 }
931
932 if !folded_wait && !pushed_sqe && !has_completions {
933 #[cfg(feature = "spin")]
957 {
958 if adaptive_idle_spin(state, idle_streak) {
959 continue;
960 }
961 }
962 #[cfg(feature = "spin")]
963 {
964 idle_streak = idle_streak.saturating_add(1);
965 }
966
967 state.is_sleeping.store(true, Ordering::SeqCst);
968 fence(Ordering::SeqCst);
975 let missed = state.queues.iter().any(|q| !q.is_empty());
976 if !any_pending && !missed {
977 let sr = ring.submit_and_wait(1);
978 io_trace!(
979 "[dtact-io] t={} loop submit_and_wait(idle) result={:?}",
980 trace_now_us(),
981 sr
982 );
983 } else if missed {
984 let sr = ring.submit();
985 io_trace!(
986 "[dtact-io] t={} loop submit(idle-missed) result={:?}",
987 trace_now_us(),
988 sr
989 );
990 }
991 state.is_sleeping.store(false, Ordering::Release);
992 }
993 }
994}
995
996#[cfg(all(target_os = "linux", feature = "spin"))]
1012fn adaptive_idle_spin(state: &WorkerState, idle_streak: u32) -> bool {
1013 const ARM_AFTER: u32 = 2;
1017 if idle_streak < ARM_AFTER {
1018 return false;
1019 }
1020
1021 let budget = 256u32.saturating_mul(idle_streak.saturating_sub(ARM_AFTER) + 1);
1025 let budget = budget.min(4096);
1026
1027 for _ in 0..budget {
1028 if state.queues.iter().any(|q| !q.is_empty()) {
1029 return true;
1030 }
1031 core::hint::spin_loop();
1032 }
1033 false
1034}
1035
1036#[cfg(target_os = "linux")]
1037unsafe fn push_sqe(
1038 ring: &mut io_uring::IoUring,
1039 sqe: &io_uring::squeue::Entry,
1040) -> Result<(), &'static str> {
1041 loop {
1042 let res = unsafe { ring.submission().push(sqe) };
1043 if res == Ok(()) {
1044 return Ok(());
1045 }
1046 let _ = ring.submit();
1047 core::hint::spin_loop();
1048 }
1049}
1050
1051#[allow(clippy::too_many_lines)]
1059#[cfg(target_os = "linux")]
1060fn submit_linux_request(state: &WorkerState, req: &IoRequest) -> Result<(), &'static str> {
1061 let ring = unsafe { &mut *state.ring.get() };
1062
1063 let sqe = match *req {
1064 IoRequest::Read {
1065 fd,
1066 direct_fd_idx,
1067 buf_ptr,
1068 len,
1069 offset,
1070 slot_idx,
1071 } => {
1072 let use_fixed = direct_fd_idx != u32::MAX;
1073 let target_fd = if use_fixed {
1074 direct_fd_idx as i32
1075 } else {
1076 fd as i32
1077 };
1078 let mut s =
1079 io_uring::opcode::Read::new(io_uring::types::Fd(target_fd), buf_ptr, len as u32)
1080 .offset(offset as u64)
1081 .build()
1082 .user_data(slot_idx as u64);
1083 if use_fixed {
1084 s = s.flags(io_uring::squeue::Flags::FIXED_FILE);
1085 }
1086 s
1087 }
1088 IoRequest::Write {
1089 fd,
1090 direct_fd_idx,
1091 buf_ptr,
1092 len,
1093 offset,
1094 slot_idx,
1095 } => {
1096 let use_fixed = direct_fd_idx != u32::MAX;
1097 let target_fd = if use_fixed {
1098 direct_fd_idx as i32
1099 } else {
1100 fd as i32
1101 };
1102 let mut s =
1103 io_uring::opcode::Write::new(io_uring::types::Fd(target_fd), buf_ptr, len as u32)
1104 .offset(offset as u64)
1105 .build()
1106 .user_data(slot_idx as u64);
1107 if use_fixed {
1108 s = s.flags(io_uring::squeue::Flags::FIXED_FILE);
1109 }
1110 s
1111 }
1112 IoRequest::Accept {
1113 fd,
1114 direct_fd_idx,
1115 slot_idx,
1116 } => {
1117 let use_fixed = direct_fd_idx != u32::MAX;
1118 let target_fd = if use_fixed {
1119 direct_fd_idx as i32
1120 } else {
1121 fd as i32
1122 };
1123 let mut s = io_uring::opcode::Accept::new(
1124 io_uring::types::Fd(target_fd),
1125 std::ptr::null_mut(),
1126 std::ptr::null_mut(),
1127 )
1128 .build()
1129 .user_data(slot_idx as u64);
1130 if use_fixed {
1131 s = s.flags(io_uring::squeue::Flags::FIXED_FILE);
1132 }
1133 s
1134 }
1135 IoRequest::Connect {
1136 fd,
1137 direct_fd_idx,
1138 addr,
1139 addr_len,
1140 slot_idx,
1141 } => {
1142 let addr_ptr = (&raw const addr).cast::<libc::sockaddr>();
1147
1148 let use_fixed = direct_fd_idx != u32::MAX;
1149 let target_fd = if use_fixed {
1150 direct_fd_idx as i32
1151 } else {
1152 fd as i32
1153 };
1154 let mut s =
1155 io_uring::opcode::Connect::new(io_uring::types::Fd(target_fd), addr_ptr, addr_len)
1156 .build()
1157 .user_data(slot_idx as u64);
1158 if use_fixed {
1159 s = s.flags(io_uring::squeue::Flags::FIXED_FILE);
1160 }
1161 s
1162 }
1163 IoRequest::SendTo {
1164 fd,
1165 direct_fd_idx,
1166 msg_ptr,
1167 slot_idx,
1168 } => {
1169 let use_fixed = direct_fd_idx != u32::MAX;
1170 let target_fd = if use_fixed {
1171 direct_fd_idx as i32
1172 } else {
1173 fd as i32
1174 };
1175 let mut s = io_uring::opcode::SendMsg::new(
1176 io_uring::types::Fd(target_fd),
1177 msg_ptr.cast_const(),
1178 )
1179 .build()
1180 .user_data(slot_idx as u64);
1181 if use_fixed {
1182 s = s.flags(io_uring::squeue::Flags::FIXED_FILE);
1183 }
1184 s
1185 }
1186 IoRequest::RecvFrom {
1187 fd,
1188 direct_fd_idx,
1189 msg_ptr,
1190 slot_idx,
1191 } => {
1192 let use_fixed = direct_fd_idx != u32::MAX;
1193 let target_fd = if use_fixed {
1194 direct_fd_idx as i32
1195 } else {
1196 fd as i32
1197 };
1198 let mut s = io_uring::opcode::RecvMsg::new(io_uring::types::Fd(target_fd), msg_ptr)
1199 .build()
1200 .user_data(slot_idx as u64);
1201 if use_fixed {
1202 s = s.flags(io_uring::squeue::Flags::FIXED_FILE);
1203 }
1204 s
1205 }
1206 IoRequest::RegisterFile { fd, slot_idx } => {
1207 if let Some(direct_idx) = state.direct_fd_free.pop() {
1208 let fds = [fd];
1209 let res = ring.submitter().register_files_update(direct_idx, &fds);
1210 let out_res = match res {
1211 Ok(_) => direct_idx as i32,
1212 Err(e) => -(e.raw_os_error().unwrap_or(libc::EINVAL)),
1213 };
1214 process_linux_completion(state, slot_idx, out_res);
1215 } else {
1216 process_linux_completion(state, slot_idx, -libc::ENFILE);
1217 }
1218 return Ok(());
1219 }
1220 IoRequest::UnregisterFile {
1221 direct_fd_idx,
1222 slot_idx,
1223 } => {
1224 let fds = [-1];
1225 let res = ring.submitter().register_files_update(direct_fd_idx, &fds);
1226 state.direct_fd_free.push(direct_fd_idx);
1227 let out_res = match res {
1228 Ok(_) => 0,
1229 Err(e) => -(e.raw_os_error().unwrap_or(libc::EINVAL)),
1230 };
1231 process_linux_completion(state, slot_idx, out_res);
1232 return Ok(());
1233 }
1234 };
1235
1236 let user_data = sqe.get_user_data();
1237 let r = unsafe { push_sqe(ring, &sqe) };
1238 io_trace!(
1239 "[dtact-io] t={} slot={} submit_linux_request pushed_local ok={}",
1240 trace_now_us(),
1241 user_data,
1242 r.is_ok()
1243 );
1244 r
1245}
1246
1247#[cfg(target_os = "linux")]
1248fn process_linux_completion(state: &WorkerState, slot_idx: usize, res: i32) {
1249 let slot = &state.slots[slot_idx];
1250
1251 io_trace!(
1252 "[dtact-io] t={} slot={} res={} B_kernel_complete",
1253 trace_now_us(),
1254 slot_idx,
1255 res
1256 );
1257
1258 slot.result.store(res, Ordering::Release);
1259
1260 slot.lock_waker();
1270 let data = slot
1271 .waker_data
1272 .swap(std::ptr::null_mut(), Ordering::Relaxed);
1273 let vtable = slot
1274 .waker_vtable
1275 .swap(std::ptr::null_mut(), Ordering::Relaxed);
1276 slot.unlock_waker();
1277
1278 slot.completed.store(true, Ordering::Release);
1279
1280 if slot.dropped.load(Ordering::Acquire) {
1281 state.free_slots.push(slot_idx as u32);
1282 wake_next_waiting_fiber(state);
1283 } else if !data.is_null() && !vtable.is_null() {
1284 let raw = RawWaker::new(data.cast_const(), unsafe { &*vtable });
1285 let w = unsafe { Waker::from_raw(raw) };
1286 w.wake();
1287 }
1288}
1289
1290#[cfg(not(target_os = "linux"))]
1294struct FdState {
1295 reader_waker: Option<Waker>,
1296 writer_waker: Option<Waker>,
1297 reader_slot: Option<usize>,
1300 writer_slot: Option<usize>,
1301 registered_interest: Option<mio::Interest>,
1306}
1307
1308#[cfg(not(target_os = "linux"))]
1309impl FdState {
1310 const fn new() -> Self {
1311 Self {
1312 reader_waker: None,
1313 writer_waker: None,
1314 reader_slot: None,
1315 writer_slot: None,
1316 registered_interest: None,
1317 }
1318 }
1319}
1320
1321#[cfg(not(target_os = "linux"))]
1324fn ensure_fd_state(fd_states: &mut Vec<FdState>, fd: usize) {
1325 if fd_states.len() <= fd {
1326 fd_states.resize_with(fd + 1, FdState::new);
1327 }
1328}
1329
1330#[cfg(not(target_os = "linux"))]
1338fn install_interest(
1339 state: &WorkerState,
1340 fd_state: &mut FdState,
1341 fd: u32,
1342 slot_idx: usize,
1343 is_reader: bool,
1344) -> bool {
1345 let slot = &state.slots[slot_idx];
1346 slot.lock_waker();
1347 let data = slot
1348 .waker_data
1349 .swap(std::ptr::null_mut(), Ordering::Relaxed);
1350 let vtable = slot
1351 .waker_vtable
1352 .swap(std::ptr::null_mut(), Ordering::Relaxed);
1353 slot.unlock_waker();
1354
1355 let waker = if !data.is_null() && !vtable.is_null() {
1356 let raw = RawWaker::new(data as *const (), unsafe { &*vtable });
1357 Some(unsafe { Waker::from_raw(raw) })
1358 } else {
1359 None
1360 };
1361
1362 if is_reader {
1363 fd_state.reader_waker = waker;
1364 fd_state.reader_slot = if fd_state.reader_waker.is_some() {
1365 Some(slot_idx)
1366 } else {
1367 None
1368 };
1369 } else {
1370 fd_state.writer_waker = waker;
1371 fd_state.writer_slot = if fd_state.writer_waker.is_some() {
1372 Some(slot_idx)
1373 } else {
1374 None
1375 };
1376 }
1377
1378 let interest = get_mio_interest(fd_state);
1379 if fd_state.registered_interest == Some(interest) {
1380 return true;
1381 }
1382
1383 let res = unsafe {
1384 let poll = &mut *state.poll.get();
1385 poll.registry().reregister(
1386 &mut mio::unix::SourceFd(&(fd as i32)),
1387 mio::Token(fd as usize),
1388 interest,
1389 )
1390 };
1391
1392 match res {
1393 Ok(()) => {
1394 fd_state.registered_interest = Some(interest);
1395 true
1396 }
1397 Err(e) => {
1398 io_trace!(
1399 "[dtact-io] t={} fd={} reregister failed: {e}",
1400 trace_now_us(),
1401 fd
1402 );
1403 let woken = if is_reader {
1408 fd_state.reader_waker.take()
1409 } else {
1410 fd_state.writer_waker.take()
1411 };
1412 if is_reader {
1413 fd_state.reader_slot = None;
1414 } else {
1415 fd_state.writer_slot = None;
1416 }
1417 if let Some(w) = woken {
1418 w.wake();
1419 }
1420 false
1421 }
1422 }
1423}
1424
1425#[cfg(not(target_os = "linux"))]
1426fn get_mio_interest(fd_state: &FdState) -> mio::Interest {
1427 let r = fd_state.reader_waker.is_some();
1428 let w = fd_state.writer_waker.is_some();
1429 if r && w {
1430 mio::Interest::READABLE | mio::Interest::WRITABLE
1431 } else if r {
1432 mio::Interest::READABLE
1433 } else if w {
1434 mio::Interest::WRITABLE
1435 } else {
1436 mio::Interest::READABLE
1437 }
1438}
1439
1440#[cfg(not(target_os = "linux"))]
1441fn run_mio_worker_loop(worker_idx: usize, state: &WorkerState) {
1442 if let Some(config) = GLOBAL_CONFIG.get() {
1443 if let Some(&cpu_id) = config.pin_cpus.get(worker_idx) {
1444 let _ = pin_thread_to_cpu(cpu_id);
1445 }
1446 }
1447
1448 let poll = unsafe { &mut *state.poll.get() };
1449 let mut events = mio::Events::with_capacity(256);
1450 let mut fd_states: Vec<FdState> = Vec::with_capacity(256);
1454
1455 loop {
1456 if SHUTDOWN.load(Ordering::Relaxed) {
1457 break;
1458 }
1459
1460 let mut processed_any = false;
1461 for q in state.queues.iter() {
1462 while let Some(req) = q.pop() {
1463 processed_any = true;
1464 process_mio_request(state, &mut fd_states, req);
1465 }
1466 }
1467
1468 while let Some(slot_idx) = state.cancel_queue.pop() {
1474 processed_any = true;
1475 cancel_mio_slot(state, &mut fd_states, slot_idx as usize);
1476 }
1477
1478 state.is_sleeping.store(true, Ordering::SeqCst);
1479 fence(Ordering::SeqCst);
1480 let mut any_pending = false;
1481 for q in state.queues.iter() {
1482 if !q.is_empty() {
1483 any_pending = true;
1484 break;
1485 }
1486 }
1487
1488 let poll_res = if !any_pending {
1489 poll.poll(&mut events, Some(std::time::Duration::from_millis(10)))
1490 } else {
1491 poll.poll(&mut events, Some(std::time::Duration::from_millis(0)))
1492 };
1493 state.is_sleeping.store(false, Ordering::Release);
1494
1495 if poll_res.is_err() {
1496 continue;
1497 }
1498
1499 for event in events.iter() {
1500 let token = event.token();
1501 if token == mio::Token(0) {
1502 continue;
1503 }
1504 let fd = token.0;
1505 process_mio_event(
1506 state,
1507 &mut fd_states,
1508 fd,
1509 event.is_readable(),
1510 event.is_writable(),
1511 );
1512 }
1513 }
1514}
1515
1516#[cfg(not(target_os = "linux"))]
1517fn process_mio_request(state: &WorkerState, fd_states: &mut Vec<FdState>, req: IoRequest) {
1518 match req {
1519 IoRequest::Read { fd, slot_idx, .. }
1520 | IoRequest::Accept { fd, slot_idx, .. }
1521 | IoRequest::RecvFrom { fd, slot_idx, .. } => {
1522 ensure_fd_state(fd_states, fd as usize);
1523 install_interest(state, &mut fd_states[fd as usize], fd, slot_idx, true);
1524 }
1525 IoRequest::Write { fd, slot_idx, .. }
1526 | IoRequest::Connect { fd, slot_idx, .. }
1527 | IoRequest::SendTo { fd, slot_idx, .. } => {
1528 ensure_fd_state(fd_states, fd as usize);
1529 install_interest(state, &mut fd_states[fd as usize], fd, slot_idx, false);
1530 }
1531 IoRequest::RegisterFile { fd, slot_idx } => {
1532 let res = unsafe {
1533 let poll = &mut *state.poll.get();
1534 poll.registry().register(
1535 &mut mio::unix::SourceFd(&fd),
1536 mio::Token(fd as usize),
1537 mio::Interest::READABLE | mio::Interest::WRITABLE,
1538 )
1539 };
1540 match res {
1541 Ok(()) => complete_mio_slot(state, slot_idx, fd),
1542 Err(e) => {
1543 let os_err = e.raw_os_error().unwrap_or(libc::EINVAL);
1544 complete_mio_slot(state, slot_idx, -os_err);
1545 }
1546 }
1547 }
1548 IoRequest::UnregisterFile {
1549 direct_fd_idx,
1550 slot_idx,
1551 } => {
1552 let _ = unsafe {
1553 let poll = &mut *state.poll.get();
1554 poll.registry()
1555 .deregister(&mut mio::unix::SourceFd(&(direct_fd_idx as i32)))
1556 };
1557 if let Some(fd_state) = fd_states.get_mut(direct_fd_idx as usize) {
1558 fd_state.reader_waker = None;
1559 fd_state.writer_waker = None;
1560 fd_state.reader_slot = None;
1561 fd_state.writer_slot = None;
1562 fd_state.registered_interest = None;
1563 }
1564 complete_mio_slot(state, slot_idx, 0);
1565 }
1566 }
1567}
1568
1569#[cfg(not(target_os = "linux"))]
1574fn cancel_mio_slot(state: &WorkerState, fd_states: &mut Vec<FdState>, slot_idx: usize) {
1575 let slot = &state.slots[slot_idx];
1576 let fd = slot.origin_fd.load(Ordering::Relaxed);
1577
1578 if fd != u32::MAX {
1579 ensure_fd_state(fd_states, fd as usize);
1580 let fd_state = &mut fd_states[fd as usize];
1581 let mut touched = false;
1582 if fd_state.reader_slot == Some(slot_idx) {
1583 fd_state.reader_waker = None;
1584 fd_state.reader_slot = None;
1585 touched = true;
1586 }
1587 if fd_state.writer_slot == Some(slot_idx) {
1588 fd_state.writer_waker = None;
1589 fd_state.writer_slot = None;
1590 touched = true;
1591 }
1592 if touched {
1593 let interest = get_mio_interest(fd_state);
1594 if fd_state.registered_interest != Some(interest) {
1595 let res = unsafe {
1596 let poll = &mut *state.poll.get();
1597 poll.registry().reregister(
1598 &mut mio::unix::SourceFd(&(fd as i32)),
1599 mio::Token(fd as usize),
1600 interest,
1601 )
1602 };
1603 if res.is_ok() {
1604 fd_state.registered_interest = Some(interest);
1605 }
1606 }
1607 }
1608 }
1609
1610 state.free_slots.push(slot_idx as u32);
1611 wake_next_waiting_fiber(state);
1612}
1613
1614#[cfg(not(target_os = "linux"))]
1615fn process_mio_event(
1616 _state: &WorkerState,
1617 fd_states: &mut Vec<FdState>,
1618 fd: usize,
1619 readable: bool,
1620 writable: bool,
1621) {
1622 ensure_fd_state(fd_states, fd);
1623 let fd_state = &mut fd_states[fd];
1624
1625 if readable {
1626 if let Some(w) = fd_state.reader_waker.take() {
1627 w.wake();
1628 }
1629 fd_state.reader_slot = None;
1630 }
1631 if writable {
1632 if let Some(w) = fd_state.writer_waker.take() {
1633 w.wake();
1634 }
1635 fd_state.writer_slot = None;
1636 }
1637
1638 let interest = get_mio_interest(fd_state);
1639 if fd_state.registered_interest != Some(interest) {
1640 let res = unsafe {
1641 let poll = &mut *_state.poll.get();
1642 poll.registry().reregister(
1643 &mut mio::unix::SourceFd(&(fd as i32)),
1644 mio::Token(fd),
1645 interest,
1646 )
1647 };
1648 if res.is_ok() {
1649 fd_state.registered_interest = Some(interest);
1650 }
1651 }
1652}
1653
1654#[cfg(not(target_os = "linux"))]
1655fn complete_mio_slot(state: &WorkerState, slot_idx: usize, res: i32) {
1656 let slot = &state.slots[slot_idx];
1657 slot.result.store(res, Ordering::Release);
1658
1659 slot.lock_waker();
1664 let data = slot
1665 .waker_data
1666 .swap(std::ptr::null_mut(), Ordering::Relaxed);
1667 let vtable = slot
1668 .waker_vtable
1669 .swap(std::ptr::null_mut(), Ordering::Relaxed);
1670 slot.unlock_waker();
1671
1672 slot.completed.store(true, Ordering::Release);
1673
1674 if !data.is_null() && !vtable.is_null() {
1675 let raw = RawWaker::new(data as *const (), unsafe { &*vtable });
1676 let w = unsafe { Waker::from_raw(raw) };
1677 w.wake();
1678 }
1679}
1680
1681pub struct DtactIoFuture {
1691 pub worker_idx: usize,
1693 pub fd: u32,
1695 pub direct_fd_idx: u32,
1698 pub op: OpCode,
1700 pub buf_ptr: *mut u8,
1702 pub len: usize,
1704 pub offset: i64,
1706 pub addr: Option<libc::sockaddr_storage>,
1708 pub addr_len: libc::socklen_t,
1710 pub slot_idx: Option<usize>,
1713 pub msg_ptr: *mut libc::msghdr,
1716}
1717
1718unsafe impl Send for DtactIoFuture {}
1719unsafe impl Sync for DtactIoFuture {}
1720
1721impl DtactIoFuture {
1722 #[allow(clippy::too_many_arguments)]
1725 pub const fn new(
1726 worker_idx: usize,
1727 fd: u32,
1728 direct_fd_idx: u32,
1729 op: OpCode,
1730 buf_ptr: *mut u8,
1731 len: usize,
1732 offset: i64,
1733 addr: Option<libc::sockaddr_storage>,
1734 addr_len: libc::socklen_t,
1735 slot_idx: Option<usize>,
1736 ) -> Self {
1737 Self {
1738 worker_idx,
1739 fd,
1740 direct_fd_idx,
1741 op,
1742 buf_ptr,
1743 len,
1744 offset,
1745 addr,
1746 addr_len,
1747 slot_idx,
1748 msg_ptr: std::ptr::null_mut(),
1749 }
1750 }
1751
1752 const fn create_io_request(&self, slot_idx: usize) -> IoRequest {
1753 match self.op {
1754 OpCode::SendTo => IoRequest::SendTo {
1755 fd: self.fd,
1756 direct_fd_idx: self.direct_fd_idx,
1757 msg_ptr: self.msg_ptr,
1758 slot_idx,
1759 },
1760 OpCode::RecvFrom => IoRequest::RecvFrom {
1761 fd: self.fd,
1762 direct_fd_idx: self.direct_fd_idx,
1763 msg_ptr: self.msg_ptr,
1764 slot_idx,
1765 },
1766 OpCode::Read => IoRequest::Read {
1767 fd: self.fd,
1768 direct_fd_idx: self.direct_fd_idx,
1769 buf_ptr: self.buf_ptr,
1770 len: self.len,
1771 offset: self.offset,
1772 slot_idx,
1773 },
1774 OpCode::Write => IoRequest::Write {
1775 fd: self.fd,
1776 direct_fd_idx: self.direct_fd_idx,
1777 buf_ptr: self.buf_ptr,
1778 len: self.len,
1779 offset: self.offset,
1780 slot_idx,
1781 },
1782 OpCode::Accept => IoRequest::Accept {
1783 fd: self.fd,
1784 direct_fd_idx: self.direct_fd_idx,
1785 slot_idx,
1786 },
1787 OpCode::Connect => IoRequest::Connect {
1788 fd: self.fd,
1789 direct_fd_idx: self.direct_fd_idx,
1790 addr: self.addr.unwrap(),
1791 addr_len: self.addr_len,
1792 slot_idx,
1793 },
1794 }
1795 }
1796
1797 #[cfg(not(target_os = "linux"))]
1798 fn execute_syscall(&self) -> std::io::Result<usize> {
1799 let res = match self.op {
1800 OpCode::Read => {
1801 let buf_ptr = self.buf_ptr;
1802 let len = self.len;
1803 unsafe { libc::read(self.fd as i32, buf_ptr as *mut libc::c_void, len) }
1804 }
1805 OpCode::Write => {
1806 let buf_ptr = self.buf_ptr;
1807 let len = self.len;
1808 unsafe { libc::write(self.fd as i32, buf_ptr as *const libc::c_void, len) }
1809 }
1810 OpCode::Accept => unsafe {
1811 libc::accept(self.fd as i32, std::ptr::null_mut(), std::ptr::null_mut()) as isize
1812 },
1813 OpCode::SendTo => unsafe { libc::sendmsg(self.fd as i32, self.msg_ptr, 0) },
1814 OpCode::RecvFrom => unsafe { libc::recvmsg(self.fd as i32, self.msg_ptr, 0) },
1815 OpCode::Connect => {
1816 let addr_ptr =
1817 &self.addr.unwrap() as *const libc::sockaddr_storage as *const libc::sockaddr;
1818 let res = unsafe { libc::connect(self.fd as i32, addr_ptr, self.addr_len) };
1819 if res < 0 {
1820 let err = std::io::Error::last_os_error();
1821 if err.raw_os_error() == Some(libc::EISCONN) {
1822 return Ok(0);
1823 }
1824 return Err(err);
1825 }
1826 res as isize
1827 }
1828 };
1829
1830 if res < 0 {
1831 Err(std::io::Error::last_os_error())
1832 } else {
1833 Ok(res as usize)
1834 }
1835 }
1836}
1837
1838impl Future for DtactIoFuture {
1839 type Output = std::io::Result<usize>;
1840
1841 #[allow(clippy::too_many_lines)]
1847 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1848 #[cfg(target_os = "linux")]
1849 {
1850 let slot_idx = if let Some(idx) = self.slot_idx {
1851 idx
1852 } else {
1853 let state = &WORKERS.get().unwrap()[self.worker_idx];
1854 let idx = match state.free_slots.pop() {
1855 Some(i) => i as usize,
1856 None => {
1857 if let Some(wait_idx) = state.free_wait_slots.pop() {
1858 let wait_slot = &state.wait_slots[wait_idx as usize];
1859 wait_slot
1860 .waker_data
1861 .store(cx.waker().data().cast_mut(), Ordering::Relaxed);
1862 wait_slot.waker_vtable.store(
1863 std::ptr::from_ref::<RawWakerVTable>(cx.waker().vtable())
1864 .cast_mut(),
1865 Ordering::Relaxed,
1866 );
1867 state.waiting_queue.push(wait_idx);
1868
1869 if let Some(i) = state.free_slots.pop() {
1870 wait_slot
1871 .waker_data
1872 .store(std::ptr::null_mut(), Ordering::Relaxed);
1873 wait_slot
1874 .waker_vtable
1875 .store(std::ptr::null_mut(), Ordering::Relaxed);
1876 i as usize
1877 } else {
1878 return Poll::Pending;
1879 }
1880 } else {
1881 cx.waker().wake_by_ref();
1882 return Poll::Pending;
1883 }
1884 }
1885 };
1886
1887 let slot = &state.slots[idx];
1888 slot.completed.store(false, Ordering::Relaxed);
1889 slot.dropped.store(false, Ordering::Relaxed);
1890 slot.origin_fd.store(self.fd, Ordering::Relaxed);
1891 slot.lock_waker();
1893 slot.waker_data
1894 .store(cx.waker().data().cast_mut(), Ordering::Relaxed);
1895 slot.waker_vtable.store(
1896 std::ptr::from_ref::<RawWakerVTable>(cx.waker().vtable()).cast_mut(),
1897 Ordering::Relaxed,
1898 );
1899 slot.unlock_waker();
1900
1901 let req = self.create_io_request(idx);
1902 let q_idx = get_or_init_local_allocator().unwrap_or(0);
1903 let queue = &state.queues[q_idx];
1904
1905 if queue.push(req).is_err() {
1906 slot.lock_waker();
1908 slot.waker_data
1909 .store(std::ptr::null_mut(), Ordering::Relaxed);
1910 slot.waker_vtable
1911 .store(std::ptr::null_mut(), Ordering::Relaxed);
1912 slot.unlock_waker();
1913 state.free_slots.push(idx as u32);
1914 wake_next_waiting_fiber(state);
1915 cx.waker().wake_by_ref();
1916 return Poll::Pending;
1917 }
1918
1919 fence(Ordering::SeqCst);
1929 if state.is_sleeping.load(Ordering::SeqCst) {
1930 unsafe {
1931 let _ = libc::write(
1932 state.wake_eventfd,
1933 std::ptr::from_ref::<u64>(&1u64).cast::<libc::c_void>(),
1934 8,
1935 );
1936 }
1937 }
1938
1939 io_trace!(
1940 "[dtact-io] t={} slot={} fd={} op={:?} A_submit",
1941 trace_now_us(),
1942 idx,
1943 self.fd,
1944 self.op
1945 );
1946
1947 self.slot_idx = Some(idx);
1948 idx
1949 };
1950
1951 let state = &WORKERS.get().unwrap()[self.worker_idx];
1952 let slot = &state.slots[slot_idx];
1953
1954 if slot.completed.load(Ordering::Acquire) {
1955 let res = slot.result.load(Ordering::Acquire);
1956 io_trace!(
1957 "[dtact-io] t={} slot={} res={} C_fiber_resumed",
1958 trace_now_us(),
1959 slot_idx,
1960 res
1961 );
1962 slot.lock_waker();
1964 slot.waker_data
1965 .store(std::ptr::null_mut(), Ordering::Relaxed);
1966 slot.waker_vtable
1967 .store(std::ptr::null_mut(), Ordering::Relaxed);
1968 slot.unlock_waker();
1969 state.free_slots.push(slot_idx as u32);
1970 self.slot_idx = None;
1971
1972 wake_next_waiting_fiber(state);
1973
1974 if res < 0 {
1975 Poll::Ready(Err(std::io::Error::from_raw_os_error(-res)))
1976 } else {
1977 Poll::Ready(Ok(res as usize))
1978 }
1979 } else {
1980 let new_data = cx.waker().data().cast_mut();
1983 let new_vtable =
1984 std::ptr::from_ref::<RawWakerVTable>(cx.waker().vtable()).cast_mut();
1985
1986 slot.lock_waker();
1987 let old_data = slot.waker_data.load(Ordering::Relaxed);
1988 let old_vtable = slot.waker_vtable.load(Ordering::Relaxed);
1989 if old_data != new_data || old_vtable != new_vtable {
1990 slot.waker_data.store(new_data, Ordering::Relaxed);
1991 slot.waker_vtable.store(new_vtable, Ordering::Relaxed);
1992 }
1993 slot.unlock_waker();
1994 Poll::Pending
1995 }
1996 }
1997
1998 #[cfg(not(target_os = "linux"))]
1999 {
2000 let res = self.execute_syscall();
2001 if self.slot_idx.is_some()
2002 && !matches!(res, Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock)
2003 {
2004 let state = &WORKERS.get().unwrap()[self.worker_idx];
2005 state.free_slots.push(self.slot_idx.unwrap() as u32);
2006 self.slot_idx = None;
2007 wake_next_waiting_fiber(state);
2008 }
2009
2010 match res {
2011 Ok(n) => Poll::Ready(Ok(n)),
2012 Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
2013 let slot_idx = match self.slot_idx {
2014 Some(idx) => idx,
2015 None => {
2016 let state = &WORKERS.get().unwrap()[self.worker_idx];
2017 let idx = match state.free_slots.pop() {
2018 Some(i) => i as usize,
2019 None => {
2020 if let Some(wait_idx) = state.free_wait_slots.pop() {
2021 let wait_slot = &state.wait_slots[wait_idx as usize];
2022 wait_slot
2023 .waker_data
2024 .store(cx.waker().data() as *mut (), Ordering::Relaxed);
2025 wait_slot.waker_vtable.store(
2026 cx.waker().vtable() as *const RawWakerVTable as *mut _,
2027 Ordering::Relaxed,
2028 );
2029 state.waiting_queue.push(wait_idx);
2030
2031 if let Some(i) = state.free_slots.pop() {
2032 wait_slot
2033 .waker_data
2034 .store(std::ptr::null_mut(), Ordering::Relaxed);
2035 wait_slot
2036 .waker_vtable
2037 .store(std::ptr::null_mut(), Ordering::Relaxed);
2038 i as usize
2039 } else {
2040 return Poll::Pending;
2041 }
2042 } else {
2043 cx.waker().wake_by_ref();
2044 return Poll::Pending;
2045 }
2046 }
2047 };
2048
2049 let slot = &state.slots[idx];
2050 slot.completed.store(false, Ordering::Relaxed);
2051 slot.dropped.store(false, Ordering::Relaxed);
2052 slot.origin_fd.store(self.fd, Ordering::Relaxed);
2053 let raw = cx.waker().as_raw();
2054 slot.lock_waker();
2055 slot.waker_data
2056 .store(raw.data() as *mut (), Ordering::Relaxed);
2057 slot.waker_vtable.store(
2058 raw.vtable() as *const RawWakerVTable as *mut _,
2059 Ordering::Relaxed,
2060 );
2061 slot.unlock_waker();
2062
2063 let req = self.create_io_request(idx);
2064 let q_idx = get_or_init_local_allocator().unwrap_or(0);
2065 let queue = &state.queues[q_idx];
2066
2067 if queue.push(req).is_err() {
2068 slot.lock_waker();
2069 slot.waker_data
2070 .store(std::ptr::null_mut(), Ordering::Relaxed);
2071 slot.waker_vtable
2072 .store(std::ptr::null_mut(), Ordering::Relaxed);
2073 slot.unlock_waker();
2074 state.free_slots.push(idx as u32);
2075 wake_next_waiting_fiber(state);
2076 cx.waker().wake_by_ref();
2077 return Poll::Pending;
2078 }
2079
2080 fence(Ordering::SeqCst);
2081 if state.is_sleeping.load(Ordering::SeqCst) {
2082 state.waker.wake();
2083 }
2084 self.slot_idx = Some(idx);
2085 idx
2086 }
2087 };
2088
2089 let state = &WORKERS.get().unwrap()[self.worker_idx];
2090 let slot = &state.slots[slot_idx];
2091 let raw = cx.waker().as_raw();
2092 let new_data = raw.data() as *mut ();
2093 let new_vtable = raw.vtable() as *const RawWakerVTable as *mut _;
2094
2095 slot.lock_waker();
2096 let old_data = slot.waker_data.load(Ordering::Relaxed);
2097 let old_vtable = slot.waker_vtable.load(Ordering::Relaxed);
2098 let mut changed = false;
2099 if old_data != new_data || old_vtable != new_vtable {
2100 slot.waker_data.store(new_data, Ordering::Relaxed);
2101 slot.waker_vtable.store(new_vtable, Ordering::Relaxed);
2102 changed = true;
2103 }
2104 slot.unlock_waker();
2105
2106 if changed {
2107 let req = self.create_io_request(slot_idx);
2108 let q_idx = get_or_init_local_allocator().unwrap_or(0);
2109 let _ = state.queues[q_idx].push(req);
2110 fence(Ordering::SeqCst);
2111 if state.is_sleeping.load(Ordering::SeqCst) {
2112 state.waker.wake();
2113 }
2114 }
2115 Poll::Pending
2116 }
2117 Err(e) => Poll::Ready(Err(e)),
2118 }
2119 }
2120 }
2121}
2122
2123impl Drop for DtactIoFuture {
2124 fn drop(&mut self) {
2125 let Some(idx) = self.slot_idx else { return };
2126 let Some(state) = WORKERS.get().and_then(|w| w.get(self.worker_idx)) else {
2127 return;
2128 };
2129
2130 let slot = &state.slots[idx];
2133 slot.lock_waker();
2134 slot.waker_data
2135 .store(std::ptr::null_mut(), Ordering::Relaxed);
2136 slot.waker_vtable
2137 .store(std::ptr::null_mut(), Ordering::Relaxed);
2138 slot.unlock_waker();
2139
2140 if slot.completed.load(Ordering::Acquire) {
2141 state.free_slots.push(idx as u32);
2145 wake_next_waiting_fiber(state);
2146 return;
2147 }
2148
2149 slot.dropped.store(true, Ordering::Release);
2158 state.cancel_queue.push(idx as u32);
2159 fence(Ordering::SeqCst);
2160
2161 #[cfg(target_os = "linux")]
2162 {
2163 if state.is_sleeping.load(Ordering::SeqCst) {
2164 unsafe {
2165 let _ = libc::write(
2166 state.wake_eventfd,
2167 std::ptr::from_ref::<u64>(&1u64).cast::<libc::c_void>(),
2168 8,
2169 );
2170 }
2171 }
2172 }
2173 #[cfg(not(target_os = "linux"))]
2174 {
2175 if state.is_sleeping.load(Ordering::SeqCst) {
2176 state.waker.wake();
2177 }
2178 }
2179 }
2180}
2181
2182pub struct DtactTcpStream {
2189 inner: std::net::TcpStream,
2190 direct_fd_idx: u32,
2191 worker_idx: usize,
2192}
2193
2194impl DtactTcpStream {
2195 pub fn from_std(stream: std::net::TcpStream) -> std::io::Result<Self> {
2213 let fd = stream.as_raw_fd();
2214 stream.set_nonblocking(true)?;
2215 stream.set_nodelay(true)?;
2222
2223 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
2224 let worker_idx = fd as usize % num_workers;
2225 let state = &WORKERS.get().unwrap()[worker_idx];
2226
2227 let direct_fd_idx = register_fd_sync(state, fd);
2228
2229 Ok(Self {
2230 inner: stream,
2231 direct_fd_idx,
2232 worker_idx,
2233 })
2234 }
2235
2236 pub async fn read(&self, buf: &mut [u8]) -> std::io::Result<usize> {
2244 if buf.is_empty() {
2245 return Ok(0);
2246 }
2247
2248 let res = unsafe {
2259 let r = libc::read(
2260 self.inner.as_raw_fd(),
2261 buf.as_mut_ptr().cast::<libc::c_void>(),
2262 buf.len(),
2263 );
2264 match r.cmp(&0) {
2265 std::cmp::Ordering::Greater => Ok(r as usize),
2266 std::cmp::Ordering::Equal => Ok(0), std::cmp::Ordering::Less => Err(std::io::Error::last_os_error()),
2268 }
2269 };
2270
2271 match res {
2272 Ok(n) => return Ok(n),
2273 Err(e) => {
2274 if e.kind() != std::io::ErrorKind::WouldBlock {
2275 return Err(e);
2276 }
2277 }
2278 }
2279
2280 DtactIoFuture {
2282 worker_idx: self.worker_idx,
2283 fd: self.inner.as_raw_fd() as u32,
2284 direct_fd_idx: self.direct_fd_idx,
2285 op: OpCode::Read,
2286 buf_ptr: buf.as_mut_ptr(),
2287 len: buf.len(),
2288 offset: 0,
2289 addr: None,
2290 addr_len: 0,
2291 slot_idx: None,
2292 msg_ptr: std::ptr::null_mut(),
2293 }
2294 .await
2295 .map(|n| n.min(buf.len()))
2296 }
2297
2298 pub async fn write(&self, buf: &[u8]) -> std::io::Result<usize> {
2307 if buf.is_empty() {
2308 return Ok(0);
2309 }
2310
2311 let res = unsafe {
2314 let r = libc::write(
2315 self.inner.as_raw_fd(),
2316 buf.as_ptr().cast::<libc::c_void>(),
2317 buf.len(),
2318 );
2319 if r >= 0 {
2320 Ok(r as usize)
2321 } else {
2322 Err(std::io::Error::last_os_error())
2323 }
2324 };
2325
2326 match res {
2327 Ok(n) => return Ok(n),
2328 Err(e) => {
2329 if e.kind() != std::io::ErrorKind::WouldBlock {
2330 return Err(e);
2331 }
2332 }
2333 }
2334
2335 DtactIoFuture {
2336 worker_idx: self.worker_idx,
2337 fd: self.inner.as_raw_fd() as u32,
2338 direct_fd_idx: self.direct_fd_idx,
2339 op: OpCode::Write,
2340 buf_ptr: buf.as_ptr().cast_mut(),
2341 len: buf.len(),
2342 offset: 0,
2343 addr: None,
2344 addr_len: 0,
2345 slot_idx: None,
2346 msg_ptr: std::ptr::null_mut(),
2347 }
2348 .await
2349 }
2350
2351 #[allow(clippy::too_many_lines)]
2371 pub async fn connect(addr: std::net::SocketAddr) -> std::io::Result<Self> {
2372 let domain = match addr {
2373 std::net::SocketAddr::V4(_) => libc::AF_INET,
2374 std::net::SocketAddr::V6(_) => libc::AF_INET6,
2375 };
2376 let fd = unsafe {
2377 libc::socket(
2378 domain,
2379 libc::SOCK_STREAM | libc::SOCK_CLOEXEC | libc::SOCK_NONBLOCK,
2380 0,
2381 )
2382 };
2383 if fd < 0 {
2384 return Err(std::io::Error::last_os_error());
2385 }
2386
2387 let stream = unsafe { std::net::TcpStream::from_raw_fd(fd) };
2389 stream.set_nodelay(true)?;
2392 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
2393 let worker_idx = fd as usize % num_workers;
2394 let state = &WORKERS.get().unwrap()[worker_idx];
2395
2396 let direct_fd_idx = register_fd_sync(state, fd);
2399
2400 let (libc_addr, addr_len) = socket_addr_to_libc(addr);
2401
2402 let connect_res = unsafe {
2404 libc::connect(
2405 fd,
2406 (&raw const libc_addr).cast::<libc::sockaddr>(),
2407 addr_len,
2408 )
2409 };
2410 if connect_res == 0 {
2411 return Ok(Self {
2412 inner: stream,
2413 direct_fd_idx,
2414 worker_idx,
2415 });
2416 }
2417 let err = std::io::Error::last_os_error();
2418 #[cfg(target_os = "windows")]
2419 let is_in_progress = err.raw_os_error()
2420 == Some(windows_sys::Win32::Networking::WinSock::WSAEWOULDBLOCK as i32);
2421 #[cfg(not(target_os = "windows"))]
2422 let is_in_progress = err.raw_os_error() == Some(libc::EINPROGRESS);
2423
2424 if !is_in_progress {
2425 return Err(err);
2426 }
2427
2428 let mut pollfd = libc::pollfd {
2433 fd,
2434 events: libc::POLLOUT,
2435 revents: 0,
2436 };
2437 let poll_res = unsafe { libc::poll(&raw mut pollfd, 1, 0) };
2438 if poll_res > 0 {
2439 if (pollfd.revents & libc::POLLOUT) != 0 {
2440 let mut err_code: libc::c_int = 0;
2441 let mut err_len = std::mem::size_of::<libc::c_int>() as libc::socklen_t;
2442 let sockopt_res = unsafe {
2443 libc::getsockopt(
2444 fd,
2445 libc::SOL_SOCKET,
2446 libc::SO_ERROR,
2447 (&raw mut err_code).cast::<libc::c_void>(),
2448 &raw mut err_len,
2449 )
2450 };
2451 if sockopt_res == 0 && err_code == 0 {
2452 return Ok(Self {
2453 inner: stream,
2454 direct_fd_idx,
2455 worker_idx,
2456 });
2457 }
2458 let os_err = if err_code != 0 {
2459 err_code
2460 } else {
2461 libc::ECONNREFUSED
2462 };
2463 return Err(std::io::Error::from_raw_os_error(os_err));
2464 } else if (pollfd.revents & (libc::POLLERR | libc::POLLHUP)) != 0 {
2465 return Err(std::io::Error::new(
2466 std::io::ErrorKind::ConnectionRefused,
2467 "connect failed",
2468 ));
2469 }
2470 }
2471
2472 let connect_res = DtactIoFuture {
2473 worker_idx,
2474 fd: fd as u32,
2475 direct_fd_idx,
2476 op: OpCode::Connect,
2477 buf_ptr: std::ptr::null_mut(),
2478 len: 0,
2479 offset: 0,
2480 addr: Some(libc_addr),
2481 addr_len,
2482 slot_idx: None,
2483 msg_ptr: std::ptr::null_mut(),
2484 }
2485 .await;
2486
2487 match connect_res {
2488 Ok(_) => Ok(Self {
2489 inner: stream,
2490 direct_fd_idx,
2491 worker_idx,
2492 }),
2493 Err(e) => Err(e),
2494 }
2495 }
2496}
2497
2498impl Drop for DtactTcpStream {
2499 fn drop(&mut self) {
2500 if let Some(workers) = WORKERS.get()
2501 && let Some(state) = workers.get(self.worker_idx)
2502 {
2503 unregister_fd_sync(state, self.direct_fd_idx);
2504 }
2505 }
2506}
2507
2508impl crate::io::AsyncRead for DtactTcpStream {
2509 async fn read(&self, buf: &mut [u8]) -> std::io::Result<usize> {
2510 self.read(buf).await
2511 }
2512}
2513
2514impl crate::io::AsyncWrite for DtactTcpStream {
2515 async fn write(&self, buf: &[u8]) -> std::io::Result<usize> {
2516 self.write(buf).await
2517 }
2518}
2519
2520pub struct DtactTcpListener {
2524 inner: std::net::TcpListener,
2525 direct_fd_idx: u32,
2526 worker_idx: usize,
2527}
2528
2529impl DtactTcpListener {
2530 pub fn from_std(listener: std::net::TcpListener) -> std::io::Result<Self> {
2544 let fd = listener.as_raw_fd();
2545 listener.set_nonblocking(true)?;
2546
2547 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
2548 let worker_idx = fd as usize % num_workers;
2549 let state = &WORKERS.get().unwrap()[worker_idx];
2550
2551 let direct_fd_idx = register_fd_sync(state, fd);
2552
2553 Ok(Self {
2554 inner: listener,
2555 direct_fd_idx,
2556 worker_idx,
2557 })
2558 }
2559
2560 pub async fn accept(&self) -> std::io::Result<(DtactTcpStream, std::net::SocketAddr)> {
2569 let res = unsafe {
2575 let mut addr: libc::sockaddr_storage = std::mem::zeroed();
2576 let mut len = std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
2577 let r = libc::accept(
2578 self.inner.as_raw_fd(),
2579 (&raw mut addr).cast::<libc::sockaddr>(),
2580 &raw mut len,
2581 );
2582 if r >= 0 {
2583 Ok((r, addr, len))
2584 } else {
2585 Err(std::io::Error::last_os_error())
2586 }
2587 };
2588
2589 match res {
2590 Ok((client_fd, addr, len)) => {
2591 let peer_addr = sockaddr_storage_to_socketaddr(&addr, len);
2594 unsafe { libc::fcntl(client_fd, libc::F_SETFL, libc::O_NONBLOCK) };
2596 let stream = unsafe { std::net::TcpStream::from_raw_fd(client_fd) };
2597 let client_stream = DtactTcpStream::from_std(stream)?;
2598 return Ok((client_stream, peer_addr));
2599 }
2600 Err(e) => {
2601 if e.kind() != std::io::ErrorKind::WouldBlock {
2602 return Err(e);
2603 }
2604 }
2605 }
2606
2607 let res = DtactIoFuture {
2608 worker_idx: self.worker_idx,
2609 fd: self.inner.as_raw_fd() as u32,
2610 direct_fd_idx: self.direct_fd_idx,
2611 op: OpCode::Accept,
2612 buf_ptr: std::ptr::null_mut(),
2613 len: 0,
2614 offset: 0,
2615 addr: None,
2616 addr_len: 0,
2617 slot_idx: None,
2618 msg_ptr: std::ptr::null_mut(),
2619 }
2620 .await?;
2621
2622 let client_fd = res as RawFd;
2623 unsafe { libc::fcntl(client_fd, libc::F_SETFL, libc::O_NONBLOCK) };
2625 let stream = unsafe { std::net::TcpStream::from_raw_fd(client_fd) };
2626 let peer_addr = stream.peer_addr()?;
2627 let client_stream = DtactTcpStream::from_std(stream)?;
2628 Ok((client_stream, peer_addr))
2629 }
2630}
2631
2632impl Drop for DtactTcpListener {
2633 fn drop(&mut self) {
2634 if let Some(workers) = WORKERS.get()
2635 && let Some(state) = workers.get(self.worker_idx)
2636 {
2637 unregister_fd_sync(state, self.direct_fd_idx);
2638 }
2639 }
2640}
2641
2642pub struct DtactUdpSocket {
2654 inner: std::net::UdpSocket,
2655 direct_fd_idx: u32,
2656 worker_idx: usize,
2657}
2658
2659impl DtactUdpSocket {
2660 pub async fn bind(addr: std::net::SocketAddr) -> std::io::Result<Self> {
2665 let sock = std::net::UdpSocket::bind(addr)?;
2666 Self::from_std(sock)
2667 }
2668
2669 pub fn from_std(socket: std::net::UdpSocket) -> std::io::Result<Self> {
2680 let fd = socket.as_raw_fd();
2681 socket.set_nonblocking(true)?;
2682 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
2683 let worker_idx = fd as usize % num_workers;
2684 let state = &WORKERS.get().unwrap()[worker_idx];
2685 let direct_fd_idx = register_fd_sync(state, fd);
2686 Ok(Self {
2687 inner: socket,
2688 direct_fd_idx,
2689 worker_idx,
2690 })
2691 }
2692
2693 pub fn local_addr(&self) -> std::io::Result<std::net::SocketAddr> {
2698 self.inner.local_addr()
2699 }
2700
2701 pub async fn send_to(
2707 &self,
2708 buf: &[u8],
2709 target: std::net::SocketAddr,
2710 ) -> std::io::Result<usize> {
2711 struct SendToState {
2722 storage: libc::sockaddr_storage,
2723 iov: libc::iovec,
2724 msg: libc::msghdr,
2725 }
2726 unsafe impl Send for SendToState {}
2734
2735 let (storage, addr_len) = socket_addr_to_libc(target);
2736 let mut state = SendToState {
2737 storage,
2738 iov: libc::iovec {
2739 iov_base: buf.as_ptr().cast_mut().cast::<libc::c_void>(),
2740 iov_len: buf.len(),
2741 },
2742 msg: unsafe { std::mem::zeroed() },
2743 };
2744 state.msg.msg_name = std::ptr::addr_of_mut!(state.storage).cast::<libc::c_void>();
2745 state.msg.msg_namelen = addr_len;
2746 state.msg.msg_iov = &raw mut state.iov;
2747 state.msg.msg_iovlen = 1;
2748
2749 let r = unsafe { libc::sendmsg(self.inner.as_raw_fd(), &raw const state.msg, 0) };
2751 if r >= 0 {
2752 return Ok(r as usize);
2753 }
2754 let e = std::io::Error::last_os_error();
2755 if e.kind() != std::io::ErrorKind::WouldBlock {
2756 return Err(e);
2757 }
2758
2759 let mut fut = DtactIoFuture::new(
2760 self.worker_idx,
2761 self.inner.as_raw_fd() as u32,
2762 self.direct_fd_idx,
2763 OpCode::SendTo,
2764 std::ptr::null_mut(),
2765 0,
2766 0,
2767 None,
2768 0,
2769 None,
2770 );
2771 fut.msg_ptr = &raw mut state.msg;
2772 fut.await
2773 }
2774
2775 pub async fn recv_from(
2781 &self,
2782 buf: &mut [u8],
2783 ) -> std::io::Result<(usize, std::net::SocketAddr)> {
2784 struct RecvFromState {
2787 storage: libc::sockaddr_storage,
2788 iov: libc::iovec,
2789 msg: libc::msghdr,
2790 }
2791 unsafe impl Send for RecvFromState {}
2795
2796 let mut state = RecvFromState {
2797 storage: unsafe { std::mem::zeroed() },
2798 iov: libc::iovec {
2799 iov_base: buf.as_mut_ptr().cast::<libc::c_void>(),
2800 iov_len: buf.len(),
2801 },
2802 msg: unsafe { std::mem::zeroed() },
2803 };
2804 state.msg.msg_name = std::ptr::addr_of_mut!(state.storage).cast::<libc::c_void>();
2805 state.msg.msg_namelen = std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
2806 state.msg.msg_iov = &raw mut state.iov;
2807 state.msg.msg_iovlen = 1;
2808
2809 let r = unsafe { libc::recvmsg(self.inner.as_raw_fd(), &raw mut state.msg, 0) };
2810 if r >= 0 {
2811 let from = sockaddr_storage_to_socketaddr(&state.storage, state.msg.msg_namelen);
2812 return Ok((r as usize, from));
2813 }
2814 let e = std::io::Error::last_os_error();
2815 if e.kind() != std::io::ErrorKind::WouldBlock {
2816 return Err(e);
2817 }
2818
2819 let mut fut = DtactIoFuture::new(
2820 self.worker_idx,
2821 self.inner.as_raw_fd() as u32,
2822 self.direct_fd_idx,
2823 OpCode::RecvFrom,
2824 std::ptr::null_mut(),
2825 0,
2826 0,
2827 None,
2828 0,
2829 None,
2830 );
2831 fut.msg_ptr = &raw mut state.msg;
2832 let n = fut.await?;
2833 let from = sockaddr_storage_to_socketaddr(&state.storage, state.msg.msg_namelen);
2834 Ok((n, from))
2835 }
2836
2837 pub async fn connect(&self, addr: std::net::SocketAddr) -> std::io::Result<()> {
2844 self.inner.connect(addr)
2845 }
2846
2847 pub async fn send(&self, buf: &[u8]) -> std::io::Result<usize> {
2852 if buf.is_empty() {
2853 return Ok(0);
2854 }
2855 let r = unsafe {
2856 libc::send(
2857 self.inner.as_raw_fd(),
2858 buf.as_ptr().cast::<libc::c_void>(),
2859 buf.len(),
2860 0,
2861 )
2862 };
2863 if r >= 0 {
2864 return Ok(r as usize);
2865 }
2866 let e = std::io::Error::last_os_error();
2867 if e.kind() != std::io::ErrorKind::WouldBlock {
2868 return Err(e);
2869 }
2870 DtactIoFuture {
2871 worker_idx: self.worker_idx,
2872 fd: self.inner.as_raw_fd() as u32,
2873 direct_fd_idx: self.direct_fd_idx,
2874 op: OpCode::Write,
2875 buf_ptr: buf.as_ptr().cast_mut(),
2876 len: buf.len(),
2877 offset: 0,
2878 addr: None,
2879 addr_len: 0,
2880 slot_idx: None,
2881 msg_ptr: std::ptr::null_mut(),
2882 }
2883 .await
2884 }
2885
2886 pub async fn recv(&self, buf: &mut [u8]) -> std::io::Result<usize> {
2892 if buf.is_empty() {
2893 return Ok(0);
2894 }
2895 let r = unsafe {
2896 libc::recv(
2897 self.inner.as_raw_fd(),
2898 buf.as_mut_ptr().cast::<libc::c_void>(),
2899 buf.len(),
2900 0,
2901 )
2902 };
2903 if r >= 0 {
2904 return Ok(r as usize);
2905 }
2906 let e = std::io::Error::last_os_error();
2907 if e.kind() != std::io::ErrorKind::WouldBlock {
2908 return Err(e);
2909 }
2910 DtactIoFuture {
2911 worker_idx: self.worker_idx,
2912 fd: self.inner.as_raw_fd() as u32,
2913 direct_fd_idx: self.direct_fd_idx,
2914 op: OpCode::Read,
2915 buf_ptr: buf.as_mut_ptr(),
2916 len: buf.len(),
2917 offset: 0,
2918 addr: None,
2919 addr_len: 0,
2920 slot_idx: None,
2921 msg_ptr: std::ptr::null_mut(),
2922 }
2923 .await
2924 }
2925}
2926
2927impl Drop for DtactUdpSocket {
2928 fn drop(&mut self) {
2929 if let Some(workers) = WORKERS.get()
2930 && let Some(state) = workers.get(self.worker_idx)
2931 {
2932 unregister_fd_sync(state, self.direct_fd_idx);
2933 }
2934 }
2935}
2936
2937pub struct DtactUnixStream {
2963 inner: std::os::unix::net::UnixStream,
2964 direct_fd_idx: u32,
2965 worker_idx: usize,
2966 read_backpressured: std::sync::atomic::AtomicBool,
2967 write_backpressured: std::sync::atomic::AtomicBool,
2968}
2969
2970impl DtactUnixStream {
2971 pub fn from_std(stream: std::os::unix::net::UnixStream) -> std::io::Result<Self> {
2985 let fd = stream.as_raw_fd();
2986 stream.set_nonblocking(true)?;
2987
2988 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
2989 let worker_idx = WORKER_ROUND_ROBIN.fetch_add(1, Ordering::Relaxed) % num_workers;
2990 let state = &WORKERS.get().unwrap()[worker_idx];
2991
2992 let direct_fd_idx = register_fd_sync(state, fd);
2993
2994 Ok(Self {
2995 inner: stream,
2996 direct_fd_idx,
2997 worker_idx,
2998 read_backpressured: std::sync::atomic::AtomicBool::new(false),
2999 write_backpressured: std::sync::atomic::AtomicBool::new(false),
3000 })
3001 }
3002
3003 pub async fn read(&self, buf: &mut [u8]) -> std::io::Result<usize> {
3012 if buf.is_empty() {
3013 return Ok(0);
3014 }
3015
3016 if !self.read_backpressured.load(Ordering::Relaxed) {
3017 let res = unsafe {
3018 let r = libc::read(
3019 self.inner.as_raw_fd(),
3020 buf.as_mut_ptr().cast::<libc::c_void>(),
3021 buf.len(),
3022 );
3023 match r.cmp(&0) {
3024 std::cmp::Ordering::Greater => Ok(r as usize),
3025 std::cmp::Ordering::Equal => Ok(0), std::cmp::Ordering::Less => Err(std::io::Error::last_os_error()),
3027 }
3028 };
3029
3030 match res {
3031 Ok(n) => return Ok(n),
3032 Err(e) => {
3033 if e.kind() != std::io::ErrorKind::WouldBlock {
3034 return Err(e);
3035 }
3036 }
3037 }
3038 self.read_backpressured.store(true, Ordering::Relaxed);
3039 }
3040
3041 let future = DtactIoFuture {
3042 worker_idx: self.worker_idx,
3043 fd: self.inner.as_raw_fd() as u32,
3044 direct_fd_idx: self.direct_fd_idx,
3045 op: OpCode::Read,
3046 buf_ptr: buf.as_mut_ptr(),
3047 len: buf.len(),
3048 offset: 0,
3049 addr: None,
3050 addr_len: 0,
3051 slot_idx: None,
3052 msg_ptr: std::ptr::null_mut(),
3053 }
3054 .await;
3055
3056 self.read_backpressured.store(false, Ordering::Relaxed);
3057 future.map(|n| n.min(buf.len()))
3058 }
3059
3060 pub async fn write(&self, buf: &[u8]) -> std::io::Result<usize> {
3069 if buf.is_empty() {
3070 return Ok(0);
3071 }
3072
3073 if !self.write_backpressured.load(Ordering::Relaxed) {
3074 let res = unsafe {
3075 let r = libc::write(
3076 self.inner.as_raw_fd(),
3077 buf.as_ptr().cast::<libc::c_void>(),
3078 buf.len(),
3079 );
3080 if r >= 0 {
3081 Ok(r as usize)
3082 } else {
3083 Err(std::io::Error::last_os_error())
3084 }
3085 };
3086
3087 match res {
3088 Ok(n) => return Ok(n),
3089 Err(e) => {
3090 if e.kind() != std::io::ErrorKind::WouldBlock {
3091 return Err(e);
3092 }
3093 }
3094 }
3095 self.write_backpressured.store(true, Ordering::Relaxed);
3096 }
3097
3098 let future = DtactIoFuture {
3099 worker_idx: self.worker_idx,
3100 fd: self.inner.as_raw_fd() as u32,
3101 direct_fd_idx: self.direct_fd_idx,
3102 op: OpCode::Write,
3103 buf_ptr: buf.as_ptr().cast_mut(),
3104 len: buf.len(),
3105 offset: 0,
3106 addr: None,
3107 addr_len: 0,
3108 slot_idx: None,
3109 msg_ptr: std::ptr::null_mut(),
3110 };
3111
3112 self.write_backpressured.store(false, Ordering::Relaxed);
3113 future.await
3114 }
3115
3116 #[allow(clippy::too_many_lines)]
3134 pub async fn connect(path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
3135 let (libc_addr, addr_len) = unix_path_to_libc(path.as_ref())?;
3136
3137 let fd = unsafe {
3138 libc::socket(
3139 libc::AF_UNIX,
3140 libc::SOCK_STREAM | libc::SOCK_CLOEXEC | libc::SOCK_NONBLOCK,
3141 0,
3142 )
3143 };
3144 if fd < 0 {
3145 return Err(std::io::Error::last_os_error());
3146 }
3147
3148 let stream = unsafe { std::os::unix::net::UnixStream::from_raw_fd(fd) };
3150 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
3151 let worker_idx = fd as usize % num_workers;
3152 let state = &WORKERS.get().unwrap()[worker_idx];
3153 let direct_fd_idx = register_fd_sync(state, fd);
3154
3155 let connect_res = unsafe {
3157 libc::connect(
3158 fd,
3159 (&raw const libc_addr).cast::<libc::sockaddr>(),
3160 addr_len,
3161 )
3162 };
3163 if connect_res == 0 {
3164 return Ok(Self {
3165 inner: stream,
3166 direct_fd_idx,
3167 worker_idx,
3168 read_backpressured: std::sync::atomic::AtomicBool::new(false),
3169 write_backpressured: std::sync::atomic::AtomicBool::new(false),
3170 });
3171 }
3172 let err = std::io::Error::last_os_error();
3173 if err.raw_os_error() != Some(libc::EINPROGRESS) {
3174 return Err(err);
3175 }
3176
3177 let mut pollfd = libc::pollfd {
3182 fd,
3183 events: libc::POLLOUT,
3184 revents: 0,
3185 };
3186 let poll_res = unsafe { libc::poll(&raw mut pollfd, 1, 0) };
3187 if poll_res > 0 {
3188 if (pollfd.revents & libc::POLLOUT) != 0 {
3189 let mut err_code: libc::c_int = 0;
3190 let mut err_len = std::mem::size_of::<libc::c_int>() as libc::socklen_t;
3191 let sockopt_res = unsafe {
3192 libc::getsockopt(
3193 fd,
3194 libc::SOL_SOCKET,
3195 libc::SO_ERROR,
3196 (&raw mut err_code).cast::<libc::c_void>(),
3197 &raw mut err_len,
3198 )
3199 };
3200 if sockopt_res == 0 && err_code == 0 {
3201 return Ok(Self {
3202 inner: stream,
3203 direct_fd_idx,
3204 worker_idx,
3205 read_backpressured: std::sync::atomic::AtomicBool::new(false),
3206 write_backpressured: std::sync::atomic::AtomicBool::new(false),
3207 });
3208 }
3209 let os_err = if err_code != 0 {
3210 err_code
3211 } else {
3212 libc::ECONNREFUSED
3213 };
3214 return Err(std::io::Error::from_raw_os_error(os_err));
3215 } else if (pollfd.revents & (libc::POLLERR | libc::POLLHUP)) != 0 {
3216 return Err(std::io::Error::new(
3217 std::io::ErrorKind::ConnectionRefused,
3218 "connect failed",
3219 ));
3220 }
3221 }
3222
3223 let connect_res = DtactIoFuture {
3224 worker_idx,
3225 fd: fd as u32,
3226 direct_fd_idx,
3227 op: OpCode::Connect,
3228 buf_ptr: std::ptr::null_mut(),
3229 len: 0,
3230 offset: 0,
3231 addr: Some(libc_addr),
3232 addr_len,
3233 slot_idx: None,
3234 msg_ptr: std::ptr::null_mut(),
3235 }
3236 .await;
3237
3238 match connect_res {
3239 Ok(_) => Ok(Self {
3240 inner: stream,
3241 direct_fd_idx,
3242 worker_idx,
3243 read_backpressured: std::sync::atomic::AtomicBool::new(false),
3244 write_backpressured: std::sync::atomic::AtomicBool::new(false),
3245 }),
3246 Err(e) => Err(e),
3247 }
3248 }
3249
3250 pub fn peer_cred(&self) -> std::io::Result<DtactUCred> {
3264 peer_cred_impl(self.inner.as_raw_fd())
3265 }
3266}
3267
3268impl Drop for DtactUnixStream {
3269 fn drop(&mut self) {
3270 if let Some(workers) = WORKERS.get()
3271 && let Some(state) = workers.get(self.worker_idx)
3272 {
3273 unregister_fd_sync(state, self.direct_fd_idx);
3274 }
3275 }
3276}
3277
3278impl crate::io::AsyncRead for DtactUnixStream {
3279 async fn read(&self, buf: &mut [u8]) -> std::io::Result<usize> {
3280 self.read(buf).await
3281 }
3282}
3283
3284impl crate::io::AsyncWrite for DtactUnixStream {
3285 async fn write(&self, buf: &[u8]) -> std::io::Result<usize> {
3286 self.write(buf).await
3287 }
3288}
3289
3290pub struct DtactUnixListener {
3293 inner: std::os::unix::net::UnixListener,
3294 direct_fd_idx: u32,
3295 worker_idx: usize,
3296}
3297
3298impl DtactUnixListener {
3299 pub fn bind(path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
3309 let listener = std::os::unix::net::UnixListener::bind(path)?;
3310 Self::from_std(listener)
3311 }
3312
3313 pub fn from_std(listener: std::os::unix::net::UnixListener) -> std::io::Result<Self> {
3325 let fd = listener.as_raw_fd();
3326 listener.set_nonblocking(true)?;
3327
3328 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
3329 let worker_idx = WORKER_ROUND_ROBIN.fetch_add(1, Ordering::Relaxed) % num_workers;
3330 let state = &WORKERS.get().unwrap()[worker_idx];
3331
3332 let direct_fd_idx = register_fd_sync(state, fd);
3333
3334 Ok(Self {
3335 inner: listener,
3336 direct_fd_idx,
3337 worker_idx,
3338 })
3339 }
3340
3341 pub async fn accept(
3358 &self,
3359 ) -> std::io::Result<(DtactUnixStream, std::os::unix::net::SocketAddr)> {
3360 let res = unsafe {
3362 libc::accept4(
3363 self.inner.as_raw_fd(),
3364 std::ptr::null_mut(),
3365 std::ptr::null_mut(),
3366 libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC,
3367 )
3368 };
3369 if res >= 0 {
3370 let stream = unsafe { std::os::unix::net::UnixStream::from_raw_fd(res) };
3371 let peer_addr = stream.peer_addr()?;
3372 let client_stream = DtactUnixStream::from_std(stream)?;
3373 return Ok((client_stream, peer_addr));
3374 }
3375 let err = std::io::Error::last_os_error();
3376 if err.kind() != std::io::ErrorKind::WouldBlock {
3377 return Err(err);
3378 }
3379
3380 let res = DtactIoFuture {
3382 worker_idx: self.worker_idx,
3383 fd: self.inner.as_raw_fd() as u32,
3384 direct_fd_idx: self.direct_fd_idx,
3385 op: OpCode::Accept,
3386 buf_ptr: std::ptr::null_mut(),
3387 len: 0,
3388 offset: 0,
3389 addr: None,
3390 addr_len: 0,
3391 slot_idx: None,
3392 msg_ptr: std::ptr::null_mut(),
3393 }
3394 .await?;
3395
3396 let client_fd = res as RawFd;
3397 let stream = unsafe { std::os::unix::net::UnixStream::from_raw_fd(client_fd) };
3399 let peer_addr = stream.peer_addr()?;
3400 let client_stream = DtactUnixStream::from_std(stream)?;
3401 Ok((client_stream, peer_addr))
3402 }
3403}
3404
3405impl Drop for DtactUnixListener {
3406 fn drop(&mut self) {
3407 if let Some(workers) = WORKERS.get()
3408 && let Some(state) = workers.get(self.worker_idx)
3409 {
3410 unregister_fd_sync(state, self.direct_fd_idx);
3411 }
3412 }
3413}
3414
3415#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3422pub struct DtactUCred {
3423 uid: u32,
3424 gid: u32,
3425 pid: Option<i32>,
3426}
3427
3428impl DtactUCred {
3429 #[must_use]
3431 pub const fn uid(&self) -> u32 {
3432 self.uid
3433 }
3434
3435 #[must_use]
3437 pub const fn gid(&self) -> u32 {
3438 self.gid
3439 }
3440
3441 #[must_use]
3444 pub const fn pid(&self) -> Option<i32> {
3445 self.pid
3446 }
3447}
3448
3449#[cfg(target_os = "linux")]
3450fn peer_cred_impl(fd: RawFd) -> std::io::Result<DtactUCred> {
3451 let mut cred: libc::ucred = unsafe { std::mem::zeroed() };
3452 let mut len = std::mem::size_of::<libc::ucred>() as libc::socklen_t;
3453 let r = unsafe {
3454 libc::getsockopt(
3455 fd,
3456 libc::SOL_SOCKET,
3457 libc::SO_PEERCRED,
3458 (&raw mut cred).cast::<libc::c_void>(),
3459 &raw mut len,
3460 )
3461 };
3462 if r != 0 {
3463 return Err(std::io::Error::last_os_error());
3464 }
3465 Ok(DtactUCred {
3466 uid: cred.uid,
3467 gid: cred.gid,
3468 pid: Some(cred.pid),
3469 })
3470}
3471
3472#[cfg(not(target_os = "linux"))]
3473fn peer_cred_impl(fd: RawFd) -> std::io::Result<DtactUCred> {
3474 let mut uid: libc::uid_t = 0;
3475 let mut gid: libc::gid_t = 0;
3476 let r = unsafe { libc::getpeereid(fd, &raw mut uid, &raw mut gid) };
3477 if r != 0 {
3478 return Err(std::io::Error::last_os_error());
3479 }
3480 Ok(DtactUCred {
3481 uid,
3482 gid,
3483 pid: None,
3484 })
3485}
3486
3487fn unix_path_to_libc(
3496 path: &std::path::Path,
3497) -> std::io::Result<(libc::sockaddr_storage, libc::socklen_t)> {
3498 use std::os::unix::ffi::OsStrExt;
3499 let bytes = path.as_os_str().as_bytes();
3500 let mut storage: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
3501 let sun_ptr = (&raw mut storage).cast::<libc::sockaddr_un>();
3507 let sun_path_cap = unsafe { (*sun_ptr).sun_path.len() };
3508 if bytes.len() >= sun_path_cap {
3512 return Err(std::io::Error::new(
3513 std::io::ErrorKind::InvalidInput,
3514 "unix socket path too long for sockaddr_un::sun_path",
3515 ));
3516 }
3517 unsafe {
3518 (*sun_ptr).sun_family = libc::AF_UNIX as libc::sa_family_t;
3519 std::ptr::copy_nonoverlapping(
3520 bytes.as_ptr(),
3521 (*sun_ptr).sun_path.as_mut_ptr().cast::<u8>(),
3522 bytes.len(),
3523 );
3524 }
3525 let len = (std::mem::size_of::<libc::sa_family_t>() + bytes.len() + 1) as libc::socklen_t;
3526 Ok((storage, len))
3527}
3528
3529#[derive(Debug, Clone, PartialEq, Eq)]
3543pub struct DtactUnixSocketAddr(Option<std::path::PathBuf>);
3544
3545impl DtactUnixSocketAddr {
3546 #[must_use]
3548 pub fn as_pathname(&self) -> Option<&std::path::Path> {
3549 self.0.as_deref()
3550 }
3551
3552 #[must_use]
3555 pub const fn is_unnamed(&self) -> bool {
3556 self.0.is_none()
3557 }
3558}
3559
3560fn sockaddr_un_to_addr(
3564 storage: &libc::sockaddr_storage,
3565 len: libc::socklen_t,
3566) -> DtactUnixSocketAddr {
3567 let family_len = std::mem::size_of::<libc::sa_family_t>();
3568 if (len as usize) <= family_len {
3569 return DtactUnixSocketAddr(None); }
3571 let sun = unsafe { &*std::ptr::from_ref(storage).cast::<libc::sockaddr_un>() };
3575 let path_len = (len as usize) - family_len;
3576 let path_len = path_len.min(sun.sun_path.len());
3577 let bytes = unsafe { std::slice::from_raw_parts(sun.sun_path.as_ptr().cast::<u8>(), path_len) };
3579 let bytes = bytes.split(|&b| b == 0).next().unwrap_or(&[]);
3583 if bytes.is_empty() {
3584 DtactUnixSocketAddr(None)
3585 } else {
3586 use std::os::unix::ffi::OsStrExt;
3587 DtactUnixSocketAddr(Some(std::path::PathBuf::from(std::ffi::OsStr::from_bytes(
3588 bytes,
3589 ))))
3590 }
3591}
3592
3593pub struct DtactUnixDatagram {
3605 inner: std::os::unix::net::UnixDatagram,
3606 direct_fd_idx: u32,
3607 worker_idx: usize,
3608 read_backpressured: std::sync::atomic::AtomicBool,
3609 write_backpressured: std::sync::atomic::AtomicBool,
3610}
3611
3612impl DtactUnixDatagram {
3613 pub fn bind(path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
3621 let sock = std::os::unix::net::UnixDatagram::bind(path)?;
3622 Self::from_std(sock)
3623 }
3624
3625 pub fn unbound() -> std::io::Result<Self> {
3633 let sock = std::os::unix::net::UnixDatagram::unbound()?;
3634 Self::from_std(sock)
3635 }
3636
3637 pub fn from_std(socket: std::os::unix::net::UnixDatagram) -> std::io::Result<Self> {
3648 let fd = socket.as_raw_fd();
3649 socket.set_nonblocking(true)?;
3650 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
3651 let worker_idx = fd as usize % num_workers;
3652 let state = &WORKERS.get().unwrap()[worker_idx];
3653 let direct_fd_idx = register_fd_sync(state, fd);
3654 Ok(Self {
3655 inner: socket,
3656 direct_fd_idx,
3657 worker_idx,
3658 read_backpressured: std::sync::atomic::AtomicBool::new(false),
3659 write_backpressured: std::sync::atomic::AtomicBool::new(false),
3660 })
3661 }
3662
3663 pub async fn send_to(
3670 &self,
3671 buf: &[u8],
3672 target: impl AsRef<std::path::Path>,
3673 ) -> std::io::Result<usize> {
3674 struct SendToState {
3675 storage: libc::sockaddr_storage,
3676 iov: libc::iovec,
3677 msg: libc::msghdr,
3678 }
3679 unsafe impl Send for SendToState {}
3680
3681 if !self.write_backpressured.load(Ordering::Relaxed) {
3682 let (storage, addr_len) = unix_path_to_libc(target.as_ref())?;
3683 let mut state = SendToState {
3684 storage,
3685 iov: libc::iovec {
3686 iov_base: buf.as_ptr().cast_mut().cast::<libc::c_void>(),
3687 iov_len: buf.len(),
3688 },
3689 msg: unsafe { std::mem::zeroed() },
3690 };
3691 state.msg.msg_name = std::ptr::addr_of_mut!(state.storage).cast::<libc::c_void>();
3692 state.msg.msg_namelen = addr_len;
3693 state.msg.msg_iov = &raw mut state.iov;
3694 state.msg.msg_iovlen = 1;
3695
3696 let r = unsafe { libc::sendmsg(self.inner.as_raw_fd(), &raw const state.msg, 0) };
3697 if r >= 0 {
3698 return Ok(r as usize);
3699 }
3700
3701 let e = std::io::Error::last_os_error();
3702 if e.kind() != std::io::ErrorKind::WouldBlock {
3703 return Err(e);
3704 }
3705 self.write_backpressured.store(true, Ordering::Relaxed);
3706 }
3707
3708 let (storage, addr_len) = unix_path_to_libc(target.as_ref())?;
3710 let mut state = SendToState {
3711 storage,
3712 iov: libc::iovec {
3713 iov_base: buf.as_ptr().cast_mut().cast::<libc::c_void>(),
3714 iov_len: buf.len(),
3715 },
3716 msg: unsafe { std::mem::zeroed() },
3717 };
3718 state.msg.msg_name = std::ptr::addr_of_mut!(state.storage).cast::<libc::c_void>();
3719 state.msg.msg_namelen = addr_len;
3720 state.msg.msg_iov = &raw mut state.iov;
3721 state.msg.msg_iovlen = 1;
3722
3723 let mut fut = DtactIoFuture::new(
3724 self.worker_idx,
3725 self.inner.as_raw_fd() as u32,
3726 self.direct_fd_idx,
3727 OpCode::SendTo,
3728 std::ptr::null_mut(),
3729 0,
3730 0,
3731 None,
3732 0,
3733 None,
3734 );
3735 fut.msg_ptr = &raw mut state.msg;
3736 let res = fut.await;
3737 self.write_backpressured.store(false, Ordering::Relaxed);
3738 res
3739 }
3740
3741 pub async fn recv_from(&self, buf: &mut [u8]) -> std::io::Result<(usize, DtactUnixSocketAddr)> {
3747 struct RecvFromState {
3748 storage: libc::sockaddr_storage,
3749 iov: libc::iovec,
3750 msg: libc::msghdr,
3751 }
3752 unsafe impl Send for RecvFromState {}
3753
3754 if !self.read_backpressured.load(Ordering::Relaxed) {
3755 let mut state = RecvFromState {
3756 storage: unsafe { std::mem::zeroed() },
3757 iov: libc::iovec {
3758 iov_base: buf.as_mut_ptr().cast::<libc::c_void>(),
3759 iov_len: buf.len(),
3760 },
3761 msg: unsafe { std::mem::zeroed() },
3762 };
3763 state.msg.msg_name = std::ptr::addr_of_mut!(state.storage).cast::<libc::c_void>();
3764 state.msg.msg_namelen =
3765 std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
3766 state.msg.msg_iov = &raw mut state.iov;
3767 state.msg.msg_iovlen = 1;
3768
3769 let r = unsafe { libc::recvmsg(self.inner.as_raw_fd(), &raw mut state.msg, 0) };
3770 if r >= 0 {
3771 let from = sockaddr_un_to_addr(&state.storage, state.msg.msg_namelen);
3772 return Ok((r as usize, from));
3773 }
3774 let e = std::io::Error::last_os_error();
3775 if e.kind() != std::io::ErrorKind::WouldBlock {
3776 return Err(e);
3777 }
3778 self.read_backpressured.store(true, Ordering::Relaxed);
3779 }
3780
3781 let mut state = RecvFromState {
3782 storage: unsafe { std::mem::zeroed() },
3783 iov: libc::iovec {
3784 iov_base: buf.as_mut_ptr().cast::<libc::c_void>(),
3785 iov_len: buf.len(),
3786 },
3787 msg: unsafe { std::mem::zeroed() },
3788 };
3789 state.msg.msg_name = std::ptr::addr_of_mut!(state.storage).cast::<libc::c_void>();
3790 state.msg.msg_namelen = std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
3791 state.msg.msg_iov = &raw mut state.iov;
3792 state.msg.msg_iovlen = 1;
3793
3794 let mut fut = DtactIoFuture::new(
3795 self.worker_idx,
3796 self.inner.as_raw_fd() as u32,
3797 self.direct_fd_idx,
3798 OpCode::RecvFrom,
3799 std::ptr::null_mut(),
3800 0,
3801 0,
3802 None,
3803 0,
3804 None,
3805 );
3806 fut.msg_ptr = &raw mut state.msg;
3807 let n = fut.await?;
3808 self.read_backpressured.store(false, Ordering::Relaxed);
3809 let from = sockaddr_un_to_addr(&state.storage, state.msg.msg_namelen);
3810 Ok((n, from))
3811 }
3812
3813 pub async fn connect(&self, target: impl AsRef<std::path::Path>) -> std::io::Result<()> {
3820 self.inner.connect(target)
3821 }
3822
3823 pub async fn send(&self, buf: &[u8]) -> std::io::Result<usize> {
3829 if buf.is_empty() {
3830 return Ok(0);
3831 }
3832 let r = unsafe {
3833 libc::send(
3834 self.inner.as_raw_fd(),
3835 buf.as_ptr().cast::<libc::c_void>(),
3836 buf.len(),
3837 0,
3838 )
3839 };
3840 if r >= 0 {
3841 return Ok(r as usize);
3842 }
3843 let e = std::io::Error::last_os_error();
3844 if e.kind() != std::io::ErrorKind::WouldBlock {
3845 return Err(e);
3846 }
3847 DtactIoFuture {
3848 worker_idx: self.worker_idx,
3849 fd: self.inner.as_raw_fd() as u32,
3850 direct_fd_idx: self.direct_fd_idx,
3851 op: OpCode::Write,
3852 buf_ptr: buf.as_ptr().cast_mut(),
3853 len: buf.len(),
3854 offset: 0,
3855 addr: None,
3856 addr_len: 0,
3857 slot_idx: None,
3858 msg_ptr: std::ptr::null_mut(),
3859 }
3860 .await
3861 }
3862
3863 pub async fn recv(&self, buf: &mut [u8]) -> std::io::Result<usize> {
3869 if buf.is_empty() {
3870 return Ok(0);
3871 }
3872 let r = unsafe {
3873 libc::recv(
3874 self.inner.as_raw_fd(),
3875 buf.as_mut_ptr().cast::<libc::c_void>(),
3876 buf.len(),
3877 0,
3878 )
3879 };
3880 if r >= 0 {
3881 return Ok(r as usize);
3882 }
3883 let e = std::io::Error::last_os_error();
3884 if e.kind() != std::io::ErrorKind::WouldBlock {
3885 return Err(e);
3886 }
3887 DtactIoFuture {
3888 worker_idx: self.worker_idx,
3889 fd: self.inner.as_raw_fd() as u32,
3890 direct_fd_idx: self.direct_fd_idx,
3891 op: OpCode::Read,
3892 buf_ptr: buf.as_mut_ptr(),
3893 len: buf.len(),
3894 offset: 0,
3895 addr: None,
3896 addr_len: 0,
3897 slot_idx: None,
3898 msg_ptr: std::ptr::null_mut(),
3899 }
3900 .await
3901 }
3902}
3903
3904impl Drop for DtactUnixDatagram {
3905 fn drop(&mut self) {
3906 if let Some(workers) = WORKERS.get()
3907 && let Some(state) = workers.get(self.worker_idx)
3908 {
3909 unregister_fd_sync(state, self.direct_fd_idx);
3910 }
3911 }
3912}
3913
3914pub struct DtactFifoReader {
3928 inner: std::fs::File,
3929 direct_fd_idx: u32,
3930 worker_idx: usize,
3931 backpressured: std::sync::atomic::AtomicBool,
3932}
3933
3934unsafe impl Send for DtactFifoReader {}
3935unsafe impl Sync for DtactFifoReader {}
3936
3937impl DtactFifoReader {
3938 pub async fn read(&self, buf: &mut [u8]) -> std::io::Result<usize> {
3943 if buf.is_empty() {
3944 return Ok(0);
3945 }
3946 if !self.backpressured.load(Ordering::Relaxed) {
3947 let res = unsafe {
3948 let r = libc::read(
3949 self.inner.as_raw_fd(),
3950 buf.as_mut_ptr().cast::<libc::c_void>(),
3951 buf.len(),
3952 );
3953 match r.cmp(&0) {
3954 std::cmp::Ordering::Greater => Ok(r as usize),
3955 std::cmp::Ordering::Equal => Ok(0),
3956 std::cmp::Ordering::Less => Err(std::io::Error::last_os_error()),
3957 }
3958 };
3959 match res {
3960 Ok(n) => return Ok(n),
3961 Err(e) => {
3962 if e.kind() != std::io::ErrorKind::WouldBlock {
3963 return Err(e);
3964 }
3965 }
3966 }
3967 self.backpressured.store(true, Ordering::Relaxed);
3968 }
3969 let future = DtactIoFuture {
3970 worker_idx: self.worker_idx,
3971 fd: self.inner.as_raw_fd() as u32,
3972 direct_fd_idx: self.direct_fd_idx,
3973 op: OpCode::Read,
3974 buf_ptr: buf.as_mut_ptr(),
3975 len: buf.len(),
3976 offset: 0,
3977 addr: None,
3978 addr_len: 0,
3979 slot_idx: None,
3980 msg_ptr: std::ptr::null_mut(),
3981 }
3982 .await;
3983 self.backpressured.store(false, Ordering::Relaxed);
3984 future.map(|n| n.min(buf.len()))
3985 }
3986}
3987
3988impl Drop for DtactFifoReader {
3989 fn drop(&mut self) {
3990 if let Some(workers) = WORKERS.get()
3991 && let Some(state) = workers.get(self.worker_idx)
3992 {
3993 unregister_fd_sync(state, self.direct_fd_idx);
3994 }
3995 }
3996}
3997
3998pub struct DtactFifoWriter {
4000 inner: std::fs::File,
4001 direct_fd_idx: u32,
4002 worker_idx: usize,
4003 backpressured: std::sync::atomic::AtomicBool,
4004}
4005
4006unsafe impl Send for DtactFifoWriter {}
4007unsafe impl Sync for DtactFifoWriter {}
4008
4009impl DtactFifoWriter {
4010 pub async fn write(&self, buf: &[u8]) -> std::io::Result<usize> {
4016 if buf.is_empty() {
4017 return Ok(0);
4018 }
4019 if !self.backpressured.load(Ordering::Relaxed) {
4020 let res = unsafe {
4021 let r = libc::write(
4022 self.inner.as_raw_fd(),
4023 buf.as_ptr().cast::<libc::c_void>(),
4024 buf.len(),
4025 );
4026 if r >= 0 {
4027 Ok(r as usize)
4028 } else {
4029 Err(std::io::Error::last_os_error())
4030 }
4031 };
4032 match res {
4033 Ok(n) => return Ok(n),
4034 Err(e) => {
4035 if e.kind() != std::io::ErrorKind::WouldBlock {
4036 return Err(e);
4037 }
4038 }
4039 }
4040 self.backpressured.store(true, Ordering::Relaxed);
4041 }
4042 let future = DtactIoFuture {
4043 worker_idx: self.worker_idx,
4044 fd: self.inner.as_raw_fd() as u32,
4045 direct_fd_idx: self.direct_fd_idx,
4046 op: OpCode::Write,
4047 buf_ptr: buf.as_ptr().cast_mut(),
4048 len: buf.len(),
4049 offset: 0,
4050 addr: None,
4051 addr_len: 0,
4052 slot_idx: None,
4053 msg_ptr: std::ptr::null_mut(),
4054 }
4055 .await;
4056 self.backpressured.store(false, Ordering::Relaxed);
4057 future.map(|n| n.min(buf.len()))
4058 }
4059}
4060
4061impl Drop for DtactFifoWriter {
4062 fn drop(&mut self) {
4063 if let Some(workers) = WORKERS.get()
4064 && let Some(state) = workers.get(self.worker_idx)
4065 {
4066 unregister_fd_sync(state, self.direct_fd_idx);
4067 }
4068 }
4069}
4070
4071fn register_fifo_fd(fd: RawFd) -> (u32, usize) {
4072 let num_workers = GLOBAL_CONFIG.get().map_or(1, |c| c.workers);
4073 let worker_idx = fd as usize % num_workers;
4074 let state = &WORKERS.get().unwrap()[worker_idx];
4075 (register_fd_sync(state, fd), worker_idx)
4076}
4077
4078pub async fn open_fifo_read(
4092 path: impl Into<std::path::PathBuf>,
4093) -> std::io::Result<DtactFifoReader> {
4094 use std::os::unix::fs::OpenOptionsExt;
4095 let path = path.into();
4096 let file = std::fs::OpenOptions::new()
4102 .read(true)
4103 .custom_flags(libc::O_NONBLOCK)
4104 .open(&path)?;
4105 let (direct_fd_idx, worker_idx) = register_fifo_fd(file.as_raw_fd());
4106 Ok(DtactFifoReader {
4107 inner: file,
4108 direct_fd_idx,
4109 worker_idx,
4110 backpressured: AtomicBool::new(false),
4111 })
4112}
4113
4114pub async fn open_fifo_write(
4126 path: impl Into<std::path::PathBuf>,
4127) -> std::io::Result<DtactFifoWriter> {
4128 use std::os::unix::fs::OpenOptionsExt;
4129 let path = path.into();
4130 let file = std::fs::OpenOptions::new()
4133 .write(true)
4134 .custom_flags(libc::O_NONBLOCK)
4135 .open(&path)?;
4136 let (direct_fd_idx, worker_idx) = register_fifo_fd(file.as_raw_fd());
4137 Ok(DtactFifoWriter {
4138 inner: file,
4139 direct_fd_idx,
4140 worker_idx,
4141 backpressured: AtomicBool::new(false),
4142 })
4143}
4144
4145const fn register_fd_sync(_state: &WorkerState, _fd: RawFd) -> u32 {
4160 u32::MAX
4161}
4162
4163const fn unregister_fd_sync(_state: &WorkerState, _direct_fd_idx: u32) {}
4165
4166const fn socket_addr_to_libc(
4170 addr: std::net::SocketAddr,
4171) -> (libc::sockaddr_storage, libc::socklen_t) {
4172 let mut storage: libc::sockaddr_storage = unsafe { std::mem::zeroed() };
4173 let len = match addr {
4174 std::net::SocketAddr::V4(a) => {
4175 let sin = libc::sockaddr_in {
4176 sin_family: libc::AF_INET as libc::sa_family_t,
4177 sin_port: a.port().to_be(),
4178 sin_addr: libc::in_addr {
4179 s_addr: u32::from_ne_bytes(a.ip().octets()),
4180 },
4181 sin_zero: [0; 8],
4182 };
4183 unsafe {
4184 std::ptr::copy_nonoverlapping(
4185 (&raw const sin).cast::<u8>(),
4186 (&raw mut storage).cast::<u8>(),
4187 std::mem::size_of::<libc::sockaddr_in>(),
4188 );
4189 }
4190 std::mem::size_of::<libc::sockaddr_in>() as libc::socklen_t
4191 }
4192 std::net::SocketAddr::V6(a) => {
4193 let sin6 = libc::sockaddr_in6 {
4194 sin6_family: libc::AF_INET6 as libc::sa_family_t,
4195 sin6_port: a.port().to_be(),
4196 sin6_flowinfo: a.flowinfo(),
4197 sin6_addr: libc::in6_addr {
4198 s6_addr: a.ip().octets(),
4199 },
4200 sin6_scope_id: a.scope_id(),
4201 };
4202 unsafe {
4203 std::ptr::copy_nonoverlapping(
4204 (&raw const sin6).cast::<u8>(),
4205 (&raw mut storage).cast::<u8>(),
4206 std::mem::size_of::<libc::sockaddr_in6>(),
4207 );
4208 }
4209 std::mem::size_of::<libc::sockaddr_in6>() as libc::socklen_t
4210 }
4211 };
4212 (storage, len)
4213}
4214
4215fn sockaddr_storage_to_socketaddr(
4218 storage: &libc::sockaddr_storage,
4219 _len: libc::socklen_t,
4220) -> std::net::SocketAddr {
4221 match libc::c_int::from(storage.ss_family) {
4222 libc::AF_INET => {
4223 let sin = unsafe { &*std::ptr::from_ref(storage).cast::<libc::sockaddr_in>() };
4225 let ip = std::net::Ipv4Addr::from(u32::from_be(sin.sin_addr.s_addr));
4226 let port = u16::from_be(sin.sin_port);
4227 std::net::SocketAddr::V4(std::net::SocketAddrV4::new(ip, port))
4228 }
4229 libc::AF_INET6 => {
4230 let sin6 = unsafe { &*std::ptr::from_ref(storage).cast::<libc::sockaddr_in6>() };
4232 let ip = std::net::Ipv6Addr::from(sin6.sin6_addr.s6_addr);
4233 let port = u16::from_be(sin6.sin6_port);
4234 std::net::SocketAddr::V6(std::net::SocketAddrV6::new(
4235 ip,
4236 port,
4237 sin6.sin6_flowinfo,
4238 sin6.sin6_scope_id,
4239 ))
4240 }
4241 _ => {
4242 panic!("Unsupported address family: {}", storage.ss_family);
4243 }
4244 }
4245}