ThreadId generation fallback path: avoid spurious yields

[RalfJung]

Fixes rust-lang/miri#4737
Alternative to #149476
Cc @orlp @joboet

[rustbot]

r? @ChrisDenton

rustbot has assigned @ChrisDenton.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

[joboet]

Looks fine, r=me with or without the nit.

View changes since this review

[joboet]

In this case, I'd use swap(true) here, as it can generate better assembly on platforms that don't have a CAS instruction, but do have swap (like RISC-V, even though it always has 64-bit atomics if I understand correctly).

[RalfJung]

Like this?

[joboet]

Yes, exactly!
@bors r+

[bors]

📌 Commit 1fccfa6 has been approved by joboet

It is now in the queue for this repository.

[bors]

⌛ Testing commit 1fccfa6 with merge 3e671f5...

[bors]

💔 Test failed - checks-actions

[RalfJung]

I build-checked it this time. 🙈

@bors r=joboet

[bors]

📌 Commit 14830bf has been approved by joboet

It is now in the queue for this repository.

[orlp]

@joboet If there is interest, I also have this lock-free version:

#![feature(allocator_api)]

use std::alloc::System;
use std::ptr;
use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};

struct State {
    base_value: u64,
    increment: AtomicUsize,
    prev: AtomicPtr<State>,
}

impl State {
    fn new_ptr(base_value: u64, prev: *mut State) -> *mut State {
        Box::into_raw_with_allocator(Box::new_in(
            State {
                base_value,
                increment: AtomicUsize::new(0),
                prev: AtomicPtr::new(prev),
            },
            System,
        ))
        .0
    }

    unsafe fn drop_ptr(mut state: *mut State) {
        unsafe {
            while state != ptr::null_mut() {
                let next = (*state).prev.load(Ordering::Acquire);
                unsafe { drop(Box::from_raw_in(state, System)) };
                state = next;
            }
        }
    }
}

// A lock-free 64-bit counter using just AtomicPtr and AtomicUsize.
pub struct LockFreeU64Counter {
    state: AtomicPtr<State>,
    in_progress_increments: AtomicUsize,
}

impl LockFreeU64Counter {
    pub fn new() -> Self {
        LockFreeU64Counter {
            state: AtomicPtr::new(State::new_ptr(0, ptr::null_mut())),
            in_progress_increments: AtomicUsize::new(0),
        }
    }

    /// Increments the counter and returns the old value.
    ///
    /// Returns None if the counter has reached its maximum value.
    pub fn increment(&self) -> Option<u64> {
        self.in_progress_increments.fetch_add(1, Ordering::Relaxed);

        loop {
            let current_state_ptr = self.state.load(Ordering::Acquire);
            let current_state = unsafe { &*current_state_ptr };
            let base_value = current_state.base_value;
            let increment = current_state.increment.load(Ordering::Relaxed);
            let cur_value = base_value.checked_add(increment as u64)?;

            if increment == usize::MAX {
                let new_state_ptr = State::new_ptr(cur_value, current_state_ptr);
                let cx = self.state.compare_exchange(
                    current_state_ptr,
                    new_state_ptr,
                    Ordering::AcqRel,
                    Ordering::Relaxed,
                );
                if cx.is_err() {
                    unsafe {
                        (*new_state_ptr).prev.store(ptr::null_mut(), Ordering::Relaxed);
                        State::drop_ptr(new_state_ptr)
                    };
                }
                continue;
            }

            let cx = current_state.increment.compare_exchange(
                increment,
                increment + 1,
                Ordering::Relaxed,
                Ordering::Relaxed,
            );
            if cx.is_err() {
                continue;
            }

            if self.in_progress_increments.fetch_sub(1, Ordering::AcqRel) == 1 {
                if current_state.prev.load(Ordering::Relaxed) != ptr::null_mut() {
                    let prev = current_state.prev.swap(ptr::null_mut(), Ordering::Acquire);
                    unsafe { State::drop_ptr(prev) };
                }
            }
            return Some(cur_value);
        }
    }
}

impl Drop for LockFreeU64Counter {
    fn drop(&mut self) {
        unsafe { State::drop_ptr(self.state.load(Ordering::Relaxed)) };
    }
}

EDIT: as ralf points out this has a bug, working on a fix.

[RalfJung]

@orlp I think that has a potential UAF since you are using current_state.prev after "releasing the lock" by decrementing in_progress_increments so another thread might come in, add a new state, and deallocate what used to be the current state.

[orlp]

@RalfJung I had already spotted another bug, but yes that is also an issue, let me fix that.

[orlp]

@RalfJung Please see the new implementation of the lock-free counter, which has a completely separate freelist:

#![feature(allocator_api)]

use std::alloc::System;
use std::cell::UnsafeCell;
use std::ptr;
use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};

struct State {
    base_value: u64,
    increment: AtomicUsize,
    // UnsafeCell because the freelist accesses this field exclusively while
    // the rest of the state may still be accessed concurrently.
    next_free: UnsafeCell<*mut State>,
}

impl State {
    fn new_ptr(base_value: u64) -> *mut State {
        Box::into_raw_with_allocator(Box::new_in(
            State {
                base_value,
                increment: AtomicUsize::new(0),
                next_free: UnsafeCell::new(ptr::null_mut()),
            },
            System,
        ))
        .0
    }

    unsafe fn drop_ptr(state: *mut State) {
        unsafe { drop(Box::from_raw_in(state, System)); }
    }
}

// A lock-free 64-bit counter using just AtomicPtr and AtomicUsize.
pub struct LockFreeU64Counter {
    state: AtomicPtr<State>,
    in_progress_increments: AtomicUsize,
    free_head: AtomicPtr<State>,
}

impl LockFreeU64Counter {
    pub fn new() -> Self {
        LockFreeU64Counter {
            state: AtomicPtr::new(State::new_ptr(0)),
            in_progress_increments: AtomicUsize::new(0),
            free_head: AtomicPtr::new(ptr::null_mut()),
        }
    }

    /// Increments the counter and returns the old value.
    ///
    /// Returns None if the counter has reached its maximum value.
    pub fn increment(&self) -> Option<u64> {
        // Acquire to ensure we don't see an old state that was reclaimed.
        self.in_progress_increments.fetch_add(1, Ordering::Acquire);

        loop {
            let current_state_ptr = self.state.load(Ordering::Acquire);
            let current_state = unsafe { &*current_state_ptr };
            let base_value = current_state.base_value;
            let increment = current_state.increment.load(Ordering::Acquire);
            let cur_value = base_value.checked_add(increment as u64)?;

            if increment == usize::MAX {
                let new_state_ptr = State::new_ptr(cur_value);
                let cx = self.state.compare_exchange(
                    current_state_ptr,
                    new_state_ptr,
                    Ordering::AcqRel,
                    Ordering::Relaxed,
                );
                match cx {
                    Ok(_) => unsafe { self.add_to_freelist(current_state_ptr) },
                    Err(_) => unsafe { State::drop_ptr(new_state_ptr) },
                }
                continue;
            }

            let cx = current_state.increment.compare_exchange(
                increment,
                increment + 1,
                Ordering::Release,
                Ordering::Relaxed,
            );
            if cx.is_err() {
                continue;
            }

            if self.in_progress_increments.fetch_sub(1, Ordering::AcqRel) == 1 {
                // Pass base_value as the limit to avoid reclaiming states that were added after
                // we established in_progress_increments == 0, which can still be accessed.
                unsafe { self.try_reclaim_freelist(base_value) };
            }
            return Some(cur_value);
        }
    }

    unsafe fn try_reclaim_freelist(&self, base_value_limit: u64) {
        if self.free_head.load(Ordering::Relaxed) == ptr::null_mut() {
            return;
        }

        unsafe {
            let mut head = self.free_head.swap(ptr::null_mut(), Ordering::Acquire);
            let mut filtered_head = ptr::null_mut();

            while head != ptr::null_mut() {
                let next = *(*head).next_free.get_mut();
                if (*head).base_value < base_value_limit {
                    State::drop_ptr(head);
                    head = next;
                } else {
                    *(*head).next_free.get_mut() = filtered_head;
                    filtered_head = head;
                    head = next;
                }
            }

            if filtered_head != ptr::null_mut() {
                self.add_to_freelist(filtered_head);
            }
        }
    }

    unsafe fn add_to_freelist(&self, state: *mut State) {
        unsafe {
            let mut state_tail = state;
            while *(*state_tail).next_free.get_mut() != ptr::null_mut() {
                state_tail = *(*state_tail).next_free.get_mut();
            }

            let mut head = self.free_head.load(Ordering::Acquire);
            loop {
                *(*state_tail).next_free.get_mut() = head;
                match self.free_head.compare_exchange(
                    head,
                    state,
                    Ordering::Release,
                    Ordering::Acquire,
                ) {
                    Ok(_) => break,
                    Err(new_head) => head = new_head,
                }
            }
        }
    }
}

impl Drop for LockFreeU64Counter {
    fn drop(&mut self) {
        unsafe {
            let mut head = *self.free_head.get_mut();
            while head != ptr::null_mut() {
                let next = *(*head).next_free.get_mut();
                State::drop_ptr(head);
                head = next;
            }
            State::drop_ptr(*self.state.get_mut());
        }
    }
}

[bors]

⌛ Testing commit 14830bf with merge f783dde...

[bors]

💥 Test timed out

[RalfJung]

@bors retry

[bors]

⌛ Testing commit 14830bf with merge 9b82a4f...

[bors]

☀️ Test successful - checks-actions
Approved by: joboet
Pushing 9b82a4f to main...

[github-actions]

What is this?

This is an experimental post-merge analysis report that shows differences in test outcomes between the merged PR and its parent PR.

Comparing dfe1b8c (parent) -> 9b82a4f (this PR)

Test differences

Show 2 test diffs

2 doctest diffs were found. These are ignored, as they are noisy.

Test dashboard

Run

cargo run --manifest-path src/ci/citool/Cargo.toml -- \
    test-dashboard 9b82a4fffe2b215f488a6dfdc0b508235f37c85c --output-dir test-dashboard

And then open test-dashboard/index.html in your browser to see an overview of all executed tests.

Job duration changes

1. dist-aarch64-apple: 5342.4s -> 7351.5s (+37.6%)
2. dist-x86_64-apple: 6364.2s -> 8513.1s (+33.8%)
3. i686-gnu-1: 6758.7s -> 7615.7s (+12.7%)
4. test-various: 6717.0s -> 6004.9s (-10.6%)
5. x86_64-gnu-debug: 6038.4s -> 6662.2s (+10.3%)
6. dist-apple-various: 3366.1s -> 3679.9s (+9.3%)
7. x86_64-gnu-llvm-20-2: 5152.1s -> 5602.2s (+8.7%)
8. x86_64-gnu-llvm-20-3: 6303.3s -> 5777.5s (-8.3%)
9. aarch64-apple: 7587.2s -> 8215.6s (+8.3%)
10. dist-x86_64-freebsd: 4832.7s -> 5204.4s (+7.7%)

How to interpret the job duration changes?

Job durations can vary a lot, based on the actual runner instance
that executed the job, system noise, invalidated caches, etc. The table above is provided
mostly for t-infra members, for simpler debugging of potential CI slow-downs.

[rust-timer]

Finished benchmarking commit (9b82a4f): comparison URL.

Overall result: ✅ improvements - no action needed

@rustbot label: -perf-regression

Instruction count

Our most reliable metric. Used to determine the overall result above. However, even this metric can be noisy.

	mean	range	count
Regressions ❌ (primary)	-	-	0
Regressions ❌ (secondary)	-	-	0
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-0.1%	[-0.1%, -0.1%]	1
All ❌✅ (primary)	-	-	0

Max RSS (memory usage)

Results (primary 2.5%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	2.5%	[2.5%, 2.5%]	1
Regressions ❌ (secondary)	-	-	0
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-	-	0
All ❌✅ (primary)	2.5%	[2.5%, 2.5%]	1

Cycles

Results (secondary 2.1%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	-	-	0
Regressions ❌ (secondary)	2.1%	[2.1%, 2.1%]	1
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-	-	0
All ❌✅ (primary)	-	-	0

Binary size

Results (primary 0.0%, secondary 0.0%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	0.0%	[0.0%, 0.0%]	4
Regressions ❌ (secondary)	0.0%	[0.0%, 0.1%]	13
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-	-	0
All ❌✅ (primary)	0.0%	[0.0%, 0.0%]	4

Bootstrap: 472.475s -> 474.873s (0.51%)
Artifact size: 386.89 MiB -> 386.96 MiB (0.02%)