1#![cfg_attr(not(feature = "udp"), allow(dead_code, unused_imports, unused_mut))]
6
7use crate::handler::{Acceptor, CloseReason, ConnInfo, Connection, Endpoint, Message, UdpInfo, UdpService};
8use crate::offload::{OffloadPool, Rejected};
9use crate::pool::BufPool;
10use crate::runtime::{Config, Shared};
11use crate::timer::{TimerId, TimerWheel};
12use simoom_sys::driver::{Driver, Event, Id, OutBuf, Waker};
13use simoom_sys::net::Socket;
14use simoom_sys::udp::Transmit;
15use std::io;
16use std::net::SocketAddr;
17use std::sync::atomic::{AtomicBool, Ordering};
18use std::sync::Arc;
19use std::time::Duration;
20
21const LISTENER_BIT: u64 = 1 << 63;
22const INITIAL_RECV_BUFFERS: usize = 8;
24const COMPACT_INPUT: usize = 512;
27const BALANCE_AFTER_GROWTH_NS: u64 = 1_000_000_000;
29#[cfg_attr(not(feature = "idle-trim"), allow(dead_code))]
33const IDLE_TRIM_NS: u64 = 1_000_000_000;
34#[cfg_attr(not(feature = "idle-trim"), allow(dead_code))]
35const IDLE_POOL_KEEP: usize = 4;
36const FLUSH_EARLY_BATCH: usize = 16;
39#[cfg_attr(not(feature = "idle-trim"), allow(dead_code))]
41const BURST_EVENTS: u64 = 4096;
42const UDP_BIT: u64 = 1 << 62;
43
44pub(crate) type Task = Box<dyn FnOnce(&mut LoopCtx<'_>) + Send>;
46
47pub(crate) enum Item {
50 Task(Task),
51 Conn(u64, Message),
52 Udp(usize, Message),
53}
54
55pub(crate) struct Queue {
59 items: simoom_sys::mpsc::Mpsc<Item>,
60 pub(crate) waker: Waker,
61 closed: AtomicBool,
62}
63
64impl Queue {
65 pub(crate) fn new(waker: Waker) -> Queue {
66 Queue { items: simoom_sys::mpsc::Mpsc::new(), waker, closed: AtomicBool::new(false) }
67 }
68
69 fn push(&self, t: Item) -> bool {
70 if self.closed.load(Ordering::Acquire) {
71 return false;
72 }
73 self.items.push(t);
74 self.waker.post();
75 !self.closed.load(Ordering::Acquire)
77 }
78
79 fn take(&self, into: &mut Vec<Item>) {
80 self.items.take(into);
81 }
82
83 fn close(&self) {
84 self.closed.store(true, Ordering::Release);
85 let mut v = Vec::new();
86 self.items.take(&mut v);
87 }
88}
89
90#[derive(Clone)]
92pub struct LoopHandle {
93 pub(crate) index: usize,
94 pub(crate) queue: Arc<Queue>,
95}
96
97impl std::fmt::Debug for LoopHandle {
98 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99 f.debug_struct("LoopHandle").field("index", &self.index).finish()
100 }
101}
102
103impl LoopHandle {
104 pub fn index(&self) -> usize {
106 self.index
107 }
108
109 pub fn run(&self, f: impl FnOnce(&mut LoopCtx<'_>) + Send + 'static) -> bool {
111 self.queue.push(Item::Task(Box::new(f)))
112 }
113}
114
115#[derive(Clone)]
118pub struct ConnHandle {
119 queue: Arc<Queue>,
120 key: u64,
121 loop_index: usize,
122}
123
124impl std::fmt::Debug for ConnHandle {
125 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126 f.debug_struct("ConnHandle").field("loop", &self.loop_index).field("key", &self.key).finish()
127 }
128}
129
130impl ConnHandle {
131 pub fn send(&self, msg: Message) -> bool {
134 let key = self.key;
135 self.queue.push(Item::Conn(key, msg))
136 }
137
138 pub fn loop_index(&self) -> usize {
140 self.loop_index
141 }
142}
143
144#[derive(Clone)]
146pub struct EndpointHandle {
147 queue: Arc<Queue>,
148 idx: usize,
149}
150
151impl std::fmt::Debug for EndpointHandle {
152 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153 f.debug_struct("EndpointHandle").field("idx", &self.idx).finish()
154 }
155}
156
157impl EndpointHandle {
158 pub fn send(&self, msg: Message) -> bool {
161 let idx = self.idx;
162 self.queue.push(Item::Udp(idx, msg))
163 }
164}
165
166struct Slab<T> {
171 slots: Vec<(u32, Option<T>)>,
172 free: Vec<u32>,
173 len: usize,
174}
175
176impl<T> Slab<T> {
177 fn new() -> Self {
178 Slab { slots: Vec::new(), free: Vec::new(), len: 0 }
179 }
180 fn next_key(&self) -> u64 {
181 match self.free.last() {
182 Some(&i) => i as u64 | (self.slots[i as usize].0 as u64) << 32,
183 None => self.slots.len() as u64,
184 }
185 }
186 fn insert(&mut self, v: T) -> u64 {
187 self.len += 1;
188 match self.free.pop() {
189 Some(i) => {
190 self.slots[i as usize].1 = Some(v);
191 i as u64 | (self.slots[i as usize].0 as u64) << 32
192 }
193 None => {
194 self.slots.push((0, Some(v)));
195 self.slots.len() as u64 - 1
196 }
197 }
198 }
199 fn get_mut(&mut self, key: u64) -> Option<&mut T> {
200 match self.slots.get_mut((key & 0xFFFF_FFFF) as usize) {
201 Some((g, Some(v))) if *g == (key >> 32) as u32 => Some(v),
202 _ => None,
203 }
204 }
205 fn remove(&mut self, key: u64) -> Option<T> {
206 let s = self.slots.get_mut((key & 0xFFFF_FFFF) as usize)?;
207 if s.0 != (key >> 32) as u32 || s.1.is_none() {
208 return None;
209 }
210 s.0 = s.0.wrapping_add(1) & 0x3FFF_FFFF; self.free.push((key & 0xFFFF_FFFF) as u32);
212 self.len -= 1;
213 s.1.take()
214 }
215 fn keys(&self) -> Vec<u64> {
216 self.slots.iter().enumerate().filter(|(_, s)| s.1.is_some()).map(|(i, s)| i as u64 | (s.0 as u64) << 32).collect()
217 }
218}
219
220#[derive(Clone, Copy)]
225pub(crate) enum TimerTarget {
226 Conn(u64, u64),
227 Udp(usize, u64),
228 Linger(u64),
230}
231
232struct Conn {
233 id: Id,
234 handler: Option<Box<dyn Connection>>,
235 input: Option<Vec<u8>>,
236 queued: usize,
237 reason: Option<CloseReason>,
238 user_paused: bool,
239 hw_paused: bool,
240 want_drain: bool,
241 closing: bool,
242 closed: bool,
243 timers: ConnTimers,
246 lingering: bool,
249 linger_left: u32,
250}
251
252#[cold]
255#[inline(never)]
256fn no_connection() -> ! {
257 panic!("connection exists during its callback")
258}
259
260#[derive(Default)]
264struct ConnTimers {
265 one: Option<TimerId>,
266 more: Option<Box<Vec<TimerId>>>,
267}
268
269impl ConnTimers {
270 fn push(&mut self, id: TimerId, wheel: &TimerWheel<TimerTarget>) {
271 if self.one.is_none_or(|t| !wheel.is_pending(t)) {
272 self.one = Some(id);
273 return;
274 }
275 let v = self.more.get_or_insert_with(Default::default);
276 if v.len() >= 4 {
277 v.retain(|t| wheel.is_pending(*t));
278 }
279 v.push(id);
280 }
281
282 fn remove(&mut self, id: TimerId) {
283 if self.one == Some(id) {
284 self.one = None;
285 } else if let Some(v) = self.more.as_mut() {
286 if let Some(i) = v.iter().position(|t| *t == id) {
287 v.swap_remove(i);
288 }
289 }
290 }
291
292 fn cancel_all(&mut self, wheel: &mut TimerWheel<TimerTarget>) {
293 if let Some(t) = self.one.take() {
294 wheel.cancel(t);
295 }
296 for t in self.more.take().into_iter().flat_map(|v| *v) {
297 wheel.cancel(t);
298 }
299 }
300}
301
302pub(crate) struct ListenerSlot {
303 pub(crate) acceptor: Arc<dyn Acceptor>,
304 pub(crate) socket: Option<Socket>,
306 handoff: bool,
309 id: Option<Id>,
310 pub(crate) bind: Option<(SocketAddr, simoom_sys::net::ListenOptions)>,
312}
313
314impl ListenerSlot {
315 pub(crate) fn new(acceptor: Arc<dyn Acceptor>, socket: Option<Socket>, handoff: bool) -> ListenerSlot {
316 ListenerSlot { acceptor, socket, handoff, id: None, bind: None }
317 }
318}
319
320pub(crate) struct UdpSlot {
321 pub(crate) service: Arc<dyn UdpService>,
322 pub(crate) binding: usize,
323 pub(crate) socket: Option<Socket>,
324 gro: bool,
326 id: Option<Id>,
327 local: Option<SocketAddr>,
328 pub(crate) bind: Option<(SocketAddr, simoom_sys::net::UdpOptions)>,
330}
331
332#[cfg_attr(not(feature = "udp"), allow(dead_code))]
334pub(crate) const GRO_BUFFER: usize = 64 * 1024;
335
336impl UdpSlot {
337 #[cfg_attr(not(feature = "udp"), allow(dead_code))]
338 pub(crate) fn new(service: Arc<dyn UdpService>, binding: usize, socket: Option<Socket>, gro: bool) -> UdpSlot {
339 UdpSlot { service, binding, socket, gro, id: None, local: None, bind: None }
340 }
341}
342
343pub(crate) struct Core {
345 pub(crate) index: usize,
346 driver: Driver,
347 pool: BufPool,
348 timers: TimerWheel<TimerTarget>,
349 conns: Slab<Conn>,
350 now_ns: u64,
351 tick_ns: u64,
352 queue: Arc<Queue>,
353 peers: Vec<LoopHandle>,
354 offload: OffloadPool,
355 cfg: Arc<Config>,
356 shared: Arc<Shared>,
357 listeners: Vec<ListenerSlot>,
358 udps: Vec<UdpSlot>,
359 endpoints: Vec<Option<Box<dyn Endpoint>>>,
360 redeliver: Vec<u64>,
361 drains: Vec<u64>,
362 rr: usize,
363 want_grow: bool,
365 grow_next: bool,
367 grown: bool,
368 parked: Vec<(Socket, Option<SocketAddr>, usize)>,
370 handoff_until: u64,
373 accepted: u64,
374 rejected: u64,
375 closed: u64,
376 #[cfg_attr(not(feature = "idle-trim"), allow(dead_code))]
378 last_io: u64,
379 #[cfg_attr(not(feature = "idle-trim"), allow(dead_code))]
381 io_since_trim: u64,
382 batch_left: usize,
384}
385
386impl Core {
387 fn conn(&mut self, key: u64) -> &mut Conn {
388 match self.conns.get_mut(key) {
389 Some(c) => c,
390 None => no_connection(),
391 }
392 }
393
394 fn deadline_tick(&self, after: Duration) -> u64 {
395 let at = self.now_ns.saturating_add(after.as_nanos().min(u64::MAX as u128) as u64);
396 at.div_ceil(self.tick_ns)
397 }
398
399 fn with_conn(&mut self, key: u64, f: impl FnOnce(&mut dyn Connection, &mut Ctx<'_>)) {
401 let mut f = Some(f);
402 self.with_conn_dyn(key, &mut |h, cx| {
403 if let Some(f) = f.take() {
404 f(h, cx)
405 }
406 });
407 }
408
409 #[inline(never)]
412 fn with_conn_dyn(&mut self, key: u64, f: &mut dyn FnMut(&mut dyn Connection, &mut Ctx<'_>)) {
413 let Some(c) = self.conns.get_mut(key) else { return };
414 if c.closed {
415 return;
416 }
417 let Some(mut h) = c.handler.take() else { return };
418 {
419 let mut cx = Ctx { core: self, key };
420 f(&mut *h, &mut cx);
421 }
422 if let Some(c) = self.conns.get_mut(key) {
423 c.handler = Some(h);
424 }
425 }
426
427 fn with_udp(&mut self, idx: usize, f: impl FnOnce(&mut dyn Endpoint, &mut UdpCtx<'_>)) {
428 if !cfg!(feature = "udp") {
430 return;
431 }
432 let Some(slot) = self.endpoints.get_mut(idx) else { return };
433 let Some(mut ep) = slot.take() else { return };
434 {
435 let mut cx = UdpCtx { core: self, idx };
436 f(&mut *ep, &mut cx);
437 }
438 self.endpoints[idx] = Some(ep);
439 }
440
441 #[inline(never)]
442 pub(crate) fn deliver(&mut self, key: u64, msg: Message) {
443 self.with_conn(key, |h, cx| h.on_message(cx, msg));
444 self.flush_conn(key);
445 }
446
447 pub(crate) fn deliver_udp(&mut self, idx: usize, msg: Message) {
448 self.with_udp(idx, |h, cx| h.on_message(cx, msg));
449 }
450
451 fn ensure_endpoint(&mut self, idx: usize) {
453 let Some(slot) = self.udps.get(idx) else { return };
454 if slot.id.is_none() || self.endpoints[idx].is_some() {
455 return;
456 }
457 let info = UdpInfo { binding: slot.binding, loop_index: self.index, local_addr: slot.local.unwrap_or(SocketAddr::from(([0, 0, 0, 0], 0))) };
458 self.endpoints[idx] = Some(slot.service.endpoint(&info));
459 }
460
461 fn close_now(&mut self, key: u64, reason: CloseReason) {
462 let Some(c) = self.conns.get_mut(key) else { return };
463 if c.closed {
464 return;
465 }
466 c.closed = true;
467 if c.reason.is_none() {
468 c.reason = Some(reason);
469 }
470 let id = c.id;
471 self.driver.close(id);
472 }
473
474 fn linger(&mut self, key: u64) {
481 let linger = self.cfg.linger;
482 let id = match self.conns.get_mut(key) {
483 Some(c) if !c.closed && !c.lingering => c.id,
484 _ => return,
485 };
486 let shut = !linger.is_zero() && self.driver.socket(id).is_some_and(|s| s.shutdown(std::net::Shutdown::Write).is_ok());
487 if !shut {
488 self.close_now(key, CloseReason::Closed);
489 return;
490 }
491 let Some(c) = self.conns.get_mut(key) else { return };
492 c.lingering = true;
493 c.linger_left = self.cfg.linger_bytes.min(u32::MAX as usize) as u32;
494 if let Some(i) = c.input.take() {
495 self.pool.put(i);
496 }
497 let at = self.deadline_tick(linger);
498 let t = self.timers.insert(at, TimerTarget::Linger(key));
499 if let Some(c) = self.conns.get_mut(key) {
500 c.timers.push(t, &self.timers);
501 }
502 self.driver.recv(id, true);
503 }
504
505 pub(crate) fn adopt(&mut self, socket: Socket, peer: Option<SocketAddr>, listener: usize) {
507 let total = self.shared.active.fetch_add(1, Ordering::AcqRel) + 1;
508 let here = self.conns.len + 1;
509 if total > self.cfg.max_connections || here > self.cfg.max_connections_per_loop {
510 self.shared.active.fetch_sub(1, Ordering::AcqRel);
512 self.rejected += 1;
513 drop(socket);
514 return;
515 }
516 let info = ConnInfo { listener, peer, loop_index: self.index };
517 let Some(handler) = self.listeners.get(listener).and_then(|l| l.acceptor.accept(&info)) else {
518 self.shared.active.fetch_sub(1, Ordering::AcqRel);
519 self.rejected += 1;
520 return;
521 };
522 if self.register(socket, handler).is_err() {
523 self.shared.active.fetch_sub(1, Ordering::AcqRel);
524 self.rejected += 1;
525 return;
526 }
527 if self.accepted == 0 {
528 self.driver.set_buffer_target(self.cfg.recv_buffers);
529 }
530 if self.cfg.lazy_loops && self.index == 0 && self.conns.len >= 3 {
534 self.want_grow |= !self.grown;
535 }
536 self.accepted += 1;
537 }
538
539 fn register(&mut self, socket: Socket, handler: Box<dyn Connection>) -> io::Result<u64> {
540 let key = self.conns.next_key();
541 let id = self.driver.stream(socket, key)?;
542 let k = self.conns.insert(Conn {
543 id,
544 handler: Some(handler),
545 input: None,
546 queued: 0,
547 reason: None,
548 user_paused: false,
549 hw_paused: false,
550 want_drain: false,
551 closing: false,
552 closed: false,
553 timers: ConnTimers::default(),
554 lingering: false,
555 linger_left: 0,
556 });
557 debug_assert_eq!(k, key);
558 self.driver.recv(id, true);
559 self.with_conn(key, |h, cx| h.on_open(cx));
560 Ok(key)
561 }
562
563}
564
565pub struct Ctx<'a> {
571 core: &'a mut Core,
572 key: u64,
573}
574
575impl std::fmt::Debug for Ctx<'_> {
576 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
577 f.debug_struct("Ctx").field("loop", &self.core.index).field("key", &self.key).finish()
578 }
579}
580
581impl Ctx<'_> {
582 pub fn send(&mut self, buf: Vec<u8>) {
585 self.send_buf(OutBuf::Owned(buf));
586 }
587
588 pub fn send_static(&mut self, bytes: &'static [u8]) {
590 self.send_buf(OutBuf::Static(bytes));
591 }
592
593 pub fn send_buf(&mut self, buf: OutBuf) {
595 let high = self.core.cfg.write_high_water;
596 let c = self.core.conn(self.key);
597 if c.closed || buf.is_empty() {
598 if let OutBuf::Owned(v) = buf {
599 self.core.pool.put(v);
600 }
601 return;
602 }
603 c.queued += buf.len();
604 let id = c.id;
605 let pause = c.queued > high && !c.hw_paused;
606 if pause {
607 c.hw_paused = true;
608 }
609 let user_paused = c.user_paused;
610 self.core.driver.send(id, buf);
611 if pause && !user_paused {
612 self.core.driver.recv(id, false);
613 }
614 }
615
616 pub fn buffer(&mut self, min: usize) -> Vec<u8> {
619 self.core.pool.get(min)
620 }
621
622 pub fn recycle(&mut self, buf: Vec<u8>) {
624 self.core.pool.put(buf);
625 }
626
627 pub fn queued(&mut self) -> usize {
629 self.core.conn(self.key).queued
630 }
631
632 pub fn pause_reading(&mut self) {
635 let c = self.core.conn(self.key);
636 if c.user_paused || c.closed {
637 return;
638 }
639 c.user_paused = true;
640 let (id, hw) = (c.id, c.hw_paused);
641 if !hw {
642 self.core.driver.recv(id, false);
643 }
644 }
645
646 pub fn resume_reading(&mut self) {
649 let c = self.core.conn(self.key);
650 if !c.user_paused || c.closed {
651 return;
652 }
653 c.user_paused = false;
654 let (id, hw, has_input) = (c.id, c.hw_paused, c.input.as_ref().is_some_and(|i| !i.is_empty()));
655 if !hw {
656 self.core.driver.recv(id, true);
657 }
658 if has_input {
659 self.core.redeliver.push(self.key);
660 }
661 }
662
663 pub fn redeliver(&mut self) {
667 let c = self.core.conn(self.key);
668 if c.user_paused || c.closed || !c.input.as_ref().is_some_and(|i| !i.is_empty()) {
669 return;
670 }
671 self.core.redeliver.push(self.key);
672 }
673
674 pub fn notify_drain(&mut self) {
677 let low = self.core.cfg.write_low_water;
678 let c = self.core.conn(self.key);
679 if c.queued <= low {
680 self.core.drains.push(self.key);
681 } else {
682 c.want_drain = true;
683 }
684 }
685
686 pub fn close(&mut self) {
689 let c = self.core.conn(self.key);
690 if c.closed || c.closing {
691 return;
692 }
693 c.closing = true;
694 if c.queued == 0 {
695 self.core.linger(self.key);
696 }
697 }
698
699 pub fn abort(&mut self) {
701 self.core.close_now(self.key, CloseReason::Aborted);
702 }
703
704 pub fn is_closing(&mut self) -> bool {
706 let c = self.core.conn(self.key);
707 c.closing || c.closed
708 }
709
710 pub fn set_timer(&mut self, after: Duration, token: u64) -> TimerId {
713 let at = self.core.deadline_tick(after);
714 let id = self.core.timers.insert(at, TimerTarget::Conn(self.key, token));
715 let core = &mut *self.core;
716 if let Some(c) = core.conns.get_mut(self.key) {
717 c.timers.push(id, &core.timers);
718 }
719 id
720 }
721
722 pub fn cancel_timer(&mut self, id: TimerId) -> bool {
724 if let Some(c) = self.core.conns.get_mut(self.key) {
725 c.timers.remove(id);
726 }
727 self.core.timers.cancel(id).is_some()
728 }
729
730 pub fn handle(&mut self) -> ConnHandle {
732 ConnHandle { queue: self.core.queue.clone(), key: self.key, loop_index: self.core.index }
733 }
734
735 pub fn offload<F>(&mut self, job: F) -> Result<(), Rejected>
738 where
739 F: FnOnce() -> Message + Send + 'static,
740 {
741 let h = self.handle();
742 self.core.offload.try_submit(Box::new(move || {
743 h.send(job());
744 }))
745 }
746
747 pub fn offload_pool(&mut self) -> &OffloadPool {
749 &self.core.offload
750 }
751
752 pub fn now_ns(&mut self) -> u64 {
754 self.core.now_ns
755 }
756
757 pub fn peer_addr(&mut self) -> io::Result<SocketAddr> {
759 let id = self.core.conn(self.key).id;
760 self.core.driver.socket(id).ok_or_else(|| io::Error::new(io::ErrorKind::NotConnected, "closed"))?.peer_addr()
761 }
762
763 pub fn local_addr(&mut self) -> io::Result<SocketAddr> {
765 let id = self.core.conn(self.key).id;
766 self.core.driver.socket(id).ok_or_else(|| io::Error::new(io::ErrorKind::NotConnected, "closed"))?.local_addr()
767 }
768
769 pub fn socket(&mut self) -> Option<&Socket> {
771 let id = self.core.conn(self.key).id;
772 self.core.driver.socket(id)
773 }
774
775 pub fn loop_index(&mut self) -> usize {
777 self.core.index
778 }
779
780 #[inline(never)]
786 pub fn connection_pressure(&mut self) -> bool {
787 let max = self.core.cfg.max_connections;
788 let budget = simoom_sys::net::descriptor_limit().map_or(max, |fds| max.min(fds.saturating_sub(32).max(1)));
789 let open = self.core.shared.active.load(Ordering::Relaxed);
790 open.saturating_mul(4) > budget.saturating_mul(3)
791 }
792}
793
794pub struct UdpCtx<'a> {
796 core: &'a mut Core,
797 idx: usize,
798}
799
800impl std::fmt::Debug for UdpCtx<'_> {
801 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
802 f.debug_struct("UdpCtx").field("loop", &self.core.index).field("idx", &self.idx).finish()
803 }
804}
805
806impl UdpCtx<'_> {
807 pub fn send(&mut self, msgs: &[Transmit<'_>]) -> io::Result<usize> {
810 let id = self.core.udps[self.idx].id.ok_or_else(|| io::Error::new(io::ErrorKind::NotConnected, "closed"))?;
811 self.core.driver.send_to(id, msgs)
812 }
813
814 pub fn local_addr(&mut self) -> Option<SocketAddr> {
816 self.core.udps[self.idx].local
817 }
818
819 pub fn set_timer(&mut self, after: Duration, token: u64) -> TimerId {
821 let at = self.core.deadline_tick(after);
822 self.core.timers.insert(at, TimerTarget::Udp(self.idx, token))
823 }
824
825 pub fn cancel_timer(&mut self, id: TimerId) -> bool {
827 self.core.timers.cancel(id).is_some()
828 }
829
830 pub fn handle(&mut self) -> EndpointHandle {
832 EndpointHandle { queue: self.core.queue.clone(), idx: self.idx }
833 }
834
835 pub fn buffer(&mut self, min: usize) -> Vec<u8> {
837 self.core.pool.get(min)
838 }
839
840 pub fn recycle(&mut self, buf: Vec<u8>) {
842 self.core.pool.put(buf);
843 }
844
845 pub fn offload_pool(&mut self) -> &OffloadPool {
847 &self.core.offload
848 }
849
850 pub fn now_ns(&mut self) -> u64 {
852 self.core.now_ns
853 }
854
855 pub fn loop_index(&mut self) -> usize {
857 self.core.index
858 }
859}
860
861pub struct LoopCtx<'a> {
863 pub(crate) core: &'a mut Core,
864}
865
866impl std::fmt::Debug for LoopCtx<'_> {
867 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
868 f.debug_struct("LoopCtx").field("loop", &self.core.index).finish()
869 }
870}
871
872impl LoopCtx<'_> {
873 pub fn index(&self) -> usize {
875 self.core.index
876 }
877
878 pub fn adopt(&mut self, socket: Socket, handler: Box<dyn Connection>) -> io::Result<ConnHandle> {
881 self.core.shared.active.fetch_add(1, Ordering::AcqRel);
882 match self.core.register(socket, handler) {
883 Ok(key) => Ok(ConnHandle { queue: self.core.queue.clone(), key, loop_index: self.core.index }),
884 Err(e) => {
885 self.core.shared.active.fetch_sub(1, Ordering::AcqRel);
886 Err(e)
887 }
888 }
889 }
890
891 pub fn connections(&self) -> usize {
893 self.core.conns.len
894 }
895
896 pub fn now_ns(&self) -> u64 {
898 self.core.now_ns
899 }
900
901 pub fn driver_stats(&self) -> simoom_sys::driver::DriverStats {
903 self.core.driver.stats()
904 }
905
906 pub fn pool_stats(&self) -> crate::pool::PoolStats {
908 self.core.pool.stats()
909 }
910}
911
912pub(crate) struct LoopSetup {
917 pub(crate) index: usize,
918 pub(crate) cfg: Arc<Config>,
919 pub(crate) shared: Arc<Shared>,
920 pub(crate) offload: OffloadPool,
921 pub(crate) listeners: Vec<ListenerSlot>,
922 pub(crate) udps: Vec<UdpSlot>,
923}
924
925pub(crate) fn build(mut setup: LoopSetup) -> io::Result<(Core, Arc<Queue>)> {
927 let cfg = setup.cfg;
928 for l in setup.listeners.iter_mut() {
929 if let Some((addr, opts)) = l.bind.take() {
930 l.socket = Some(Socket::tcp_listen(addr, &opts)?);
931 }
932 }
933 #[cfg(feature = "udp")]
934 for u in setup.udps.iter_mut() {
935 if let Some((addr, opts)) = u.bind.take() {
936 u.socket = Some(Socket::udp_bind(addr, &opts)?);
937 }
938 }
939 let mut driver = Driver::new(&cfg.driver)?;
940 driver.set_buffer_target(if setup.udps.is_empty() { cfg.recv_buffers.min(INITIAL_RECV_BUFFERS) } else { cfg.recv_buffers });
942 let queue = Arc::new(Queue::new(driver.waker()));
943 let tick_ns = (cfg.timer_tick.as_nanos() as u64).max(1);
944 let now_ns = simoom_sys::time::now_ns();
945 let core = Core {
946 index: setup.index,
947 driver,
948 pool: BufPool::new(cfg.buffer_classes, cfg.pool_cache),
949 timers: TimerWheel::new(now_ns / tick_ns),
950 conns: Slab::new(),
951 now_ns,
952 tick_ns,
953 queue: queue.clone(),
954 peers: Vec::new(),
955 offload: setup.offload,
956 shared: setup.shared,
957 listeners: setup.listeners,
958 endpoints: if cfg!(feature = "udp") { setup.udps.iter().map(|_| None).collect() } else { Vec::new() },
959 udps: setup.udps,
960 redeliver: Vec::new(),
961 drains: Vec::new(),
962 rr: setup.index,
963 want_grow: false,
964 grow_next: false,
965 grown: false,
966 parked: Vec::new(),
967 handoff_until: 0,
968 accepted: 0,
969 rejected: 0,
970 closed: 0,
971 last_io: now_ns,
972 io_since_trim: 0,
973 batch_left: 0,
974 cfg,
975 };
976 Ok((core, queue))
977}
978
979impl Core {
980 pub(crate) fn backend(&self) -> (simoom_sys::driver::Backend, Option<String>) {
981 (self.driver.backend(), self.driver.fallback_reason().map(str::to_string))
982 }
983
984 fn start_io(&mut self) -> io::Result<()> {
985 for i in 0..self.listeners.len() {
986 if let Some(s) = self.listeners[i].socket.take() {
987 self.listeners[i].id = Some(self.driver.listen(s, LISTENER_BIT | i as u64)?);
988 }
989 }
990 #[cfg(feature = "udp")]
991 for i in 0..self.udps.len() {
992 if let Some(s) = self.udps[i].socket.take() {
993 let local = s.local_addr().ok();
994 let key = UDP_BIT | i as u64;
995 let id = if self.udps[i].gro { self.driver.udp_large(s, key, GRO_BUFFER)? } else { self.driver.udp(s, key)? };
996 self.driver.recv(id, true);
997 let slot = &mut self.udps[i];
998 slot.id = Some(id);
999 slot.local = local;
1000 self.ensure_endpoint(i);
1001 }
1002 }
1003 Ok(())
1004 }
1005
1006 fn spread_parked(&mut self) {
1008 self.peers = self.shared.handles();
1009 let n = self.peers.len().max(1);
1010 for (j, (socket, peer, li)) in std::mem::take(&mut self.parked).into_iter().enumerate() {
1011 let p = (self.rr + j) % n;
1012 match self.peers.get(p) {
1013 Some(h) if h.index != self.index => {
1014 if !h.run(move |lcx| lcx.core.adopt(socket, peer, li)) {
1015 self.rejected += 1;
1016 }
1017 }
1018 _ => self.adopt(socket, peer, li),
1019 }
1020 }
1021 }
1022
1023 fn pick_loop(&mut self) -> Option<usize> {
1027 if self.shared.started.load(Ordering::Acquire) != self.peers.len() {
1028 self.peers = self.shared.handles();
1030 }
1031 let n = self.peers.len();
1032 if n <= 1 {
1033 return None;
1034 }
1035 self.rr = (self.rr + 1) % n;
1036 let a = self.rr;
1037 let b = (a + 1 + self.rr * 7 % (n - 1).max(1)) % n;
1038 let la = self.shared.loop_active[self.peers[a].index].load(Ordering::Relaxed);
1039 let lb = self.shared.loop_active[self.peers[b].index].load(Ordering::Relaxed);
1040 let p = if lb < la { b } else { a };
1041 (self.peers[p].index != self.index).then_some(p)
1042 }
1043
1044 fn dispatch(&mut self, ev: Event) {
1045 match ev {
1046 Event::Accepted { key, socket, peer } => {
1047 let li = (key & !LISTENER_BIT) as usize;
1048 if self.want_grow || self.grow_next || (self.grown && self.shared.growing.load(Ordering::Acquire)) {
1049 self.parked.push((socket, peer, li));
1053 return;
1054 }
1055 let handoff = self.listeners.get(li).is_some_and(|l| l.handoff) || self.now_ns < self.handoff_until;
1056 match if handoff { self.pick_loop() } else { None } {
1057 None => self.adopt(socket, peer, li),
1058 Some(p) => {
1059 let ok = self.peers[p].run(move |lcx| lcx.core.adopt(socket, peer, li));
1060 if !ok {
1061 self.rejected += 1;
1062 }
1063 }
1064 }
1065 }
1066 Event::Data { key, buf } => self.on_data(key, buf),
1067 Event::Eof { key } => {
1068 if self.conns.get_mut(key).is_some_and(|c| c.lingering) {
1069 self.close_now(key, CloseReason::Closed);
1071 return;
1072 }
1073 self.with_conn(key, |h, cx| {
1074 if !cx.is_closing() {
1075 h.on_eof(cx)
1076 }
1077 })
1078 }
1079 Event::Error { key, error } => {
1080 if key & (LISTENER_BIT | UDP_BIT) != 0 {
1081 return; }
1083 self.close_now(key, CloseReason::Error(error));
1084 }
1085 Event::Sent { key, buf } => self.on_sent(key, buf),
1086 Event::Closed { key } => self.on_closed(key),
1087 #[cfg(not(feature = "udp"))]
1088 Event::Datagram { buf, .. } => self.pool.put(buf),
1089 #[cfg(feature = "udp")]
1090 Event::Datagram { key, buf, meta } => {
1091 let idx = (key & !UDP_BIT) as usize;
1092 let data = buf;
1093 self.with_udp(idx, |h, cx| h.on_datagram(cx, &data, &meta));
1094 if self.udps.get(idx).is_some_and(|u| u.gro) {
1095 self.driver.provide(data); } else {
1097 self.pool.put(data);
1098 }
1099 }
1100 Event::Wake => {}
1101 }
1102 }
1103
1104 fn on_data(&mut self, key: u64, buf: Vec<u8>) {
1105 let max_input = self.cfg.max_input;
1106 let Some(c) = self.conns.get_mut(key) else {
1107 self.pool.put(buf);
1108 return;
1109 };
1110 if c.closed || c.closing {
1111 let n = buf.len();
1112 self.pool.put(buf);
1113 if c.lingering {
1114 if n > c.linger_left as usize {
1115 self.close_now(key, CloseReason::Closed);
1116 } else {
1117 c.linger_left -= n as u32;
1118 }
1119 }
1120 return;
1121 }
1122 match c.input.as_mut() {
1123 None => c.input = Some(buf),
1124 Some(cur) => {
1125 if cur.capacity() - cur.len() < buf.len() {
1126 let mut bigger = self.pool.get(cur.len() + buf.len());
1127 bigger.extend_from_slice(cur);
1128 let old = std::mem::replace(cur, bigger);
1129 self.pool.put(old);
1130 }
1131 cur.extend_from_slice(&buf);
1132 self.pool.put(buf);
1133 }
1134 }
1135 let c = self.conns.get_mut(key).unwrap();
1136 if !c.user_paused && c.input.as_ref().map_or(0, |i| i.len()) > max_input {
1140 self.close_now(key, CloseReason::InputLimit);
1141 return;
1142 }
1143 if c.user_paused {
1144 return;
1145 }
1146 self.call_data(key);
1147 self.flush_conn(key);
1148 }
1149
1150 #[inline]
1154 fn flush_conn(&mut self, key: u64) {
1155 if self.batch_left < FLUSH_EARLY_BATCH {
1156 return;
1157 }
1158 if let Some(c) = self.conns.get_mut(key) {
1159 if !c.closed {
1160 let id = c.id;
1161 self.driver.flush_now(id);
1162 }
1163 }
1164 }
1165
1166 fn call_data(&mut self, key: u64) {
1167 let Some(c) = self.conns.get_mut(key) else { return };
1168 if c.closed || c.closing || c.user_paused {
1169 return;
1170 }
1171 let Some(input) = c.input.take() else { return };
1172 let Some(mut h) = c.handler.take() else {
1173 c.input = Some(input);
1174 return;
1175 };
1176 let n = {
1177 let mut cx = Ctx { core: self, key };
1178 h.on_data(&mut cx, &input)
1179 };
1180 let Some(c) = self.conns.get_mut(key) else {
1181 self.pool.put(input);
1182 return;
1183 };
1184 c.handler = Some(h);
1185 let mut input = input;
1186 if n >= input.len() {
1187 self.pool.put(input);
1189 } else {
1190 let rest = input.len() - n;
1191 if rest <= COMPACT_INPUT && input.capacity() >= 4 * rest {
1192 let tail = input[n..].to_vec();
1196 self.pool.put(input);
1197 if c.input.is_none() {
1198 c.input = Some(tail);
1199 }
1200 return;
1201 }
1202 if n > 0 {
1203 let len = input.len();
1204 input.copy_within(n..len, 0);
1205 input.truncate(len - n);
1206 }
1207 if c.input.is_none() {
1208 c.input = Some(input);
1209 } else {
1210 self.pool.put(input);
1211 }
1212 }
1213 }
1214
1215 fn on_sent(&mut self, key: u64, buf: OutBuf) {
1216 let len = buf.len();
1217 if let OutBuf::Owned(v) = buf {
1218 self.pool.put(v);
1219 }
1220 let low = self.cfg.write_low_water;
1221 let Some(c) = self.conns.get_mut(key) else { return };
1222 c.queued = c.queued.saturating_sub(len);
1223 if c.closed {
1224 return;
1225 }
1226 if c.closing {
1227 if c.queued == 0 {
1228 self.linger(key);
1229 }
1230 return;
1231 }
1232 let mut drain = false;
1233 if c.hw_paused && c.queued <= low {
1234 c.hw_paused = false;
1235 drain = true;
1236 if !c.user_paused {
1237 let id = c.id;
1238 let has_input = c.input.as_ref().is_some_and(|i| !i.is_empty());
1239 self.driver.recv(id, true);
1240 if has_input {
1241 self.redeliver.push(key);
1242 }
1243 }
1244 }
1245 let c = self.conns.get_mut(key).unwrap();
1246 if c.want_drain && c.queued <= low {
1247 c.want_drain = false;
1248 drain = true;
1249 }
1250 if drain {
1251 self.with_conn(key, |h, cx| h.on_drain(cx));
1252 }
1253 }
1254
1255 fn on_closed(&mut self, key: u64) {
1256 if key & LISTENER_BIT != 0 {
1257 let i = (key & !LISTENER_BIT) as usize;
1258 if let Some(l) = self.listeners.get_mut(i) {
1259 l.id = None;
1260 }
1261 return;
1262 }
1263 if cfg!(feature = "udp") && key & UDP_BIT != 0 {
1264 let i = (key & !UDP_BIT) as usize;
1265 if let Some(u) = self.udps.get_mut(i) {
1266 u.id = None;
1267 self.endpoints[i] = None;
1268 }
1269 return;
1270 }
1271 let Some(mut c) = self.conns.remove(key) else { return };
1272 c.timers.cancel_all(&mut self.timers);
1273 if let Some(i) = c.input.take() {
1274 self.pool.put(i);
1275 }
1276 self.shared.active.fetch_sub(1, Ordering::AcqRel);
1277 self.closed += 1;
1278 let reason = c.reason.take().unwrap_or(CloseReason::Closed);
1279 if let Some(mut h) = c.handler.take() {
1280 h.on_close(reason);
1281 }
1282 }
1283
1284 fn fire_timers(&mut self) {
1285 let now_tick = self.now_ns / self.tick_ns;
1286 self.timers.collect(now_tick);
1287 while let Some(t) = self.timers.pop_ready() {
1288 match t {
1289 TimerTarget::Conn(key, token) => self.with_conn(key, |h, cx| h.on_timer(cx, token)),
1290 TimerTarget::Udp(idx, token) => self.with_udp(idx, |h, cx| h.on_timer(cx, token)),
1291 TimerTarget::Linger(key) => self.close_now(key, CloseReason::Closed),
1292 }
1293 }
1294 }
1295
1296 #[cfg(feature = "idle-trim")]
1299 #[cold]
1300 #[inline(never)]
1301 fn idle_trim(&mut self, timeout: Option<Duration>) -> Option<Duration> {
1302 let at = self.last_io.saturating_add(IDLE_TRIM_NS);
1303 if self.now_ns >= at {
1304 self.pool.trim(IDLE_POOL_KEEP);
1305 self.driver.trim_buffers(INITIAL_RECV_BUFFERS);
1306 self.io_since_trim = 0;
1307 simoom_sys::alloc::release_free_memory();
1308 return timeout;
1309 }
1310 let left = Duration::from_nanos(at - self.now_ns);
1311 Some(timeout.map_or(left, |t| t.min(left)))
1312 }
1313
1314 fn after_callbacks(&mut self) {
1315 while !self.redeliver.is_empty() || !self.drains.is_empty() {
1316 for key in std::mem::take(&mut self.redeliver) {
1317 self.call_data(key);
1318 }
1319 for key in std::mem::take(&mut self.drains) {
1320 self.with_conn(key, |h, cx| h.on_drain(cx));
1321 }
1322 }
1323 }
1324
1325 fn publish(&self) {
1326 let s = &self.shared.loops[self.index];
1327 let d = self.driver.stats();
1328 s.store(d, self.pool.stats(), self.accepted, self.rejected, self.closed, self.conns.len as u64);
1329 s.timers.store(self.timers.len() as u64, Ordering::Relaxed);
1330 self.shared.loop_active[self.index].store(self.conns.len, Ordering::Relaxed);
1331 }
1332
1333 fn stop_done(&self, aborted: bool) -> bool {
1336 self.conns.len == 0 && !self.listeners.iter().any(|l| l.id.is_some()) && (aborted || !(0..self.udp_count()).any(|i| self.endpoint_busy(i)))
1337 }
1338
1339 fn udp_count(&self) -> usize {
1342 if cfg!(feature = "udp") {
1343 self.udps.len()
1344 } else {
1345 0
1346 }
1347 }
1348
1349 fn endpoint_busy(&self, i: usize) -> bool {
1350 self.endpoints[i].as_ref().is_some_and(|e| e.is_busy())
1351 }
1352
1353 fn begin_stop(&mut self) {
1354 for i in 0..self.listeners.len() {
1355 if let Some(id) = self.listeners[i].id {
1356 self.driver.close(id);
1357 }
1358 }
1359 for key in self.conns.keys() {
1360 self.with_conn(key, |h, cx| {
1361 if !cx.is_closing() {
1362 h.on_stop(cx)
1363 }
1364 });
1365 }
1366 for i in 0..self.udp_count() {
1367 self.with_udp(i, |h, cx| h.on_stop(cx));
1368 }
1369 }
1370
1371 pub(crate) fn run(&mut self) -> io::Result<()> {
1373 self.start_io()?;
1374 let mut events = Vec::with_capacity(256);
1375 let mut tasks: Vec<Item> = Vec::new();
1376 let mut stop_at: Option<u64> = None;
1377 let mut aborted = false;
1378 let recv_size = self.cfg.buffer_classes[0];
1379 loop {
1380 self.now_ns = simoom_sys::time::now_ns();
1381 self.fire_timers();
1382 self.after_callbacks();
1383 if stop_at.is_some() && self.stop_done(aborted) {
1386 break;
1387 }
1388 for _ in 0..self.driver.buffers_wanted() {
1389 let b = self.pool.get(recv_size);
1390 self.driver.provide(b);
1391 }
1392 let mut timeout = self.timers.next_deadline().map(|t| Duration::from_nanos((t * self.tick_ns).saturating_sub(self.now_ns)));
1393 if let Some(at) = stop_at {
1394 let left = Duration::from_nanos(at.saturating_sub(self.now_ns));
1395 timeout = Some(timeout.map_or(left, |t| t.min(left)));
1396 if aborted {
1397 timeout = Some(timeout.unwrap().min(Duration::from_millis(10)));
1398 }
1399 }
1400 if self.grow_next || (!self.parked.is_empty() && self.grown && !self.shared.growing.load(Ordering::Acquire)) {
1401 timeout = Some(Duration::ZERO);
1405 }
1406 #[cfg(feature = "idle-trim")]
1407 if self.pool.stats().cached > 3 * IDLE_POOL_KEEP || self.io_since_trim >= BURST_EVENTS {
1408 timeout = self.idle_trim(timeout);
1409 }
1410 self.driver.turn(timeout, &mut events)?;
1411 self.now_ns = simoom_sys::time::now_ns();
1412 #[cfg(feature = "idle-trim")]
1413 if !events.is_empty() {
1414 self.last_io = self.now_ns;
1415 self.io_since_trim += events.len() as u64;
1416 }
1417 let mut woke = false;
1418 let n = events.len();
1419 for (i, ev) in events.drain(..).enumerate() {
1420 if matches!(ev, Event::Wake) {
1421 woke = true;
1422 }
1423 self.batch_left = n - i;
1424 self.dispatch(ev);
1425 }
1426 if woke {
1427 self.queue.take(&mut tasks);
1428 let n = tasks.len();
1429 for (i, t) in tasks.drain(..).enumerate() {
1430 self.batch_left = n - i;
1431 match t {
1432 Item::Task(t) => t(&mut LoopCtx { core: self }),
1433 Item::Conn(key, msg) => self.deliver(key, msg),
1434 Item::Udp(idx, msg) => self.deliver_udp(idx, msg),
1435 }
1436 }
1437 }
1438 self.batch_left = 0;
1439 if self.grown && self.handoff_until == 0 && !self.shared.growing.load(Ordering::Acquire) {
1440 self.handoff_until = self.now_ns.saturating_add(BALANCE_AFTER_GROWTH_NS).max(1);
1443 self.spread_parked();
1444 simoom_sys::cpu::keep_current_thread_off_net_rx();
1447 }
1448 if self.grow_next {
1449 self.grow_next = false;
1452 self.grown = true;
1453 self.shared.grow();
1454 } else if self.want_grow {
1455 self.want_grow = false;
1457 self.grow_next = true;
1458 }
1459 self.after_callbacks();
1460 self.publish();
1461
1462 if stop_at.is_none() && self.shared.stop.load(Ordering::Acquire) {
1463 let grace = self.shared.grace_ns.load(Ordering::Acquire);
1464 stop_at = Some(self.now_ns.saturating_add(grace));
1465 self.begin_stop();
1466 }
1467 if let Some(at) = stop_at {
1468 if !aborted && self.now_ns >= at {
1469 aborted = true;
1470 for key in self.conns.keys() {
1471 self.close_now(key, CloseReason::Stopped);
1472 }
1473 for i in 0..self.udp_count() {
1474 if self.endpoint_busy(i) {
1475 self.with_udp(i, |h, cx| h.on_abort(cx));
1476 }
1477 }
1478 }
1479 if self.stop_done(aborted) {
1480 break;
1481 }
1482 }
1483 }
1484 for i in 0..self.udp_count() {
1485 if let Some(id) = self.udps[i].id.take() {
1486 self.driver.close(id);
1487 }
1488 }
1489 let end = simoom_sys::time::now_ns() + 200_000_000;
1491 while !self.driver.is_empty() && simoom_sys::time::now_ns() < end {
1492 self.driver.turn(Some(Duration::from_millis(10)), &mut events)?;
1493 for ev in events.drain(..) {
1494 if let Event::Sent { buf: OutBuf::Owned(v), .. } = ev {
1495 self.pool.put(v);
1496 }
1497 }
1498 }
1499 self.queue.close();
1500 self.publish();
1501 Ok(())
1502 }
1503}