stringtheory/interning/fixed_size.rs
1#[cfg(not(feature = "loom"))]
2use std::sync::{atomic::AtomicUsize, Arc, Mutex};
3use std::{
4 num::NonZeroUsize,
5 ptr::NonNull,
6 sync::atomic::Ordering::{AcqRel, Acquire},
7};
8
9#[cfg(feature = "loom")]
10use loom::sync::{atomic::AtomicUsize, Arc, Mutex};
11
12use super::{
13 helpers::{aligned, aligned_string, hash_string, layout_for_data, PackedLengthCapacity},
14 InternedString, Interner,
15};
16use crate::interning::helpers::{ReclaimedEntries, ReclaimedEntry};
17
18const HEADER_LEN: usize = std::mem::size_of::<EntryHeader>();
19const HEADER_ALIGN: usize = std::mem::align_of::<EntryHeader>();
20
21/// The minimum possible length of an entry.
22///
23/// For any entry in the interner, there is already an `EntryHeader` followed by the string data itself. In order to
24/// ensure that entries can be written contiguously, we additionally ensure that the number of bytes we utilize for the
25/// string data is aligned at least as much as `EntryHeader` itself.
26///
27/// This means that the minimum possible length of an entry, or the minimum number of bytes a valid entry could consume,
28/// is the length of the header plus the alignment of the header.
29const MINIMUM_ENTRY_LEN: usize = HEADER_LEN + HEADER_ALIGN;
30
31#[derive(Debug)]
32pub(crate) struct StringState {
33 interner: Arc<Mutex<InternerShardState>>,
34 header: NonNull<EntryHeader>,
35}
36
37impl StringState {
38 #[inline]
39 pub const fn as_str(&self) -> &str {
40 // SAFETY: We ensure `self.header` is well-aligned and points to an initialized `EntryHeader` value when creating `StringState`.
41 unsafe { get_entry_string(self.header) }
42 }
43}
44
45impl PartialEq for StringState {
46 fn eq(&self, other: &Self) -> bool {
47 self.header == other.header
48 }
49}
50
51impl Clone for StringState {
52 fn clone(&self) -> Self {
53 // SAFETY: The caller that creates `StringState` is responsible for ensuring that `self.header` is well-aligned
54 // and points to an initialized `EntryHeader` value.
55 let header = unsafe { self.header.as_ref() };
56 header.increment_active_refs();
57
58 Self {
59 interner: self.interner.clone(),
60 header: self.header,
61 }
62 }
63}
64
65impl Drop for StringState {
66 fn drop(&mut self) {
67 // SAFETY: The caller that creates `StringState` is responsible for ensuring that `self.header` is well-aligned
68 // and points to an initialized `EntryHeader` value.
69 let header = unsafe { self.header.as_ref() };
70 if header.decrement_active_refs() {
71 // We decremented the reference count to zero, so try to mark this entry for reclamation.
72 let mut interner = self.interner.lock().unwrap();
73 interner.mark_for_reclamation(self.header);
74 }
75 }
76}
77
78// SAFETY: We don't take references to the entry header pointer that outlast `StringState`, and the only modification we
79// do to the entry header is through atomic operations, so it's safe to both send and share `StringState` between
80// threads.
81unsafe impl Send for StringState {}
82unsafe impl Sync for StringState {}
83
84/// Metadata about an interner entry.
85///
86/// `EntryHeader` represents the smallest amount of information about an interned entry that's needed to support both
87/// lookup of existing interned strings, as well as the ability to reclaim space in the interner when an entry is no
88/// longer in use.
89struct EntryHeader {
90 /// The hash of the string that this entry represents.
91 hash: u64,
92
93 /// The number of active references to this entry.
94 ///
95 /// Only incremented by the interner itself, and decremented by `InternedString` when it's dropped.
96 refs: AtomicUsize,
97
98 /// Combined length/capacity of the entry, in terms of the string itself.
99 ///
100 /// Notably, this does _not_ include the length of the header itself. For example, an entry holding the string
101 /// "hello, world!" has a string length of 13 bytes, but since we've to pad out to meet our alignment requirements
102 /// for `EntryHeader`, we would end up with a capacity of 16 bytes. As such, `EntryHeader::len` would report `13`,
103 /// while `EntryHeader::capacity` would report `16`. Likewise, `EntryHeader::entry_len` would report `40`,
104 /// accounting for the string capacity (16) as well as the header length itself (24).
105 ///
106 /// As explained in the description of `PackedLengthCapacity`, this does mean strings can't be larger than ~4 GB on
107 /// 64-bit platforms, which isn't a problem we've.
108 len_cap: PackedLengthCapacity,
109}
110
111impl EntryHeader {
112 /// Creates a tombstone entry with the given capacity.
113 ///
114 /// This is to allow for updating a region in the data buffer, which has been reclaimed, such that it's
115 /// identifiable as being unused.
116 fn tombstone(entry: ReclaimedEntry) -> Self {
117 // The usable capacity for a reclaimed entry is the full capacity minus the size of `EntryHeader` itself, as
118 // reclaimed entries represent the _entire_ region in the data buffer, but `EntryHeader` only cares about the
119 // string portion itself.
120 let cap = Self::usable_from_reclaimed(entry);
121
122 Self {
123 hash: 0,
124 refs: AtomicUsize::new(0),
125 len_cap: PackedLengthCapacity::new(cap, 0),
126 }
127 }
128
129 /// Creates a new entry for the given string.
130 fn from_string(hash: u64, s: &str) -> Self {
131 // We're dictating the necessary capacity here, which is the length of the string rounded to the nearest
132 // multiple of the alignment of `EntryHeader`, which ensures that any subsequent entry will be properly aligned.
133 let cap = aligned_string::<Self>(s);
134
135 Self {
136 hash,
137 refs: AtomicUsize::new(1),
138 len_cap: PackedLengthCapacity::new(cap, s.len()),
139 }
140 }
141
142 /// Creates a new entry for the given string, based on the given reclaimed entry.
143 ///
144 /// This maps the entry header to the underlying capacity of the given reclaimed entry, which is done in cases where
145 /// a reclaimed entry is being used when interning a new string, and the reclaimed entry is larger than the string
146 /// being interned, but not large enough that we could split the excess capacity into a new reclaimed entry.
147 fn from_reclaimed_entry(mut entry: ReclaimedEntry, hash: u64, s: &str) -> (Self, Option<ReclaimedEntry>) {
148 // The usable capacity for a reclaimed entry is the full capacity minus the size of `EntryHeader` itself, as
149 // reclaimed entries represent the _entire_ region in the data buffer, but `EntryHeader` only cares about the
150 // string portion itself.
151 let entry_cap = EntryHeader::usable_from_reclaimed(entry);
152 let required_cap = aligned_string::<Self>(s);
153
154 // If the reclaimed entry has enough additional space beyond what we need for the string, we'll split it off and
155 // return it for the caller to keep around in the reclaimed entries list.
156 let remainder = entry_cap - required_cap;
157 let (adjusted_cap, maybe_split_entry) = if remainder >= MINIMUM_ENTRY_LEN {
158 let entry_len = EntryHeader::len_for(s);
159 let split_entry = entry.split_off(entry_len);
160
161 (entry_len - HEADER_LEN, Some(split_entry))
162 } else {
163 (entry_cap, None)
164 };
165
166 let header = Self {
167 hash,
168 refs: AtomicUsize::new(1),
169 len_cap: PackedLengthCapacity::new(adjusted_cap, s.len()),
170 };
171
172 (header, maybe_split_entry)
173 }
174
175 /// Returns the computed length of a complete entry, in bytes, for the given string.
176 ///
177 /// This includes the size of the entry header itself and the string data, when padded for alignment, and represents
178 /// the number of bytes that would be consumed in the data buffer.
179 const fn len_for(s: &str) -> usize {
180 HEADER_LEN + aligned_string::<Self>(s)
181 }
182
183 /// Returns the usable capacity of a reclaimed entry, in bytes.
184 ///
185 /// Usable refers to the number of bytes in a reclaimed entry that could be used for string data, after accounting
186 /// for the size of `EntryHeader` itself.
187 const fn usable_from_reclaimed(entry: ReclaimedEntry) -> usize {
188 entry.capacity() - HEADER_LEN
189 }
190
191 /// Returns the total number of bytes that this entry takes up in the data buffer.
192 const fn entry_len(&self) -> usize {
193 HEADER_LEN + self.capacity()
194 }
195
196 /// Returns the size of the string, in bytes, that this entry can hold.
197 const fn capacity(&self) -> usize {
198 self.len_cap.capacity()
199 }
200
201 /// Returns the size of the string, in bytes, that this entry _actually_ holds.
202 const fn len(&self) -> usize {
203 self.len_cap.len()
204 }
205
206 /// Returns `true` if this entry is currently referenced.
207 fn is_active(&self) -> bool {
208 self.refs.load(Acquire) != 0
209 }
210
211 /// Increments the active reference count by one.
212 fn increment_active_refs(&self) {
213 self.refs.fetch_add(1, AcqRel);
214 }
215
216 /// Decrements the active reference count by one.
217 ///
218 /// Returns `true` if the active reference count is zero _after_ calling this method.
219 fn decrement_active_refs(&self) -> bool {
220 self.refs.fetch_sub(1, AcqRel) == 1
221 }
222
223 /// Attempts to increment the active reference count, but only if it is currently non-zero.
224 ///
225 /// Returns `true` if the reference count was incremented. A return value of `false` means the entry has already
226 /// been released -- its reference count reached zero -- and is pending reclamation, so it must not be reused.
227 fn try_increment_active_refs(&self) -> bool {
228 self.refs
229 .fetch_update(AcqRel, Acquire, |refs| (refs != 0).then_some(refs + 1))
230 .is_ok()
231 }
232}
233
234#[derive(Debug)]
235struct InternerShardState {
236 // Direct pieces of our buffer allocation.
237 ptr: NonNull<u8>,
238 offset: usize,
239 capacity: NonZeroUsize,
240
241 // Number of active entries (strings) in the interner.
242 entries: usize,
243
244 // Length of all active entries, in bytes.
245 //
246 // This is equivalent to `self.offset` minus the total size of all reclaimed entries.
247 len: usize,
248
249 // Markers for entries that can be reused.
250 reclaimed: ReclaimedEntries,
251}
252
253impl InternerShardState {
254 /// Creates a new `InternerShardState` with a pre-allocated buffer that has the given capacity.
255 pub fn with_capacity(capacity: NonZeroUsize) -> Self {
256 assert!(
257 aligned::<EntryHeader>(capacity.get()) <= isize::MAX as usize,
258 "capacity would overflow isize::MAX, which violates layout constraints"
259 );
260
261 // Allocate our data buffer. This is the main backing allocation for all interned strings, and is well-aligned
262 // for `EntryHeader`.
263 //
264 // SAFETY: `layout_for_data` ensures the layout is non-zero.
265 let data_layout = layout_for_data::<EntryHeader>(capacity);
266 let data_ptr = unsafe { std::alloc::alloc(data_layout) };
267 let ptr = match NonNull::new(data_ptr) {
268 Some(ptr) => ptr,
269 None => std::alloc::handle_alloc_error(data_layout),
270 };
271
272 Self {
273 ptr,
274 offset: 0,
275 capacity,
276 entries: 0,
277 len: 0,
278 reclaimed: ReclaimedEntries::new(),
279 }
280 }
281
282 /// Returns the total number of unused bytes that are available for interning.
283 fn available(&self) -> usize {
284 self.capacity.get() - self.offset
285 }
286
287 fn get_entry_ptr(&self, offset: usize) -> NonNull<EntryHeader> {
288 debug_assert!(
289 offset + MINIMUM_ENTRY_LEN <= self.capacity.get(),
290 "offset would point to entry that cannot possibly avoid extending past end of data buffer"
291 );
292
293 // SAFETY: The caller is responsible for ensuring that `offset` is within the bounds of the data buffer, and
294 // that `offset` is well-aligned for `EntryHeader`.
295 let entry_ptr = unsafe { self.ptr.as_ptr().add(offset).cast::<EntryHeader>() };
296 debug_assert!(entry_ptr.is_aligned(), "entry header pointer must be well-aligned");
297
298 // SAFETY: `entry_ptr` is derived from `self.ptr`, which itself is `NonNull<u8>`, and the caller is responsible
299 // for ensuring that `offset` is within the bounds of the data buffer, so we know `entry_ptr` is non-null.
300 unsafe { NonNull::new_unchecked(entry_ptr) }
301 }
302
303 fn find_entry(&self, hash: u64, s: &str) -> Option<NonNull<EntryHeader>> {
304 let mut offset = 0;
305
306 while offset < self.offset {
307 // Construct a pointer to the entry at `offset`, and get a reference to the header value.
308 let header_ptr = self.get_entry_ptr(offset);
309 let header = unsafe { header_ptr.as_ref() };
310
311 // See if this entry is active or not. If it's active, then we'll quickly check the hash/length of the
312 // string to see if this is likely to be a match for `s`.
313 if header.is_active() && header.hash == hash && header.len() == s.len() {
314 // As a final check, we make sure that the entry string and `s` are equal. If they are, then we
315 // have an exact match and will return the entry.
316 //
317 // SAFETY: We know that our header is valid and initialized.
318 let s_entry = unsafe { get_entry_string(header_ptr) };
319 if s_entry == s {
320 // Reuse this entry, but only if it's still live. The `is_active()` check above is not enough on its
321 // own: a concurrent `StringState::drop` decrements the reference count _before_ taking this lock, so
322 // it can drop the count to zero in the window between that check and here. Resurrecting a
323 // zero-refcount entry would let two droppers each observe a 1 -> 0 transition and each reclaim the
324 // same slot -- a double reclaim that corrupts the data buffer. If the conditional increment fails,
325 // the entry is already pending reclamation by its dropper, so we skip it and let the caller intern a
326 // fresh entry instead.
327 if header.try_increment_active_refs() {
328 return Some(header_ptr);
329 }
330 }
331 }
332
333 // Either this was a reclaimed entry or we didn't have a match, so we move on to the next entry.
334 offset += header.entry_len();
335 }
336
337 None
338 }
339
340 fn write_entry(&mut self, offset: usize, entry_header: EntryHeader, s: &str) -> NonNull<EntryHeader> {
341 debug_assert_eq!(
342 entry_header.len(),
343 s.len(),
344 "entry header length must match string length"
345 );
346
347 let entry_ptr = self.get_entry_ptr(offset);
348 let entry_len = entry_header.entry_len();
349
350 // Write the entry header.
351 unsafe { entry_ptr.as_ptr().write(entry_header) };
352
353 let s_buf = s.as_bytes();
354
355 // Write the string.
356 let entry_s_buf = unsafe {
357 // Take the entry pointer and add 1, which sets our pointer to right _after_ the header.
358 let entry_s_ptr = entry_ptr.as_ptr().add(1).cast::<u8>();
359 std::slice::from_raw_parts_mut(entry_s_ptr, s_buf.len())
360 };
361 entry_s_buf.copy_from_slice(s_buf);
362
363 // Update our internal statistics.
364 self.entries += 1;
365 self.len += entry_len;
366
367 entry_ptr
368 }
369
370 fn write_to_unoccupied(&mut self, s_hash: u64, s: &str) -> NonNull<EntryHeader> {
371 let entry_header = EntryHeader::from_string(s_hash, s);
372
373 // Write the entry to the end of the data buffer.
374 let entry_offset = self.offset;
375 self.offset += entry_header.entry_len();
376
377 self.write_entry(entry_offset, entry_header, s)
378 }
379
380 fn write_to_reclaimed_entry(&mut self, entry: ReclaimedEntry, s_hash: u64, s: &str) -> NonNull<EntryHeader> {
381 let entry_offset = entry.offset();
382 let (entry_header, maybe_split_entry) = EntryHeader::from_reclaimed_entry(entry, s_hash, s);
383
384 // If we had enough capacity in the reclaimed entry to hold this string _and_ potentially hold another entry, we
385 // split it off and store that remainder entry.
386 if let Some(split_entry) = maybe_split_entry {
387 self.add_reclaimed(split_entry);
388 }
389
390 // Write the entry in place of the reclaimed entry.
391 self.write_entry(entry_offset, entry_header, s)
392 }
393
394 fn add_reclaimed_from_header(&mut self, header_ptr: NonNull<EntryHeader>) {
395 // Get the offset of the header within the data buffer.
396 //
397 // SAFETY: The caller is responsible for ensuring the entry header reference belongs to this interner. If that
398 // is upheld, then we know that entry header belongs to our data buffer, and that the pointer to the entry
399 // header is not less than the base pointer of the data buffer, ensuring the offset is non-negative.
400 let entry_offset = unsafe {
401 header_ptr
402 .cast::<u8>()
403 .as_ptr()
404 .offset_from(self.ptr.as_ptr().cast_const())
405 };
406 debug_assert!(entry_offset >= 0, "entry offset must be non-negative");
407
408 let header = unsafe { header_ptr.as_ref() };
409
410 let entry = ReclaimedEntry::new(entry_offset as usize, header.entry_len());
411 self.len -= entry.capacity();
412 self.add_reclaimed(entry);
413 }
414
415 fn add_reclaimed(&mut self, entry: ReclaimedEntry) {
416 // Reclamation is a two-step process: first, we have to actually keep track of the reclaimed entry, which
417 // potentially involves merging adjacent reclaimed entries, and then once all of that has happened, we tombstone
418 // the entry (whether merged or not).
419 //
420 // However, if the merged reclaimed entry immediately precedes any available capacity, we can skip tombstoning
421 // it, since we can just wind back `offset` to reclaim the space.
422 let merged_entry = self.reclaimed.insert(entry);
423 if merged_entry.offset() + merged_entry.capacity() == self.offset {
424 self.offset -= merged_entry.capacity();
425 self.reclaimed.remove(&merged_entry);
426 return;
427 }
428
429 self.clear_reclaimed_entry(merged_entry);
430 }
431
432 fn mark_for_reclamation(&mut self, header_ptr: NonNull<EntryHeader>) {
433 // Reclaim the entry only if its reference count is zero, and rely on `find_entry` to never resurrect it.
434 //
435 // `StringState::drop` decrements the reference count _before_ taking this lock, so a concurrent `try_intern`
436 // can observe this entry during that window. But the only reuse path, `find_entry`, increments via
437 // `try_increment_active_refs`, which refuses to increment a zero reference count -- so once the count reaches
438 // zero it stays zero, and exactly one dropper (the one that transitioned it 1 -> 0) reaches this method. That
439 // makes reclamation exactly-once: no second thread can resurrect the entry and then drop it again, which would
440 // otherwise double-reclaim the same slot and corrupt the data buffer.
441 //
442 // The reference count also cannot rise while we hold this lock (cloning needs a live handle, and `try_intern`
443 // needs this lock), so observing zero here means it stays zero through the reclamation below -- no
444 // use-after-free.
445 //
446 // SAFETY: The caller is responsible for ensuring that `header_ptr` is well-aligned and points to an initialized
447 // `EntryHeader` value.
448 let header = unsafe { header_ptr.as_ref() };
449 if !header.is_active() {
450 self.entries -= 1;
451 self.add_reclaimed_from_header(header_ptr);
452 }
453 }
454
455 fn clear_reclaimed_entry(&mut self, entry: ReclaimedEntry) {
456 let entry_ptr = self.get_entry_ptr(entry.offset);
457
458 // Write the entry tombstone itself, which clears out the hash and sets the reference count to zero.
459 //
460 // SAFETY: We know that `entry_ptr` is valid for writes (reclaimed entries are, by definition, inactive regions
461 // in the data buffer) and is well-aligned for `EntryHeader`.
462 let tombstone = EntryHeader::tombstone(entry);
463 let str_cap = tombstone.capacity();
464
465 unsafe {
466 entry_ptr.as_ptr().write(EntryHeader::tombstone(entry));
467 }
468
469 // Write a magic value to the entire string capacity for the entry. This ensures that there's a known repeating
470 // value which, in the case of debugging issues, can be a signal that offsets/reclaimed entries are incorrect
471 // and overlapping with active entries.
472 //
473 // SAFETY: Like above, the caller is responsible for ensuring that `offset` is within the bounds of the data
474 // buffer, and that `offset + capacity` does not extend past the bounds of the data buffer.
475 unsafe {
476 // Take the entry pointer and add 1, which sets our pointer to right _after_ the header.
477 let str_ptr = entry_ptr.as_ptr().add(1).cast::<u8>();
478 let str_buf = std::slice::from_raw_parts_mut(str_ptr, str_cap);
479 str_buf.fill(0xAA);
480 }
481 }
482
483 fn try_intern(&mut self, s_hash: u64, s: &str) -> Option<NonNull<EntryHeader>> {
484 // We can only intern strings with a size that fits within a packed length/capacity value, so if `s` is larger
485 // than that, we can't intern it, and there's existing entry we could have for it either.
486 if s.len() > PackedLengthCapacity::maximum_value() {
487 return None;
488 }
489
490 // Try and find an existing entry for this string.
491 if self.entries != 0 {
492 if let Some(existing_entry) = self.find_entry(s_hash, s) {
493 return Some(existing_entry);
494 }
495 }
496
497 let required_cap = EntryHeader::len_for(s);
498
499 // We didn't find an existing entry, so we're going to intern it.
500 //
501 // First, try and see if we have a reclaimed entry that can fit this string. If nothing suitable is found, or we
502 // have no reclaimed entries, then we'll just try to fit it in the remaining capacity of our data buffer.
503 if !self.reclaimed.is_empty() {
504 let maybe_reclaimed_entry = self.reclaimed.take_if(|entry| entry.capacity() >= required_cap);
505 if let Some(reclaimed_entry) = maybe_reclaimed_entry {
506 return Some(self.write_to_reclaimed_entry(reclaimed_entry, s_hash, s));
507 }
508 }
509
510 if required_cap <= self.available() {
511 Some(self.write_to_unoccupied(s_hash, s))
512 } else {
513 None
514 }
515 }
516}
517
518impl Drop for InternerShardState {
519 fn drop(&mut self) {
520 // SAFETY: We allocated this buffer with the global allocator, and we're generating the same layout that was
521 // used to allocate it in the first place.
522 unsafe {
523 std::alloc::dealloc(self.ptr.as_ptr(), layout_for_data::<EntryHeader>(self.capacity));
524 }
525 }
526}
527
528// SAFETY: We don't take references to the data buffer pointer that outlast `InternerShardState`, and all access to
529// `InternerShardState` itself is mediated through a mutex, so we're safe to send it around and share it between
530// threads.
531unsafe impl Send for InternerShardState {}
532unsafe impl Sync for InternerShardState {}
533
534#[derive(Debug)]
535struct InternerState<const SHARD_FACTOR: usize> {
536 shards: [Arc<Mutex<InternerShardState>>; SHARD_FACTOR],
537 capacity: usize,
538}
539
540impl<const SHARD_FACTOR: usize> InternerState<SHARD_FACTOR> {
541 // Ensure our shard factor is a power of two at compile, so that we can just use a mask to get the shard index.
542 const _POWER_OF_TWO_SHARD_FACTOR: () = {
543 if !SHARD_FACTOR.is_power_of_two() {
544 panic!("shard factor must be a power of two")
545 }
546 };
547
548 pub fn with_capacity(capacity: NonZeroUsize) -> Self {
549 let shard_capacity = match NonZeroUsize::new(capacity.get() / SHARD_FACTOR) {
550 Some(shard_capacity) => shard_capacity,
551 // If we can't divide the capacity evenly, we just specify a capacity of one which will force every shard to
552 // upsize the capacity so that a single entry can fit, and satisfies the need for `NonZeroUsize`.
553 //
554 // SAFETY: One is obviously not zero.
555 None => unsafe { NonZeroUsize::new_unchecked(1) },
556 };
557
558 let shards = std::iter::repeat_with(|| Arc::new(Mutex::new(InternerShardState::with_capacity(shard_capacity))))
559 .take(SHARD_FACTOR)
560 .collect::<Vec<_>>();
561 let capacity = shards
562 .iter()
563 .map(|shard| {
564 let shard = shard.lock().unwrap();
565 shard.capacity.get()
566 })
567 .sum();
568 let shards: [Arc<Mutex<InternerShardState>>; SHARD_FACTOR] =
569 shards.try_into().expect("should not fail to convert to array");
570
571 Self { shards, capacity }
572 }
573
574 fn is_empty(&self) -> bool {
575 self.shards.iter().any(|shard| shard.lock().unwrap().entries != 0)
576 }
577
578 fn len(&self) -> usize {
579 self.shards.iter().map(|shard| shard.lock().unwrap().entries).sum()
580 }
581
582 fn len_bytes(&self) -> usize {
583 self.shards.iter().map(|shard| shard.lock().unwrap().offset).sum()
584 }
585
586 fn capacity_bytes(&self) -> usize {
587 self.capacity
588 }
589
590 fn try_intern(&self, s: &str) -> Option<InternedString> {
591 let hash = hash_string(s);
592 let shard_idx = (hash as usize) & (SHARD_FACTOR - 1);
593
594 let shard = &self.shards[shard_idx];
595 intern_with_shard_and_hash(shard, hash, s)
596 }
597}
598
599/// A string interner based on a single, fixed-size backing buffer with support for reclamation.
600///
601/// ## Overview
602///
603/// This interner uses a single, fixed-size backing buffer where interned strings are stored contiguously. This provides
604/// bounded memory usage, and the interner won't allocate additional memory for new strings once the buffer is full.
605/// Since interned strings aren't likely to need to live for the life of the program, the interner supports
606/// reclamation. Once all references to an interned string have been dropped, the storage for that string is reclaimed
607/// and can be used to hold new strings.
608///
609/// ## Storage layout
610///
611/// The backing buffer stores strings contiguously, with an entry "header" prepended to each string. The header contains
612/// relevant data -- hash of the string, reference count, and length of the string -- needed to work with the entry
613/// either when searching for existing entries or when using the entry itself.
614///
615/// The layout of an entry is as follows:
616///
617/// ```text
618/// ┌───────────────────────── entry #1 ──────────────────────────┐ ┌─ entry #2 ─┐ ┌─ entry .. ─┐
619/// ▼ ▼ ▼ ▼ ▼ ▼
620/// ┏━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━┓
621/// ┃ str hash │ ref cnt │ str len │ str data │ padding ┃ ┃ header ┃ ┃ header ┃
622/// ┃ (8 bytes) │ (8 bytes) │ (8 bytes) │ (N bytes) │ (1-7 bytes) ┃ ┃ & string ┃ ┃ & string ┃
623/// ┗━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━┛
624/// ▲ ▲ ▲
625/// └────────── `EntryHeader` ──────────┘ └── aligned for `EntryHeader`
626/// (8 byte alignment) via trailing padding
627/// ```
628///
629/// The backing buffer is always aligned properly for `EntryHeader`, so that the first entry can be referenced
630/// correctly. However, when appending additional entries to the buffer, we need to ensure that those entries also have
631/// an aligned start for accessing the header. This is complicated due to the variable number of bytes for the string
632/// data.
633///
634/// Alignment padding is added to the end of the entry to ensure that when appending the next entry, the start of the
635/// entry is properly aligned for `EntryHeader`. In the worst case, up to 7 bytes could be added (and thus wasted) on
636/// this alignment padding.
637///
638/// ## `InternedString`
639///
640/// The `InternedString` type is a handle to the entry header, and thus the string data, for an interned string. It's
641/// designed to be small -- 8 bytes. -- and cheap to clone, as it contains an atomic reference to the entry header and a
642/// reference to the interner that owns the string. It dereferences to the underlying string with relatively low
643/// overhead: two pointer indirections.
644///
645/// When an `InternedString` is dropped, it decrements the reference count for the entry it points to. If the reference
646/// count drops to zero, it will attempt to mark the entry for reclamation.
647///
648/// ## Reclamation
649///
650/// As we want to bound the memory used by the interner, but also not allow it to be filled up with strings that
651/// eventually end up going entirely unused, we need a way to remove those unused strings so their underlying storage
652/// can be used for new strings. This is where reclamation comes in.
653///
654/// When a string is interned, the entry header tracks how many active references there are to it. When that reference
655/// count drops to zero, the last reference to the string attempts to mark the entry for reclamation. Assuming no other
656/// reference has been taken out on the entry in the meantime, the entry gets added to a list of "reclaimed" entries.
657///
658/// Reclaimed entries are simple markers -- start and end position in the data buffer -- which are stored in a freelist.
659/// When attempting to intern a new string, this freelist is searched to see if there's an entry large enough to fit the
660/// new string, and if so, it's used.
661///
662/// Additionally, when entries are reclaimed, adjacent entries are merged together where possible. This helps to avoid
663/// unnecessary fragmentation over time, although not as effectively as reconstructing the data buffer to re-pack
664/// entries.
665#[derive(Clone, Debug)]
666pub struct FixedSizeInterner<const SHARD_FACTOR: usize> {
667 state: Arc<InternerState<SHARD_FACTOR>>,
668}
669
670impl<const SHARD_FACTOR: usize> FixedSizeInterner<SHARD_FACTOR> {
671 /// Creates a new `FixedSizeInterner` with the given capacity.
672 ///
673 /// The given capacity will potentially be rounded up by a small number of bytes (up to 7) in order to ensure the
674 /// backing buffer is properly aligned.
675 pub fn new(capacity: NonZeroUsize) -> Self {
676 Self {
677 state: Arc::new(InternerState::with_capacity(capacity)),
678 }
679 }
680}
681
682impl<const SHARD_FACTOR: usize> Interner for FixedSizeInterner<SHARD_FACTOR> {
683 fn is_empty(&self) -> bool {
684 self.state.is_empty()
685 }
686
687 fn len(&self) -> usize {
688 self.state.len()
689 }
690
691 fn len_bytes(&self) -> usize {
692 self.state.len_bytes()
693 }
694
695 fn capacity_bytes(&self) -> usize {
696 self.state.capacity_bytes()
697 }
698
699 fn try_intern(&self, s: &str) -> Option<InternedString> {
700 self.state.try_intern(s)
701 }
702}
703
704#[inline]
705const unsafe fn get_entry_string_parts(header_ptr: NonNull<EntryHeader>) -> (NonNull<u8>, usize) {
706 // SAFETY: The caller is responsible for ensuring that `header_ptr` is well-aligned and points to an initialized
707 // `EntryHeader` value.
708 let header = header_ptr.as_ref();
709
710 // Advance past the header and get the pointer to the string.
711 //
712 // SAFETY: We know that we're simply skipping over the header by advancing the pointer by one when it's still typed
713 // as `*mut EntryHeader`.
714 let s_ptr = header_ptr.add(1).cast::<u8>();
715 (s_ptr, header.len())
716}
717
718#[inline]
719const unsafe fn get_entry_string<'a>(header_ptr: NonNull<EntryHeader>) -> &'a str {
720 let (s_ptr, s_len) = get_entry_string_parts(header_ptr);
721
722 // SAFETY: We depend on `get_entry_string_parts` to give us a valid pointer and length for the string.
723 std::str::from_utf8_unchecked(std::slice::from_raw_parts(s_ptr.as_ptr() as *const _, s_len))
724}
725
726fn intern_with_shard_and_hash(shard: &Arc<Mutex<InternerShardState>>, hash: u64, s: &str) -> Option<InternedString> {
727 let header = {
728 let mut shard = shard.lock().unwrap();
729 shard.try_intern(hash, s)?
730 };
731
732 Some(InternedString::from(StringState {
733 interner: Arc::clone(shard),
734 header,
735 }))
736}
737
738#[cfg(test)]
739mod tests {
740 use std::{
741 collections::HashSet,
742 ops::{Deref as _, RangeInclusive},
743 };
744
745 use prop::sample::Index;
746 use proptest::{
747 collection::{hash_set, vec as arb_vec},
748 prelude::*,
749 };
750
751 use super::*;
752 use crate::interning::InternedStringState;
753
754 pub(super) fn create_shard(capacity: NonZeroUsize) -> Arc<Mutex<InternerShardState>> {
755 Arc::new(Mutex::new(InternerShardState::with_capacity(capacity)))
756 }
757
758 pub(super) fn intern_for_shard(shard: &Arc<Mutex<InternerShardState>>, s: &str) -> Option<InternedString> {
759 let hash = hash_string(s);
760 intern_with_shard_and_hash(shard, hash, s)
761 }
762
763 fn shard_capacity(shard: &Arc<Mutex<InternerShardState>>) -> usize {
764 shard.lock().unwrap().capacity.get()
765 }
766
767 fn shard_available(shard: &Arc<Mutex<InternerShardState>>) -> usize {
768 shard.lock().unwrap().available()
769 }
770
771 fn shard_entries(shard: &Arc<Mutex<InternerShardState>>) -> usize {
772 shard.lock().unwrap().entries
773 }
774
775 fn shard_reclaimed_len(shard: &Arc<Mutex<InternerShardState>>) -> usize {
776 shard.lock().unwrap().reclaimed.len()
777 }
778
779 fn shard_first_reclaimed_entry(shard: &Arc<Mutex<InternerShardState>>) -> ReclaimedEntry {
780 shard.lock().unwrap().reclaimed.first().unwrap()
781 }
782
783 pub(super) fn get_reclaimed_entry_for_string(s: &InternedString) -> ReclaimedEntry {
784 let state = match &s.state {
785 InternedStringState::FixedSize(state) => state,
786 _ => panic!("unexpected string state"),
787 };
788
789 let ptr = state.interner.lock().unwrap().ptr.as_ptr();
790 let header = unsafe { state.header.as_ref() };
791 let offset = unsafe { state.header.as_ptr().cast::<u8>().offset_from(ptr) as usize };
792 ReclaimedEntry::new(offset, header.entry_len())
793 }
794
795 fn entry_len(s: &str) -> usize {
796 EntryHeader::len_for(s)
797 }
798
799 fn arb_alphanum_strings(
800 str_len: RangeInclusive<usize>, unique_strs: RangeInclusive<usize>,
801 ) -> impl Strategy<Value = Vec<String>> {
802 // Create characters between 0x20 (32) and 0x7E (126), which are all printable ASCII characters.
803 let char_gen = any::<u8>().prop_map(|c| std::cmp::max(c % 127, 32));
804
805 let str_gen = any::<usize>()
806 .prop_map(move |n| std::cmp::max(n % *str_len.end(), *str_len.start()))
807 .prop_flat_map(move |len| arb_vec(char_gen.clone(), len))
808 // SAFETY: We know our characters are all valid UTF-8 because they're in the ASCII range.
809 .prop_map(|xs| unsafe { String::from_utf8_unchecked(xs) });
810
811 // Create a hash set, which handles the deduplication aspect for us, ensuring we have N unique strings where N
812 // is within the `unique_strs` range... and then convert it to `Vec<String>` for easier consumption.
813 hash_set(str_gen, unique_strs).prop_map(|unique_strs| unique_strs.into_iter().collect::<Vec<_>>())
814 }
815
816 #[test]
817 fn basic() {
818 let interner = FixedSizeInterner::<1>::new(NonZeroUsize::new(1024).unwrap());
819
820 let s1 = interner.try_intern("hello").unwrap();
821 let s2 = interner.try_intern("world").unwrap();
822 let s3 = interner.try_intern("hello").unwrap();
823
824 assert_eq!(s1.deref(), "hello");
825 assert_eq!(s2.deref(), "world");
826 assert_eq!(s3.deref(), "hello");
827
828 // The pointers from the interned strings should be the same, but not between the interned string and a pointer
829 // to an equivalent (but not interned) string:
830 assert!(std::ptr::eq(s1.deref() as *const _, s3.deref() as *const _));
831
832 let local_hello = "hello";
833 assert!(!std::ptr::eq(s1.deref() as *const _, local_hello as *const _));
834 }
835
836 #[test]
837 fn try_intern_without_capacity() {
838 // Big enough to fit a single "hello world!" string, but not big enough to fit two.
839 let interner = FixedSizeInterner::<1>::new(NonZeroUsize::new(64).unwrap());
840
841 let s1 = interner.try_intern("hello world!");
842 assert!(s1.is_some());
843
844 let s2 = interner.try_intern("hello, world");
845 assert!(s2.is_none());
846 }
847
848 #[test]
849 fn reclaim_after_dropped() {
850 let shard = create_shard(NonZeroUsize::new(1024).unwrap());
851
852 let s1 = intern_for_shard(&shard, "hello").expect("should not fail to intern");
853 let s1_entry_len = entry_len(&s1);
854
855 assert_eq!(shard_entries(&shard), 1);
856 assert_eq!(shard_available(&shard), 1024 - s1_entry_len);
857 assert_eq!(shard_reclaimed_len(&shard), 0);
858
859 // Drop the interned string, which should decrement the reference count to zero and then reclaim the entry.
860 drop(s1);
861
862 assert_eq!(shard_entries(&shard), 0);
863 assert_eq!(shard_available(&shard), 1024);
864 assert_eq!(shard_reclaimed_len(&shard), 0);
865 }
866
867 #[test]
868 fn interns_to_reclaimed_entry_with_leftover() {
869 // We want to intern a string initially, which takes up almost all of the capacity, and then drop it so it gets
870 // reclaimed. After that, we'll intern a much smaller string which should lead to utilizing that reclaimed
871 // entry, but only a part of it. Finally, we'll intern another new string.
872 //
873 // The point is to demonstrate that our reclamation logic is sound in terms of allowing reclaimed entries to be
874 // split while the search/insertion logic is operating.
875 let shard_capacity = NonZeroUsize::new(256).unwrap();
876 let shard = create_shard(shard_capacity);
877
878 // We craft four strings such that the first two (`s_large` and `s_medium1`) will take up enough capacity that
879 // `s_small` can't possibly be interned in the available capacity. We'll also craft `s_medium2` so it can fit
880 // within the reclaimed entry for `s_large` but takes enough capacity that `s_small` cannot fit in the leftover
881 // reclaimed entry that we split off.
882 let s_large = "99 bottles of beer on the wall, 99 bottles of beer! take one down, pass it around, 98 bottles of beer on the wall!";
883 let s_medium1 = "no act of kindness, no matter how small, is ever wasted";
884 let s_medium2 = "if you want to go fast, go alone; if you want to go far, go together";
885 let s_small = "are you there god? it's me, margaret";
886
887 let phase1_available_capacity = shard_capacity.get() - entry_len(s_large) - entry_len(s_medium1);
888 assert!(phase1_available_capacity < entry_len(s_small));
889 assert!((entry_len(s_large) - entry_len(s_medium2)) < entry_len(s_small));
890 assert!(entry_len(s_medium2) < entry_len(s_large));
891
892 // Phase 1: intern our two larger strings.
893 let s1 = intern_for_shard(&shard, s_large).expect("should not fail to intern");
894 let s2 = intern_for_shard(&shard, s_medium1).expect("should not fail to intern");
895
896 assert_eq!(shard_entries(&shard), 2);
897 assert_eq!(shard_available(&shard), phase1_available_capacity);
898 assert_eq!(shard_reclaimed_len(&shard), 0);
899
900 // Phase 2: drop `s_large` so it gets reclaimed.
901 drop(s1);
902 assert_eq!(shard_entries(&shard), 1);
903 assert_eq!(shard_reclaimed_len(&shard), 1);
904
905 // Phase 3: intern `s_medium2`, which should fit in the reclaimed entry for `s_large`. This should leave a
906 // small, split off reclaimed entry.
907 let s3 = intern_for_shard(&shard, s_medium2).expect("should not fail to intern");
908
909 assert_eq!(shard_entries(&shard), 2);
910 assert_eq!(shard_reclaimed_len(&shard), 1);
911
912 // Phase 4: intern `s_small`, which should not fit in the leftover reclaimed entry from `s_large` _or_ the
913 // available capacity.
914 let s4 = intern_for_shard(&shard, s_small);
915 assert_eq!(s4, None);
916
917 assert_eq!(shard_entries(&shard), 2);
918 assert_eq!(shard_reclaimed_len(&shard), 1);
919
920 // And make sure we can still dereference the interned strings we _do_ have left:
921 assert_eq!(s2.deref(), s_medium1);
922 assert_eq!(s3.deref(), s_medium2);
923 }
924
925 #[test]
926 fn has_reclaimed_entries_string_fits_exactly() {
927 // The shard is large enough to fit two of the identically-sized strings, but not all three. We show that when we drop
928 // an entry and it is reclaimed, a string of identical size should always be able to reuse that reclaimed entry.
929 const S1_VALUE: &str = "hello world!";
930 const S2_VALUE: &str = "hello, world";
931 const S3_VALUE: &str = "hello--world";
932
933 let shard = create_shard(NonZeroUsize::new(80).unwrap());
934
935 // Intern the first two strings, which should fit without issue.
936 let s1 = intern_for_shard(&shard, S1_VALUE).expect("should not fail to intern");
937 let s1_reclaimed_expected = get_reclaimed_entry_for_string(&s1);
938 let _s2 = intern_for_shard(&shard, S2_VALUE).expect("should not fail to intern");
939
940 assert_eq!(shard_entries(&shard), 2);
941 assert_eq!(shard_reclaimed_len(&shard), 0);
942
943 // Try to intern the third string, which should fail as we don't have the space.
944 let s3 = intern_for_shard(&shard, S3_VALUE);
945 assert_eq!(s3, None);
946
947 // Drop the first string, which should decrement the reference count to zero and then reclaim the entry.
948 drop(s1);
949
950 assert_eq!(shard_entries(&shard), 1);
951 assert_eq!(shard_reclaimed_len(&shard), 1);
952
953 let s1_reclaimed = shard_first_reclaimed_entry(&shard);
954 assert_eq!(s1_reclaimed_expected, s1_reclaimed);
955
956 // Try again to intern the third string, which should now succeed and take over the reclaimed entry entirely
957 // as the strings are identical in length.
958 let _s3 = intern_for_shard(&shard, S3_VALUE).expect("should not fail to intern");
959
960 assert_eq!(shard_entries(&shard), 2);
961 assert_eq!(shard_reclaimed_len(&shard), 0);
962 }
963
964 #[test]
965 fn reclaimed_entry_reuse_split_too_small() {
966 // This situation is slightly contrived, but: we want to test that when there's a reclaimed entry of a certain
967 // size, reusing that reclaimed entry won't lead to it being split if the resulting split entry would be too
968 // "small": unable to hold another minimum-sized entry.
969 //
970 // We have to intern three strings to do this because we only track reclaimed entries when they're followed by
971 // in-use entries, and the string we drop to create a reclaimed entry has to be big enough, but not _too_ big,
972 // to hold the string we want to intern after it.
973 let shard = create_shard(NonZeroUsize::new(128).unwrap());
974
975 // Declare our strings to intern and just check some preconditions by hand.
976 let s_one = "a horse, a horse, my kingdom for a horse!";
977 let s_one_entry_len = entry_len(s_one);
978 let s_two = "why hello there, beautiful";
979 let s_two_entry_len = entry_len(s_two);
980 let s_three = "real gs move in silence like lasagna";
981 let s_three_entry_len = entry_len(s_three);
982
983 assert!(s_one_entry_len <= shard_capacity(&shard));
984 assert!(s_two_entry_len <= shard_capacity(&shard));
985 assert!(s_one_entry_len + s_two_entry_len + s_three_entry_len > shard_capacity(&shard));
986 assert!(s_one_entry_len > s_two_entry_len);
987 assert!(s_three_entry_len > s_two_entry_len);
988 assert!((s_one_entry_len - s_three_entry_len) < MINIMUM_ENTRY_LEN);
989
990 // Intern the first two strings, which should fit without issue.
991 let s1 = intern_for_shard(&shard, s_one).expect("should not fail to intern");
992 let s1_reclaimed_expected = get_reclaimed_entry_for_string(&s1);
993 let _s2 = intern_for_shard(&shard, s_two).expect("should not fail to intern");
994
995 assert_eq!(shard_entries(&shard), 2);
996 assert_eq!(shard_reclaimed_len(&shard), 0);
997
998 // Try to intern the third string, which should fail as we don't have the space.
999 let s3 = intern_for_shard(&shard, s_three);
1000 assert_eq!(s3, None);
1001
1002 // Drop the first string, which should decrement the reference count to zero and then reclaim the entry.
1003 drop(s1);
1004
1005 assert_eq!(shard_entries(&shard), 1);
1006 assert_eq!(shard_reclaimed_len(&shard), 1);
1007
1008 let s1_reclaimed = shard_first_reclaimed_entry(&shard);
1009 assert_eq!(s1_reclaimed_expected, s1_reclaimed);
1010
1011 // Try again to intern the third string, which should now succeed and take over the reclaimed entry, but since
1012 // the remainder of the reclaimed entry after taking the necessary capacity for `s_three` is not large enough
1013 // (`MINIMUM_ENTRY_LEN`), we shouldn't end up splitting the reclaimed entry, and instead, `s3` should consume
1014 // the entire reclaimed entry.
1015 let s3 = intern_for_shard(&shard, s_three).expect("should not fail to intern");
1016 let s3_reclaimed_expected = get_reclaimed_entry_for_string(&s3);
1017
1018 assert_eq!(shard_entries(&shard), 2);
1019 assert_eq!(shard_reclaimed_len(&shard), 0);
1020 assert_eq!(s1_reclaimed_expected, s3_reclaimed_expected);
1021 }
1022
1023 #[test]
1024 fn reclaimed_entry_adjacent_to_spare_capacity() {
1025 let shard = create_shard(NonZeroUsize::new(128).unwrap());
1026
1027 // Intern two smallish strings that fit without issue.
1028 let s1 = intern_for_shard(&shard, "hello, world!").expect("should not fail to intern");
1029 let s2 = intern_for_shard(&shard, "cheeeeeehooooo!").expect("should not fail to intern");
1030 let s1_entry_len = entry_len(&s1);
1031 let s2_entry_len = entry_len(&s2);
1032
1033 assert_eq!(shard_reclaimed_len(&shard), 0);
1034 assert_eq!(shard_available(&shard), 128 - s1_entry_len - s2_entry_len);
1035
1036 drop(s2);
1037 assert_eq!(shard_reclaimed_len(&shard), 0);
1038 assert_eq!(shard_available(&shard), 128 - s1_entry_len);
1039
1040 drop(s1);
1041 assert_eq!(shard_reclaimed_len(&shard), 0);
1042 assert_eq!(shard_available(&shard), 128);
1043 }
1044
1045 proptest! {
1046 #[test]
1047 #[cfg_attr(miri, ignore)]
1048 fn property_test_entry_count_accurate(
1049 strs in arb_alphanum_strings(1..=128, 16..=512),
1050 indices in arb_vec(any::<Index>(), 1..=1000),
1051 ) {
1052 // We ask `proptest` to generate a bunch of unique strings of varying lengths (1-128 bytes, 16-512 unique
1053 // strings) which we then randomly select out of those strings which strings we want to intern. The goal
1054 // here is to potentially select the same string multiple times, to exercise the actual interning logic...
1055 // but practically, to ensure that when we intern a string that has already been interned, we're not
1056 // incrementing the entries count again.
1057
1058 // Create an interner with enough capacity to hold all of the strings we've generated. This is the maximum
1059 // string size multiplied by the number of strings we've generated... plus a little constant factor per
1060 // string to account for the entry header.
1061 const ENTRY_SIZE: usize = 128 + HEADER_LEN;
1062 let interner = FixedSizeInterner::<1>::new(NonZeroUsize::new(ENTRY_SIZE * indices.len()).unwrap());
1063
1064 // For each index, pull out the string and both track it in `unique_strs` and intern it. We hold on to the
1065 // interned string handle to make sure the interned string is actually kept alive, keeping the entry count
1066 // stable.
1067 let mut interned = Vec::new();
1068 let mut unique_strs = HashSet::new();
1069 for index in &indices {
1070 let s = index.get(&strs);
1071 unique_strs.insert(s);
1072
1073 let s_interned = interner.try_intern(s).expect("should never fail to intern");
1074 interned.push(s_interned);
1075 }
1076
1077 assert_eq!(unique_strs.len(), interner.len());
1078 }
1079 }
1080}
1081
1082#[cfg(all(test, feature = "loom"))]
1083mod loom_tests {
1084 use std::{num::NonZeroUsize, ops::Deref};
1085
1086 use loom::sync::{Arc, Mutex};
1087
1088 use super::{
1089 tests::{create_shard, get_reclaimed_entry_for_string, intern_for_shard},
1090 *,
1091 };
1092
1093 #[test]
1094 fn concurrent_drop_and_intern() {
1095 fn shard_reclaimed_entries(shard: &Arc<Mutex<InternerShardState>>) -> Vec<ReclaimedEntry> {
1096 shard.lock().unwrap().reclaimed.iter().copied().collect()
1097 }
1098
1099 fn do_reclaimed_entries_overlap(a: ReclaimedEntry, b: ReclaimedEntry) -> bool {
1100 let a_start = a.offset;
1101 let a_end = a.offset + a.capacity - 1;
1102
1103 let b_start = b.offset;
1104 let b_end = b.offset + b.capacity - 1;
1105
1106 (a_start <= b_start && b_start <= a_end) || (a_start <= b_end && b_end <= a_end)
1107 }
1108
1109 const STRING_TO_INTERN: &str = "hello, world!";
1110
1111 // This test is meant to explore the thread orderings when one thread is trying to drop (and thus reclaim) the
1112 // last active reference to an interned string, and another thread is trying to intern that very same string.
1113 //
1114 // We accept, as a caveat, that a possible outcome is that we intern the "new" string again, even though an
1115 // existing entry to that string may have existed in an alternative ordering.
1116 loom::model(|| {
1117 let shard = create_shard(NonZeroUsize::new(1024).unwrap());
1118 let t2_shard = Arc::clone(&shard);
1119
1120 // Intern the string from thread T1.
1121 let t1_interned_s = intern_for_shard(&shard, STRING_TO_INTERN).expect("should not fail to intern");
1122 assert_eq!(t1_interned_s.deref(), STRING_TO_INTERN);
1123 let t1_reclaimed_entry = get_reclaimed_entry_for_string(&t1_interned_s);
1124
1125 // Spawn thread T2, which tries to intern the same string and returns the handle to it.
1126 let t2_result = loom::thread::spawn(move || {
1127 let interned_s = intern_for_shard(&t2_shard, STRING_TO_INTERN).expect("should not fail to intern");
1128 let reclaimed_entry = get_reclaimed_entry_for_string(&interned_s);
1129
1130 (interned_s, reclaimed_entry)
1131 });
1132
1133 drop(t1_interned_s);
1134
1135 let (t2_interned_s, t2_reclaimed_entry) = t2_result.join().expect("should not fail to join T2");
1136 assert_eq!(t2_interned_s.deref(), STRING_TO_INTERN);
1137
1138 // What we're checking for here is that either:
1139 // - there's no reclaimed entries (T2 found the existing entry for the string before T1 dropped it)
1140 // - there's a reclaimed entry (T2 didn't find the existing entry for the string before T1 marked it as
1141 // inactive) but the reclaimed entry does _not_ overlap with the interned string from T2, meaning we
1142 // didn't get confused and allow T2 to use an existing entry that T1 then later marked as reclaimed
1143 let reclaimed_entries = shard_reclaimed_entries(&shard);
1144 assert!(reclaimed_entries.len() <= 1, "should have at most one reclaimed entry");
1145
1146 if !reclaimed_entries.is_empty() {
1147 // If we do have a reclaimed entry, it needs to match exactly with only one of the interned strings.
1148 let is_t1_entry = reclaimed_entries.first().unwrap() == &t1_reclaimed_entry;
1149 let is_t2_entry = reclaimed_entries.first().unwrap() == &t2_reclaimed_entry;
1150
1151 assert!(
1152 (is_t1_entry || is_t2_entry) && !(is_t1_entry && is_t2_entry),
1153 "should only match one interned string"
1154 );
1155
1156 // Additionally, we ensure that the reclaimed entry does not overlap with the other interned string.
1157 assert!(
1158 !do_reclaimed_entries_overlap(t1_reclaimed_entry, t2_reclaimed_entry),
1159 "reclaimed entry should not overlap with remaining interned string"
1160 );
1161 }
1162 });
1163 }
1164
1165 #[test]
1166 fn concurrent_resurrect_and_double_drop() {
1167 // Regression test for a double-reclamation race, mirroring the `GenericMapInterner` test of the same name.
1168 //
1169 // Unlike `concurrent_drop_and_intern` -- where the second thread keeps its handle alive -- here the second
1170 // thread interns the string _and_ drops it within the race window. This models the dangerous interleaving:
1171 //
1172 // 1. T1 drops the last reference: `refs` 1 -> 0 (but T1 hasn't taken the shard lock yet).
1173 // 2. T2 calls `try_intern`; `find_entry` observes the entry as active and resurrects it: `refs` 0 -> 1.
1174 // 3. T2 drops its handle: `refs` 1 -> 0.
1175 // 4. T1 and T2 each take the lock and observe `refs == 0`, so BOTH mark the same slot for reclamation.
1176 //
1177 // A double reclamation either underflows `entries`/`len` (a panic under overflow-checks, which the loom test
1178 // profile enables) or inserts a spurious/overlapping reclaimed entry. After every handle to the single interned
1179 // string has been dropped, the only correct end state is a fully empty shard.
1180 const STRING_TO_INTERN: &str = "hello, world!";
1181
1182 loom::model(|| {
1183 let shard = create_shard(NonZeroUsize::new(1024).unwrap());
1184 let t2_shard = Arc::clone(&shard);
1185
1186 // T1 interns the string and holds the only reference.
1187 let t1_interned_s = intern_for_shard(&shard, STRING_TO_INTERN).expect("should not fail to intern");
1188
1189 // T2 interns the same string and immediately drops it. The interesting orderings are those where this runs
1190 // while T1's drop has decremented the refcount to zero but not yet taken the shard lock.
1191 let t2 = loom::thread::spawn(move || {
1192 let t2_interned_s = intern_for_shard(&t2_shard, STRING_TO_INTERN).expect("should not fail to intern");
1193 assert_eq!(t2_interned_s.deref(), STRING_TO_INTERN);
1194 drop(t2_interned_s);
1195 });
1196
1197 // T1 drops its (originally last) reference.
1198 drop(t1_interned_s);
1199
1200 t2.join().expect("should not fail to join T2");
1201
1202 // Every handle has been dropped, so the slot must have been reclaimed exactly once: the shard is fully
1203 // empty, with consistent accounting and no leftover reclaimed entries. A double reclaim breaks at least one
1204 // of these (and typically panics first via the `entries`/`len` underflow).
1205 let shard = shard.lock().unwrap();
1206 assert_eq!(
1207 shard.entries, 0,
1208 "entries should be zero (it is double-subtracted on a double reclaim)"
1209 );
1210 assert_eq!(
1211 shard.len, 0,
1212 "len should be zero (it is double-subtracted on a double reclaim)"
1213 );
1214 assert_eq!(shard.offset, 0, "offset should wind back to zero");
1215 assert!(
1216 shard.reclaimed.is_empty(),
1217 "no reclaimed entries should remain (a double reclaim inserts a spurious/overlapping entry)"
1218 );
1219 });
1220 }
1221}