Skip to content

Commit 4ea9e84

Browse files
youknowoneJMLX42
andauthored
Give the strong count the bits an unused weak field held (#8574)
* Give the strong count the bits an unused weak field held RefCount packs its state into one usize and splits the non-flag bits evenly between a strong and a weak count, so a 32-bit target left the strong count 15 bits. An object reaching 32 767 references then called refcount_overflow, which aborts the process; on wasm that surfaces as an unreachable trap with no Python-level error. Two ordinary module imports pass that total. The weak half is never read and never written. WEAK_COUNT appears once outside its own definition, in RefCount::new, and nothing in the repository observes those bits: weak references live in the object's WeakRefList, which Py::weak_count walks. The strong count now takes every bit the three flags leave, which is 29 on a 32-bit target and 61 on a 64-bit one. The tests run on wasm32-wasip2, where the old layout traps and the new one passes. A 64-bit host passes either way, since 31 bits already covered them. Fixes #8469 Assisted-by: Claude Code:claude-opus-5 * Return the strong count as usize so a 64-bit target cannot truncate it The count is 61 bits wide on a 64-bit target once it takes the bits the weak field held, so a u32 return truncates above 4 294 967 295 references: inc would read a wrapped 0 and re-increment, and dec would read a wrapped 1 and free a live object. A 32-bit target is unaffected, where 29 bits fit u32. Assisted-by: Claude Code:claude-opus-5 * Cover the refcount ceiling through inc and dec inc_by writes its own overflow check, so a test built on it leaves inc and dec unexercised. The reported aborts came through inc, one reference at a time. Assisted-by: Claude Code:claude-opus-5 --------- Co-authored-by: Jean-Marc Le Roux <jeanmarc@lx.industries>
1 parent ed62d4a commit 4ea9e84

1 file changed

Lines changed: 68 additions & 16 deletions

File tree

‎crates/common/src/refcount.rs‎

Lines changed: 68 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,23 @@
11
use crate::atomic::{Ordering, PyAtomic, Radium};
22

33
// State layout (usize):
4-
// [1 bit: destructed] [1 bit: published] [1 bit: leaked] [N bits: weak_count] [M bits: strong_count]
5-
// 64-bit: N=30, M=31. 32-bit: N=14, M=15.
4+
// [1 bit: destructed] [1 bit: published] [1 bit: leaked] [M bits: strong_count]
5+
// 64-bit: M=61. 32-bit: M=29.
6+
//
7+
// Weak references live in the object's `WeakRefList`, not in this word, so the
8+
// strong count takes every bit the flags leave. A 32-bit target reaches its
9+
// ceiling at 536 870 911 references rather than the 32 767 that half the word
10+
// would allow — a number two ordinary module imports pass on `wasm32`.
611
const FLAG_BITS: u32 = 3;
712
const DESTRUCTED: usize = 1 << (usize::BITS - 1);
813
/// Object was published to a lock-free cache; memory reclamation is
914
/// deferred through QSBR so concurrent try-incref readers never touch
1015
/// freed memory. Sticky once set.
1116
const PUBLISHED: usize = 1 << (usize::BITS - 2);
1217
const LEAKED: usize = 1 << (usize::BITS - 3);
13-
const TOTAL_COUNT_WIDTH: u32 = usize::BITS - FLAG_BITS;
14-
const WEAK_WIDTH: u32 = TOTAL_COUNT_WIDTH / 2;
15-
const STRONG_WIDTH: u32 = TOTAL_COUNT_WIDTH - WEAK_WIDTH;
18+
const STRONG_WIDTH: u32 = usize::BITS - FLAG_BITS;
1619
const STRONG: usize = (1 << STRONG_WIDTH) - 1;
1720
const COUNT: usize = 1;
18-
const WEAK_COUNT: usize = 1 << STRONG_WIDTH;
1921

2022
#[inline(never)]
2123
#[cold]
@@ -48,8 +50,8 @@ impl State {
4850
}
4951

5052
#[inline]
51-
fn strong(self) -> u32 {
52-
((self.inner & STRONG) / COUNT) as u32
53+
fn strong(self) -> usize {
54+
(self.inner & STRONG) / COUNT
5355
}
5456

5557
#[inline]
@@ -76,8 +78,8 @@ impl State {
7678
/// Reference count using state layout with LEAKED support.
7779
///
7880
/// State layout (usize):
79-
/// 64-bit: [1 bit: destructed] [1 bit: published] [1 bit: leaked] [30 bits: weak_count] [31 bits: strong_count]
80-
/// 32-bit: [1 bit: destructed] [1 bit: published] [1 bit: leaked] [14 bits: weak_count] [15 bits: strong_count]
81+
/// 64-bit: [1 bit: destructed] [1 bit: published] [1 bit: leaked] [61 bits: strong_count]
82+
/// 32-bit: [1 bit: destructed] [1 bit: published] [1 bit: leaked] [29 bits: strong_count]
8183
pub struct RefCount {
8284
state: PyAtomic<usize>,
8385
}
@@ -92,23 +94,22 @@ impl RefCount {
9294
/// Create a new RefCount with strong count = 1
9395
#[must_use]
9496
pub fn new() -> Self {
95-
// Initial state: strong=1, weak=1 (implicit weak for strong refs)
9697
Self {
97-
state: Radium::new(COUNT + WEAK_COUNT),
98+
state: Radium::new(COUNT),
9899
}
99100
}
100101

101102
/// Get current strong count
102103
#[inline]
103104
pub fn get(&self) -> usize {
104-
State::from_raw(self.state.load(Ordering::Relaxed)).strong() as usize
105+
State::from_raw(self.state.load(Ordering::Relaxed)).strong()
105106
}
106107

107108
/// Increment strong count
108109
#[inline]
109110
pub fn inc(&self) {
110111
let val = State::from_raw(self.state.fetch_add(COUNT, Ordering::Relaxed));
111-
if val.destructed() || (val.strong() as usize) > STRONG - 1 {
112+
if val.destructed() || val.strong() > STRONG - 1 {
112113
refcount_overflow();
113114
}
114115
if val.strong() == 0 {
@@ -121,7 +122,7 @@ impl RefCount {
121122
pub fn inc_by(&self, n: usize) {
122123
debug_assert!(n <= STRONG);
123124
let val = State::from_raw(self.state.fetch_add(n * COUNT, Ordering::Relaxed));
124-
if val.destructed() || (val.strong() as usize) > STRONG - n {
125+
if val.destructed() || val.strong() > STRONG - n {
125126
refcount_overflow();
126127
}
127128
}
@@ -135,7 +136,7 @@ impl RefCount {
135136
if old.destructed() || old.strong() == 0 {
136137
return false;
137138
}
138-
if (old.strong() as usize) >= STRONG {
139+
if old.strong() >= STRONG {
139140
refcount_overflow();
140141
}
141142
let new_state = old.add_strong(1);
@@ -299,6 +300,57 @@ pub fn flush_deferred_drops() {
299300
mod tests {
300301
use super::*;
301302

303+
/// The strong count reaches far past a 16-bit ceiling on every target.
304+
///
305+
/// The count shares its word with the flag bits, so its width follows the
306+
/// pointer width. A 32-bit target is the one this guards: a second counter
307+
/// packed beside the strong count once left it 15 bits, and `wasm32`
308+
/// aborted at 32 767 references — a total two ordinary module imports
309+
/// pass. The check is a no-op on a 64-bit host, where 31 bits already
310+
/// covered this; run the crate's tests against a 32-bit target to exercise
311+
/// it.
312+
#[test]
313+
fn strong_count_reaches_past_a_16_bit_ceiling() {
314+
const REFERENCES: usize = 1 << 20;
315+
316+
let rc = RefCount::new();
317+
rc.inc_by(REFERENCES);
318+
assert_eq!(rc.get(), REFERENCES + 1);
319+
}
320+
321+
/// `inc` and `dec` reach the same ceiling as `inc_by`.
322+
///
323+
/// The aborts reported against this layout came one reference at a time
324+
/// through `inc`, whose overflow check is written separately from
325+
/// `inc_by`'s, and the count has to come back down through `dec` without
326+
/// reporting the object collectable before the last reference goes.
327+
#[test]
328+
fn inc_and_dec_reach_past_a_16_bit_ceiling() {
329+
const REFERENCES: usize = 1 << 20;
330+
331+
let rc = RefCount::new(); // strong = 1
332+
for _ in 1..REFERENCES {
333+
rc.inc();
334+
}
335+
assert_eq!(rc.get(), REFERENCES);
336+
for _ in 1..REFERENCES {
337+
assert!(!rc.dec());
338+
}
339+
assert_eq!(rc.get(), 1);
340+
assert!(rc.dec());
341+
}
342+
343+
/// A fresh count holds exactly one strong reference and no stray bits.
344+
///
345+
/// `get` masks the flags away, so a spare field left in the word would not
346+
/// show up there. Reading the raw state keeps the layout honest.
347+
#[test]
348+
fn a_new_refcount_holds_one_strong_reference_and_nothing_else() {
349+
let rc = RefCount::new();
350+
assert_eq!(rc.get(), 1);
351+
assert_eq!(rc.state.load(Ordering::Relaxed), COUNT);
352+
}
353+
302354
#[test]
303355
fn published_bit_survives_refcount_traffic() {
304356
let rc = RefCount::new(); // strong = 1

0 commit comments

Comments
 (0)