nuclei 0.5.4

Proactive IO & runtime system
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
//! IO buffers: each shard's pool, recycled, and charged before they exist.
//!
//! A shard reads into buffers of its own pool when the bytes cannot go
//! straight into the caller's slice (the read had to wait), and copies
//! writes that cannot go out at once into them. Buffers come in four size
//! classes; a request takes the smallest class that holds it. A freed buffer
//! goes back to its class's free list for the next read, so a pool that has
//! grown to what its connections need allocates nothing more. A datagram
//! larger than the largest class gets a buffer made for it, charged like any
//! other and freed when it comes back.
//!
//! # Accounting
//!
//! The pool counts every byte it took from the allocator, in use or free,
//! and takes none past its cap. At the cap it frees idle buffers of other
//! classes to make room; when every byte is in use, [`Pool::get`] answers
//! `None` and the side that asked waits ([`Pool::wait_for_return`]) until a
//! buffer comes back. So a shard's buffer memory is bounded by its cap,
//! whatever its connections do.
//!
//! The shard's receive group (io_uring's provided buffers, see
//! `rt::shard::recv`) is charged here too, for as long as it lives
//! ([`Pool::reserve`]): the pool's own buffers get what is left of the cap,
//! and a group that leaves no room for one of the smallest class is refused.
//!
//! # Where a buffer goes back
//!
//! A buffer dropped on its pool's own thread goes straight to the free list.
//! Dropped anywhere else (its read-ahead moved with a handle to another
//! thread, its write-behind written by a blocking-pool job), it goes onto
//! the pool's return list, linked through the returned buffers themselves:
//! pushed with a compare-and-swap, taken whole by the pool when its free
//! list runs dry, so nothing is allocated for it and nothing waits for
//! reclamation (the push and the take of proofs/tla/remote, without its
//! close: the home outlives every buffer). The shared part of a pool
//! ([`Home`]) lives as long as any of its buffers, and frees what comes back
//! after the pool itself is gone. Either way a buffer coming back wakes
//! whoever waits for one; nothing that drops a buffer may hold a source's
//! guard (see `io::source`), since the wake reaches into sources.

pub(crate) mod aligned;
pub(crate) mod owned;
#[cfg(any(ring, unix))]
pub(crate) mod spare;
pub(crate) mod vectored;

use core::cell::{Cell, RefCell};
use core::ptr::NonNull;
use core::task::Waker;
use std::alloc::{Layout, alloc_zeroed, dealloc};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering, fence};

use kovan_queue::seg_queue::SegQueue;

/// The size classes, smallest first.
pub(crate) const CLASSES: [usize; 4] = [512, 4 * 1024, 16 * 1024, 64 * 1024];
/// The largest class.
pub(crate) const LARGEST: usize = CLASSES[CLASSES.len() - 1];
/// The class tag of a buffer made for one request larger than every class.
const WHOLE: u8 = u8::MAX;
/// What a shard's pool may hold unless its runtime says otherwise.
pub(crate) const DEFAULT_CAP: usize = 64 << 20;

/// The largest class a pool whose own buffers may hold `room` bytes takes.
fn top_for(room: usize) -> Option<usize> {
    CLASSES.iter().rposition(|&size| size <= room)
}

/// The class a request of `len` bytes takes: the smallest that holds it, or
/// the largest.
fn class_for(len: usize) -> usize {
    CLASSES
        .iter()
        .position(|&size| size >= len)
        .unwrap_or(CLASSES.len() - 1)
}

/// How a class's buffers are laid out: page-aligned from 4 KiB, so the
/// larger ones can serve direct IO; cache-line aligned below.
fn layout(class: usize) -> Layout {
    let size = CLASSES[class];
    let align = if size >= 4096 { 4096 } else { 64 };
    // SAFETY: every class size is a non-zero multiple of its alignment, and
    // the alignment is a power of two.
    unsafe { Layout::from_size_align_unchecked(size, align) }
}

/// How a whole buffer of `size` bytes (a multiple of 4096) is laid out.
fn layout_whole(size: usize) -> Layout {
    // SAFETY: `size` is a non-zero multiple of the page alignment, and was
    // checked to fit a `u32` when the buffer was made.
    unsafe { Layout::from_size_align_unchecked(size, 4096) }
}

/// A buffer's memory, without its length.
struct Raw(NonNull<u8>);

// SAFETY: plain memory nobody else points at.
unsafe impl Send for Raw {}

/// A buffer on its way back to its pool.
struct Returned {
    raw: Raw,
    class: u8,
    cap: u32,
}

/// A returned buffer as it lies on the return list, written into its first
/// bytes: every buffer is at least 512 bytes, aligned to 64. No padding, so
/// every byte of the buffer stays initialized (a slice over it is sound).
#[repr(C)]
struct Back {
    next: *mut Back,
    cap: u32,
    class: u32,
}

/// The part of a pool its buffers go back to from any thread.
pub(crate) struct Home {
    /// The buffers given back, last first.
    returned: AtomicPtr<Back>,
    /// Someone waits for a buffer to come back.
    waiting: AtomicBool,
    /// Who: woken, all of them, by the next buffer that comes back.
    waiters: SegQueue<Waker>,
}

impl Home {
    /// Takes back a buffer dropped on a thread that is not its pool's.
    fn give_back(&self, returned: Returned) {
        let back = returned.raw.0.cast::<Back>().as_ptr();
        let mut head = self.returned.load(Ordering::Relaxed);
        loop {
            // SAFETY: the buffer is no one's but this call's until the
            // exchange publishes it, and holds a `Back`.
            unsafe {
                back.write(Back {
                    next: head,
                    cap: returned.cap,
                    class: u32::from(returned.class),
                });
            }
            match self.returned.compare_exchange_weak(
                head,
                back,
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(now) => head = now,
            }
        }
        // Pairs with the fence of `wait_for_return`: either this sees the
        // waiter's flag, or the waiter sees this buffer in the queue.
        fence(Ordering::SeqCst);
        if self.waiting.load(Ordering::Relaxed) {
            self.wake_waiters();
        }
    }

    /// Every buffer given back so far, the last first.
    fn take_all(&self) -> Taken {
        if self.returned.load(Ordering::Relaxed).is_null() {
            return Taken(core::ptr::null_mut());
        }
        Taken(self.returned.swap(core::ptr::null_mut(), Ordering::Acquire))
    }

    fn wake_waiters(&self) {
        if self.waiting.swap(false, Ordering::AcqRel) {
            while let Ok(waker) = self.waiters.try_pop() {
                waker.wake();
            }
        }
    }
}

impl Drop for Home {
    fn drop(&mut self) {
        // The pool is gone: what came back after it goes to the allocator.
        for r in self.take_all() {
            // SAFETY: allocated with this layout, owned by nobody else.
            unsafe { dealloc(r.raw.0.as_ptr(), returned_layout(r.class, r.cap)) };
        }
    }
}

/// Buffers taken off a return list, theirs alone now.
struct Taken(*mut Back);

impl Iterator for Taken {
    type Item = Returned;

    fn next(&mut self) -> Option<Returned> {
        let back = NonNull::new(self.0)?;
        // SAFETY: a buffer given back, this list's since the swap that took
        // it (Acquire, against the push's Release).
        let Back { next, cap, class } = unsafe { back.as_ptr().read() };
        self.0 = next;
        Some(Returned {
            raw: Raw(back.cast()),
            // Written from a `u8` by `give_back`.
            class: class as u8,
            cap,
        })
    }
}

fn returned_layout(class: u8, cap: u32) -> Layout {
    if class == WHOLE {
        layout_whole(cap as usize)
    } else {
        layout(usize::from(class))
    }
}

/// A shard's buffer pool.
pub(crate) struct Pool {
    home: Arc<Home>,
    free: [RefCell<Vec<NonNull<u8>>>; CLASSES.len()],
    /// Bytes the pool may hold, in use and free, the reservation included.
    cap: Cell<usize>,
    /// Of `charged`, the bytes [`reserve`](Self::reserve) took for memory the
    /// shard holds outside the pool.
    reserved: Cell<usize>,
    /// The largest class that fits beside the reservation, if any does.
    top: Cell<Option<usize>>,
    /// Bytes it holds.
    charged: Cell<usize>,
    /// Times a buffer was wanted with every byte under the cap in use.
    exhausted: Cell<u64>,
}

// SAFETY: the free buffers are memory the pool owns outright and the rest
// is shared through `Arc<Home>`; a pool moved before it is entered (a shard
// built on one thread, run on its worker) leaves nothing behind. While
// entered it is borrowed, so it cannot move.
unsafe impl Send for Pool {}

impl Pool {
    /// A pool holding at most `cap` bytes, and at least one buffer of the
    /// smallest class whatever `cap` says: IO through a pool that could hold
    /// none would never move.
    pub(crate) fn new(cap: usize) -> Self {
        let cap = cap.max(CLASSES[0]);
        Self {
            home: Arc::new(Home {
                returned: AtomicPtr::new(core::ptr::null_mut()),
                waiting: AtomicBool::new(false),
                waiters: SegQueue::new(),
            }),
            free: Default::default(),
            cap: Cell::new(cap),
            reserved: Cell::new(0),
            top: Cell::new(top_for(cap)),
            charged: Cell::new(0),
            exhausted: Cell::new(0),
        }
    }

    /// Bytes the pool holds, in use and free, the reservation included.
    pub(crate) fn charged(&self) -> usize {
        self.charged.get()
    }

    /// Charges `bytes` of memory the shard holds outside the pool (its
    /// receive group) for as long as it holds them, through the pool's own
    /// ledger: the pool's buffers get the rest of the cap. False, charging
    /// nothing, when that would leave no room for a buffer of the smallest
    /// class.
    #[cfg_attr(
        all(not(ring), not(test)),
        expect(dead_code, reason = "only a ring's receive group is reserved")
    )]
    pub(crate) fn reserve(&self, bytes: usize) -> bool {
        let reserved = self.reserved.get().saturating_add(bytes);
        if reserved.saturating_add(CLASSES[0]) > self.cap.get() || !self.charge(bytes) {
            return false;
        }
        self.reserved.set(reserved);
        self.top.set(top_for(self.room()));
        true
    }

    /// Gives back `bytes` that [`reserve`](Self::reserve) charged.
    pub(crate) fn unreserve(&self, bytes: usize) {
        let bytes = bytes.min(self.reserved.get());
        self.release(bytes);
        self.reserved.set(self.reserved.get() - bytes);
        self.top.set(top_for(self.room()));
    }

    /// The bytes [`reserve`](Self::reserve) holds.
    #[cfg(test)]
    pub(crate) fn reserved(&self) -> usize {
        self.reserved.get()
    }

    /// The bytes the pool's own buffers may hold: the cap less the
    /// reservation.
    pub(crate) fn room(&self) -> usize {
        self.cap.get() - self.reserved.get()
    }

    /// Times a buffer was wanted while every byte under the cap was in use.
    pub(crate) fn exhausted(&self) -> u64 {
        self.exhausted.get()
    }

    /// Whether `home` is this pool's.
    fn is_home(&self, home: &Arc<Home>) -> bool {
        Arc::ptr_eq(&self.home, home)
    }

    /// A buffer of at least `len` bytes (or of the largest class that fits
    /// the cap), empty, or `None` when every byte under the cap is in use.
    pub(crate) fn get(&self, len: usize) -> Option<PoolBuf> {
        self.get_class(class_for(len).min(self.top.get()?))
    }

    fn get_class(&self, class: usize) -> Option<PoolBuf> {
        // The free list is let go before the fallbacks, which reach it.
        let free = self.free[class].borrow_mut().pop();
        let ptr = free
            .or_else(|| self.take_returned(class))
            .or_else(|| self.allocate(class))?;
        Some(self.buf(ptr, class as u8, CLASSES[class]))
    }

    /// Whether a buffer holding all of `len` bytes could ever come from the
    /// pool, every other buffer back.
    pub(crate) fn could_hold_whole(&self, len: usize) -> bool {
        let size = if len <= LARGEST {
            CLASSES[class_for(len)]
        } else {
            len.next_multiple_of(4096)
        };
        size <= self.room()
    }

    /// The bytes the pool may hold, the reservation included.
    pub(crate) fn cap(&self) -> usize {
        self.cap.get()
    }

    /// Holds the pool to `cap` from now on (at least one buffer of the smallest class), freeing
    /// its idle buffers past it now; a pool still past it takes no new memory until it is back
    /// under it.
    fn hold_to(&self, cap: usize) {
        self.cap
            .set(cap.max(CLASSES[0]).max(self.reserved.get() + CLASSES[0]));
        self.top.set(top_for(self.room()));
        let _ = self.make_room(0);
    }

    /// A buffer that holds all of `len` bytes, whatever their number: a class
    /// buffer when one is large enough, else one made for them.
    pub(crate) fn get_whole(&self, len: usize) -> Option<PoolBuf> {
        if len <= LARGEST {
            return self.get_class(class_for(len));
        }
        let size = len.checked_next_multiple_of(4096)?;
        u32::try_from(size).ok()?;
        // Wholes that came back are freed here, which may make the room.
        let _ = self.take_returned(usize::MAX);
        if !self.charge(size) {
            return None;
        }
        // SAFETY: the layout has a non-zero size.
        let Some(ptr) = NonNull::new(unsafe { alloc_zeroed(layout_whole(size)) }) else {
            self.release(size);
            return None;
        };
        Some(self.buf(ptr, WHOLE, size))
    }

    fn buf(&self, ptr: NonNull<u8>, class: u8, cap: usize) -> PoolBuf {
        PoolBuf {
            ptr,
            class,
            cap: cap as u32,
            start: 0,
            end: 0,
            home: Arc::clone(&self.home),
        }
    }

    /// A buffer of `class` that came back from another thread; others found
    /// on the way go to their free lists, and wholes to the allocator.
    fn take_returned(&self, class: usize) -> Option<NonNull<u8>> {
        let mut found = None;
        for r in self.home.take_all() {
            if r.class == WHOLE {
                self.free_whole(r.raw.0, r.cap);
            } else if found.is_none() && usize::from(r.class) == class {
                found = Some(r.raw.0);
            } else {
                self.free[usize::from(r.class)].borrow_mut().push(r.raw.0);
            }
        }
        found
    }

    /// A new buffer of `class`, charged first; none past the cap.
    fn allocate(&self, class: usize) -> Option<NonNull<u8>> {
        let size = CLASSES[class];
        if !self.charge(size) {
            return None;
        }
        // Zeroed once, when the pool grows: every byte of a buffer is
        // initialized from then on, so a slice over it is always sound, and
        // a recycled buffer costs nothing.
        // SAFETY: the layout has a non-zero size.
        let ptr = NonNull::new(unsafe { alloc_zeroed(layout(class)) });
        if ptr.is_none() {
            self.release(size);
        }
        ptr
    }

    /// Counts `size` more bytes against the cap, freeing idle buffers to make
    /// the room; false, counting nothing, when even that leaves too little.
    /// With [`release`](Self::release), the one place the ledger moves.
    fn charge(&self, size: usize) -> bool {
        if !self.make_room(size) {
            return false;
        }
        self.charged.set(self.charged.get() + size);
        true
    }

    fn release(&self, size: usize) {
        self.charged.set(self.charged.get() - size);
    }

    /// Frees idle buffers, the largest first, until `size` more bytes fit
    /// under the cap. False when even all of them would not make the room.
    fn make_room(&self, size: usize) -> bool {
        let fits = |charged: usize| charged.saturating_add(size) <= self.cap.get();
        if fits(self.charged.get()) {
            return true;
        }
        for class in (0..CLASSES.len()).rev() {
            let mut free = self.free[class].borrow_mut();
            while let Some(ptr) = free.pop() {
                // SAFETY: allocated with this class's layout; free, so owned
                // by the pool alone.
                unsafe { dealloc(ptr.as_ptr(), layout(class)) };
                self.release(CLASSES[class]);
                if fits(self.charged.get()) {
                    return true;
                }
            }
        }
        self.exhausted.set(self.exhausted.get() + 1);
        false
    }

    fn free_whole(&self, ptr: NonNull<u8>, cap: u32) {
        // SAFETY: made with this layout, and owned by nobody else.
        unsafe { dealloc(ptr.as_ptr(), layout_whole(cap as usize)) };
        self.release(cap as usize);
    }

    /// Takes back a buffer dropped on this pool's own thread.
    fn take_back(&self, ptr: NonNull<u8>, class: u8, cap: u32) {
        if class == WHOLE {
            self.free_whole(ptr, cap);
        } else {
            self.free[usize::from(class)].borrow_mut().push(ptr);
        }
        // Only this thread raises the flag, so a plain load sees it.
        if self.home.waiting.load(Ordering::Relaxed) {
            self.home.wake_waiters();
        }
    }

    /// After a [`get`](Self::get) found nothing: `waker` is woken when a
    /// buffer comes back to this pool, from any thread; at once if one came
    /// back since, or if none is out to come back (the request is larger
    /// than the cap allows).
    pub(crate) fn wait_for_return(&self, waker: Waker) {
        if self.charged.get() == self.reserved.get() {
            waker.wake();
            return;
        }
        self.home.waiters.push(waker);
        self.home.waiting.store(true, Ordering::SeqCst);
        // Pairs with the fence of `Home::give_back`.
        fence(Ordering::SeqCst);
        let idle = self.free.iter().any(|f| !f.borrow().is_empty());
        if idle || !self.home.returned.load(Ordering::SeqCst).is_null() {
            self.home.wake_waiters();
        }
    }
}

impl Drop for Pool {
    fn drop(&mut self) {
        for (class, free) in self.free.iter_mut().enumerate() {
            for ptr in free.get_mut().drain(..) {
                // SAFETY: allocated with this class's layout; free, so owned
                // by the pool alone.
                unsafe { dealloc(ptr.as_ptr(), layout(class)) };
            }
        }
        // Buffers still out come back to `Home`, which frees them; whoever
        // waits for one is woken, to find the pool gone.
        self.home.wake_waiters();
    }
}

thread_local! {
    /// The pool of a thread that runs no shard (a task of another executor writing to a handle,
    /// a blocking pool's thread, a worker whose shard was handed over), made on first need. A
    /// thread that serves a runtime holds it to that runtime's cap ([`serve_with_cap`]).
    static FALLBACK: Pool = Pool::new(DEFAULT_CAP);
}

/// Holds the pool of this thread, which serves a runtime and may run no shard (a blocking
/// pool's thread, a worker, a thread inside the runtime with `enter`), to that runtime's
/// `io_buffer_cap`, as a worker's own pool is (proofs/tla/bufcap/BufCap.tla). Called as the
/// thread starts serving, before its IO; a pool that holds more than the cap meanwhile frees
/// its idle buffers past it now and takes no new memory until it is back under it. Here and not
/// in the pool's lazy start, so the IO steps that reach this thread's pool carry none of it.
pub(crate) fn serve_with_cap(cap: usize) {
    let _ = FALLBACK.try_with(|pool| pool.hold_to(cap));
}

/// Runs `f` with this thread's pool: the pool of the shard it runs (a call
/// from a task's code comes back into the shard first, see `rt::context`),
/// or its own if it runs none. `None` only while the thread exits, after its
/// pool is gone.
pub(crate) fn with_pool<R>(f: impl FnOnce(Option<&Pool>) -> R) -> R {
    // The held form: one call of `f`, which is inlined into the IO steps that ask.
    let current = crate::rt::context::current_held();
    match current.shard() {
        Some(shard) => f(Some(shard.pool())),
        None => {
            let fallback = FALLBACK
                .try_with(core::ptr::from_ref)
                .unwrap_or(core::ptr::null());
            // SAFETY: the fallback lives until the thread's destructors run,
            // which `try_with` would have reported.
            f(unsafe { fallback.as_ref() })
        }
    }
}

/// [`with_pool`], unless the caller already holds this thread's pool.
pub(crate) fn with_pool_or<R>(known: Option<&Pool>, f: impl FnOnce(Option<&Pool>) -> R) -> R {
    match known {
        Some(pool) => f(Some(pool)),
        None => with_pool(f),
    }
}

/// Copies the front of `src` into `bufs`, in order; returns how many bytes.
pub(crate) fn copy_vectored(src: &[u8], bufs: &mut [std::io::IoSliceMut<'_>]) -> usize {
    let mut n = 0;
    for buf in bufs.iter_mut() {
        let rest = &src[n..];
        if rest.is_empty() {
            break;
        }
        let k = rest.len().min(buf.len());
        buf[..k].copy_from_slice(&rest[..k]);
        n += k;
    }
    n
}

/// Copies the front of `src` onto the end of `buf`, as much as its spare
/// capacity holds; returns how many bytes.
pub(crate) fn append_spare(src: &[u8], buf: &mut owned::OwnedBuf) -> usize {
    buf.append(src)
}

/// A buffer of a shard's pool, with the bytes `start..end` filled and not
/// read yet.
pub(crate) struct PoolBuf {
    ptr: NonNull<u8>,
    class: u8,
    cap: u32,
    start: u32,
    end: u32,
    home: Arc<Home>,
}

// SAFETY: the buffer owns its memory; going back to its pool is either a
// thread-local free list (on the pool's own thread) or the lock-free return
// queue.
unsafe impl Send for PoolBuf {}
// SAFETY: `&PoolBuf` only reads the filled bytes.
unsafe impl Sync for PoolBuf {}

impl PoolBuf {
    /// The buffer's size.
    pub(crate) fn capacity(&self) -> usize {
        self.cap as usize
    }

    /// The filled bytes not read yet.
    pub(crate) fn filled(&self) -> &[u8] {
        // SAFETY: `start..end` is within the buffer and was written.
        unsafe {
            core::slice::from_raw_parts(
                self.ptr.as_ptr().add(self.start as usize),
                (self.end - self.start) as usize,
            )
        }
    }

    /// Whether every filled byte was read.
    pub(crate) fn is_empty(&self) -> bool {
        self.start == self.end
    }

    /// The whole buffer to fill from its start, discarding what it held; then
    /// [`set_filled`](Self::set_filled) says how much was written.
    pub(crate) fn spare(&mut self) -> &mut [u8] {
        self.start = 0;
        self.end = 0;
        // SAFETY: the buffer is ours, and every byte of it is initialized (it
        // was zeroed when it was allocated).
        unsafe { core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.capacity()) }
    }

    /// The first `n` bytes of the buffer were written.
    pub(crate) fn set_filled(&mut self, n: usize) {
        debug_assert!(n <= self.capacity());
        self.start = 0;
        self.end = n.min(self.capacity()) as u32;
    }

    /// The `n` bytes from `from` were written: they are the filled ones.
    pub(crate) fn set_filled_at(&mut self, from: usize, n: usize) {
        let from = from.min(self.capacity());
        self.start = from as u32;
        self.end = (from + n).min(self.capacity()) as u32;
    }

    /// Fills the buffer from `from` with `bytes` (as many as fit); returns
    /// how many.
    pub(crate) fn fill_at(&mut self, from: usize, bytes: &[u8]) -> usize {
        let n = bytes.len().min(self.capacity().saturating_sub(from));
        if let Some(room) = self.spare().get_mut(from..from + n) {
            room.copy_from_slice(&bytes[..n]);
        }
        self.set_filled_at(from, n);
        n
    }

    /// Fills the buffer with `bytes` (as many as fit); returns how many.
    #[cfg(any(completion, test))]
    pub(crate) fn fill_from(&mut self, bytes: &[u8]) -> usize {
        let n = bytes.len().min(self.capacity());
        self.spare()[..n].copy_from_slice(&bytes[..n]);
        self.set_filled(n);
        n
    }

    /// Adds `bytes` after the filled ones, as many as fit, moving the unread
    /// bytes to the front first when that makes room; returns how many.
    pub(crate) fn append(&mut self, bytes: &[u8]) -> usize {
        let cap = self.capacity();
        let (start, end) = (self.start as usize, self.end as usize);
        if start > 0 && end + bytes.len() > cap {
            // SAFETY: both ranges lie within the buffer; `copy` allows the
            // overlap.
            unsafe {
                core::ptr::copy(self.ptr.as_ptr().add(start), self.ptr.as_ptr(), end - start)
            };
            self.start = 0;
            self.end = (end - start) as u32;
        }
        let end = self.end as usize;
        let n = (cap - end).min(bytes.len());
        // SAFETY: `end..end + n` lies within the buffer, which `bytes` (a
        // shared borrow) cannot overlap.
        unsafe { core::ptr::copy_nonoverlapping(bytes.as_ptr(), self.ptr.as_ptr().add(end), n) };
        self.end += n as u32;
        n
    }

    /// `n` filled bytes were read.
    pub(crate) fn consume(&mut self, n: usize) {
        let n = n.min((self.end - self.start) as usize);
        self.start += n as u32;
    }

    /// Copies filled bytes into `out` and consumes them; returns how many.
    pub(crate) fn read_into(&mut self, out: &mut [u8]) -> usize {
        let n = self.peek_into(out);
        self.consume(n);
        n
    }

    /// Copies filled bytes into `bufs`, in order, and consumes them;
    /// returns how many.
    pub(crate) fn read_into_vectored(&mut self, bufs: &mut [std::io::IoSliceMut<'_>]) -> usize {
        let n = copy_vectored(self.filled(), bufs);
        self.consume(n);
        n
    }

    /// Adds the bytes of `bufs` after the filled ones, in order, as many as
    /// fit; returns how many.
    pub(crate) fn append_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> usize {
        let mut n = 0;
        for buf in bufs {
            let taken = self.append(buf);
            n += taken;
            if taken < buf.len() {
                break;
            }
        }
        n
    }

    /// Copies filled bytes into `out` without consuming them; returns how
    /// many.
    pub(crate) fn peek_into(&self, out: &mut [u8]) -> usize {
        let src = self.filled();
        let n = src.len().min(out.len());
        out[..n].copy_from_slice(&src[..n]);
        n
    }
}

impl Drop for PoolBuf {
    fn drop(&mut self) {
        let current = crate::rt::context::current_held();
        let returned = match current.shard() {
            Some(shard) if shard.pool().is_home(&self.home) => {
                shard.pool().take_back(self.ptr, self.class, self.cap);
                false
            }
            _ => true,
        };
        drop(current);
        if returned {
            self.home.give_back(Returned {
                raw: Raw(self.ptr),
                class: self.class,
                cap: self.cap,
            });
        }
    }
}

impl core::fmt::Debug for PoolBuf {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("PoolBuf")
            .field("capacity", &self.capacity())
            .field("filled", &(self.end - self.start))
            .finish()
    }
}

#[cfg(all(test, not(loom)))]
mod tests;