ffi_support/
handle_map.rs

1/* Copyright 2018-2019 Mozilla Foundation
2 *
3 * Licensed under the Apache License (Version 2.0), or the MIT license,
4 * (the "Licenses") at your option. You may not use this file except in
5 * compliance with one of the Licenses. You may obtain copies of the
6 * Licenses at:
7 *
8 *    http://www.apache.org/licenses/LICENSE-2.0
9 *    http://opensource.org/licenses/MIT
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the Licenses is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the Licenses for the specific language governing permissions and
15 * limitations under the Licenses. */
16
17//! This module provides a [`Handle`] type, which you can think of something
18//! like a dynamically checked, type erased reference/pointer type. Depending on
19//! the usage pattern a handle can behave as either a borrowed reference, or an
20//! owned pointer.
21//!
22//! They can be losslessly converted [to](Handle::into_u64) and
23//! [from](Handle::from_u64) a 64 bit integer, for ease of passing over the FFI
24//! (and they implement [`IntoFfi`] using these primitives for this purpose).
25//!
26//! The benefit is primarially that they can detect common misuse patterns that
27//! would otherwise be silent bugs, such as use-after-free, double-free, passing
28//! a wrongly-typed pointer to a function, etc.
29//!
30//! Handles are provided when inserting an item into either a [`HandleMap`] or a
31//! [`ConcurrentHandleMap`].
32//!
33//! # Comparison to types from other crates
34//!
35//! [`HandleMap`] is similar to types offered by other crates, such as
36//! `slotmap`, or `slab`. However, it has a number of key differences which make
37//! it better for our purposes as compared to the types in those crates:
38//!
39//! 1. Unlike `slab` (but like `slotmap`), we implement versioning, detecting
40//!    ABA problems, which allows us to detect use after free.
41//! 2. Unlike `slotmap`, we don't have the `T: Copy` restriction.
42//! 3. Unlike either, we can detect when you use a Key in a map that did not
43//!    allocate the key. This is true even when the map is from a `.so` file
44//!    compiled separately.
45//! 3. Our implementation of doesn't use any `unsafe` (at the time of this
46//!    writing).
47//!
48//! However, it comes with the following drawbacks:
49//!
50//! 1. `slotmap` holds its version information in a `u32`, and so it takes
51//!    2<sup>31</sup> colliding insertions and deletions before it could
52//!    potentially fail to detect an ABA issue, wheras we use a `u16`, and are
53//!    limited to 2<sup>15</sup>.
54//! 2. Similarly, we can only hold 2<sup>16</sup> items at once, unlike
55//!    `slotmap`'s 2<sup>32</sup>. (Considering these items are typically things
56//!    like database handles, this is probably plenty).
57//! 3. Our implementation is slower, and uses slightly more memory than
58//!    `slotmap` (which is in part due to the lack of `unsafe` mentioned above)
59//!
60//! The first two issues seem exceptionally unlikely, even for extremely
61//! long-lived `HandleMap`, and we're still memory safe even if they occur (we
62//! just might fail to notice a bug). The third issue also seems unimportant for
63//! our use case.
64
65use crate::error::{ErrorCode, ExternError};
66use crate::into_ffi::IntoFfi;
67use std::error::Error as StdError;
68use std::fmt;
69use std::ops;
70use std::sync::atomic::{AtomicUsize, Ordering};
71use std::sync::{Mutex, RwLock};
72
73/// `HandleMap` is a collection type which can hold any type of value, and
74/// offers a stable handle which can be used to retrieve it on insertion. These
75/// handles offer methods for converting [to](Handle::into_u64) and
76/// [from](Handle::from_u64) 64 bit integers, meaning they're very easy to pass
77/// over the FFI (they also implement [`IntoFfi`] for the same purpose).
78///
79/// See the [module level docs](index.html) for more information.
80///
81/// Note: In FFI code, most usage of `HandleMap` will be done through the
82/// [`ConcurrentHandleMap`] type, which is a thin wrapper around a
83/// `RwLock<HandleMap<Mutex<T>>>`.
84#[derive(Debug, Clone)]
85pub struct HandleMap<T> {
86    // The value of `map_id` in each `Handle`.
87    id: u16,
88
89    // Index to the start of the free list. Always points to a free item --
90    // we never allow our free list to become empty.
91    first_free: u16,
92
93    // The number of entries with `data.is_some()`. This is never equal to
94    // `entries.len()`, we always grow before that point to ensure we always have
95    // a valid `first_free` index to add entries onto. This is our `len`.
96    num_entries: usize,
97
98    // The actual data. Note: entries.len() is our 'capacity'.
99    entries: Vec<Entry<T>>,
100}
101
102#[derive(Debug, Clone)]
103struct Entry<T> {
104    // initially 1, incremented on insertion and removal. Thus,
105    // if version is even, state should always be EntryState::Active.
106    version: u16,
107    state: EntryState<T>,
108}
109
110#[derive(Debug, Clone)]
111enum EntryState<T> {
112    // Not part of the free list
113    Active(T),
114    // The u16 is the next index in the free list.
115    InFreeList(u16),
116    // Part of the free list, but the sentinel.
117    EndOfFreeList,
118}
119
120impl<T> EntryState<T> {
121    #[cfg(any(debug_assertions, test))]
122    fn is_end_of_list(&self) -> bool {
123        match self {
124            EntryState::EndOfFreeList => true,
125            _ => false,
126        }
127    }
128
129    #[inline]
130    fn is_occupied(&self) -> bool {
131        self.get_item().is_some()
132    }
133
134    #[inline]
135    fn get_item(&self) -> Option<&T> {
136        match self {
137            EntryState::Active(v) => Some(v),
138            _ => None,
139        }
140    }
141
142    #[inline]
143    fn get_item_mut(&mut self) -> Option<&mut T> {
144        match self {
145            EntryState::Active(v) => Some(v),
146            _ => None,
147        }
148    }
149}
150
151// Small helper to check our casts.
152#[inline]
153fn to_u16(v: usize) -> u16 {
154    use std::u16::MAX as U16_MAX;
155    // Shouldn't ever happen.
156    assert!(v <= (U16_MAX as usize), "Bug: Doesn't fit in u16: {}", v);
157    v as u16
158}
159
160/// The maximum capacity of a [`HandleMap`]. Attempting to instantiate one with
161/// a larger capacity will cause a panic.
162///
163/// Note: This could go as high as `(1 << 16) - 2`, but doing is seems more
164/// error prone. For the sake of paranoia, we limit it to this size, which is
165/// already quite a bit larger than it seems like we're likely to ever need.
166pub const MAX_CAPACITY: usize = (1 << 15) - 1;
167
168// Never having to worry about capacity == 0 simplifies the code at the cost of
169// worse memory usage. It doesn't seem like there's any reason to make this
170// public.
171const MIN_CAPACITY: usize = 4;
172
173/// An error representing the ways a `Handle` may be invalid.
174#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
175pub enum HandleError {
176    /// Identical to invalid handle, but has a slightly more helpful
177    /// message for the most common case 0.
178    NullHandle,
179
180    /// Returned from [`Handle::from_u64`] if [`Handle::is_valid`] fails.
181    InvalidHandle,
182
183    /// Returned from get/get_mut/delete if the handle is stale (this indicates
184    /// something equivalent to a use-after-free / double-free, etc).
185    StaleVersion,
186
187    /// Returned if the handle index references an index past the end of the
188    /// HandleMap.
189    IndexPastEnd,
190
191    /// The handle has a map_id for a different map than the one it was
192    /// attempted to be used with.
193    WrongMap,
194}
195
196impl StdError for HandleError {}
197
198impl fmt::Display for HandleError {
199    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
200        use HandleError::*;
201        match self {
202            NullHandle => {
203                f.write_str("Tried to use a null handle (this object has probably been closed)")
204            }
205            InvalidHandle => f.write_str("u64 could not encode a valid Handle"),
206            StaleVersion => f.write_str("Handle has stale version number"),
207            IndexPastEnd => f.write_str("Handle references a index past the end of this HandleMap"),
208            WrongMap => f.write_str("Handle is from a different map"),
209        }
210    }
211}
212
213impl From<HandleError> for ExternError {
214    fn from(e: HandleError) -> Self {
215        ExternError::new_error(ErrorCode::INVALID_HANDLE, e.to_string())
216    }
217}
218
219impl<T> HandleMap<T> {
220    /// Create a new `HandleMap` with the default capacity.
221    pub fn new() -> Self {
222        Self::new_with_capacity(MIN_CAPACITY)
223    }
224
225    /// Allocate a new `HandleMap`. Note that the actual capacity may be larger
226    /// than the requested value.
227    ///
228    /// Panics if `request` is greater than [`handle_map::MAX_CAPACITY`](MAX_CAPACITY)
229    pub fn new_with_capacity(request: usize) -> Self {
230        assert!(
231            request <= MAX_CAPACITY,
232            "HandleMap capacity is limited to {} (request was {})",
233            MAX_CAPACITY,
234            request
235        );
236
237        let capacity = request.max(MIN_CAPACITY);
238        let id = next_handle_map_id();
239        let mut entries = Vec::with_capacity(capacity);
240
241        // Initialize each entry with version 1, and as a member of the free list
242        for i in 0..(capacity - 1) {
243            entries.push(Entry {
244                version: 1,
245                state: EntryState::InFreeList(to_u16(i + 1)),
246            });
247        }
248
249        // And the final entry is at the end of the free list
250        // (but still has version 1).
251        entries.push(Entry {
252            version: 1,
253            state: EntryState::EndOfFreeList,
254        });
255        Self {
256            id,
257            first_free: 0,
258            num_entries: 0,
259            entries,
260        }
261    }
262
263    /// Get the number of entries in the `HandleMap`.
264    #[inline]
265    pub fn len(&self) -> usize {
266        self.num_entries
267    }
268
269    /// Returns true if the HandleMap is empty.
270    #[inline]
271    pub fn is_empty(&self) -> bool {
272        self.len() == 0
273    }
274
275    /// Returns the number of slots allocated in the handle map.
276    #[inline]
277    pub fn capacity(&self) -> usize {
278        // It's not a bug that this isn't entries.capacity() -- We're returning
279        // how many slots exist, not something about the backing memory allocation
280        self.entries.len()
281    }
282
283    fn ensure_capacity(&mut self, cap_at_least: usize) {
284        assert_ne!(self.len(), self.capacity(), "Bug: should have grown by now");
285        assert!(cap_at_least <= MAX_CAPACITY, "HandleMap overfilled");
286        if self.capacity() > cap_at_least {
287            return;
288        }
289
290        let mut next_cap = self.capacity();
291        while next_cap <= cap_at_least {
292            next_cap *= 2;
293        }
294        next_cap = next_cap.min(MAX_CAPACITY);
295
296        let need_extra = next_cap.saturating_sub(self.entries.capacity());
297        self.entries.reserve(need_extra);
298
299        assert!(
300            !self.entries[self.first_free as usize].state.is_occupied(),
301            "Bug: HandleMap.first_free points at occupied index"
302        );
303
304        // Insert new entries at the front of our list.
305        while self.entries.len() < next_cap - 1 {
306            // This is a little wasteful but whatever. Add each new entry to the
307            // front of the free list one at a time.
308            self.entries.push(Entry {
309                version: 1,
310                state: EntryState::InFreeList(self.first_free),
311            });
312            self.first_free = to_u16(self.entries.len() - 1);
313        }
314
315        self.debug_check_valid();
316    }
317
318    #[inline]
319    fn debug_check_valid(&self) {
320        // Run the expensive validity check in tests and in debug builds.
321        #[cfg(any(debug_assertions, test))]
322        {
323            self.assert_valid();
324        }
325    }
326
327    #[cfg(any(debug_assertions, test))]
328    fn assert_valid(&self) {
329        assert_ne!(self.len(), self.capacity());
330        assert!(self.capacity() <= MAX_CAPACITY, "Entries too large");
331        // Validate that our free list is correct.
332
333        let number_of_ends = self
334            .entries
335            .iter()
336            .filter(|e| e.state.is_end_of_list())
337            .count();
338        assert_eq!(
339            number_of_ends, 1,
340            "More than one entry think's it's the end of the list, or no entries do"
341        );
342
343        // Check that the free list hits every unoccupied item.
344        // The tuple is: `(should_be_in_free_list, is_in_free_list)`.
345        let mut free_indices = vec![(false, false); self.capacity()];
346        for (i, e) in self.entries.iter().enumerate() {
347            if !e.state.is_occupied() {
348                free_indices[i].0 = true;
349            }
350        }
351
352        let mut next = self.first_free;
353        loop {
354            let ni = next as usize;
355
356            assert!(
357                ni <= free_indices.len(),
358                "Free list contains out of bounds index!"
359            );
360
361            assert!(
362                free_indices[ni].0,
363                "Free list has an index that shouldn't be free! {}",
364                ni
365            );
366
367            assert!(
368                !free_indices[ni].1,
369                "Free list hit an index ({}) more than once! Cycle detected!",
370                ni
371            );
372
373            free_indices[ni].1 = true;
374
375            match &self.entries[ni].state {
376                EntryState::InFreeList(next_index) => next = *next_index,
377                EntryState::EndOfFreeList => break,
378                // Hitting `Active` here is probably not possible because of the checks above, but who knows.
379                EntryState::Active(..) => unreachable!("Bug: Active item in free list at {}", next),
380            }
381        }
382        let mut occupied_count = 0;
383        for (i, &(should_be_free, is_free)) in free_indices.iter().enumerate() {
384            assert_eq!(
385                should_be_free, is_free,
386                "Free list missed item, or contains an item it shouldn't: {}",
387                i
388            );
389            if !should_be_free {
390                occupied_count += 1;
391            }
392        }
393        assert_eq!(
394            self.num_entries, occupied_count,
395            "num_entries doesn't reflect the actual number of entries"
396        );
397    }
398
399    /// Insert an item into the map, and return a handle to it.
400    pub fn insert(&mut self, v: T) -> Handle {
401        let need_cap = self.len() + 1;
402        self.ensure_capacity(need_cap);
403        let index = self.first_free;
404        let result = {
405            // Scoped mutable borrow of entry.
406            let entry = &mut self.entries[index as usize];
407            let new_first_free = match entry.state {
408                EntryState::InFreeList(i) => i,
409                _ => panic!("Bug: next_index pointed at non-free list entry (or end of list)"),
410            };
411            entry.version += 1;
412            if entry.version == 0 {
413                entry.version += 2;
414            }
415            entry.state = EntryState::Active(v);
416            self.first_free = new_first_free;
417            self.num_entries += 1;
418
419            Handle {
420                map_id: self.id,
421                version: entry.version,
422                index,
423            }
424        };
425        self.debug_check_valid();
426        result
427    }
428
429    // Helper to contain the handle validation boilerplate. Returns `h.index as usize`.
430    fn check_handle(&self, h: Handle) -> Result<usize, HandleError> {
431        if h.map_id != self.id {
432            log::info!(
433                "HandleMap access with handle having wrong map id: {:?} (our map id is {})",
434                h,
435                self.id
436            );
437            return Err(HandleError::WrongMap);
438        }
439        let index = h.index as usize;
440        if index >= self.entries.len() {
441            log::info!("HandleMap accessed with handle past end of map: {:?}", h);
442            return Err(HandleError::IndexPastEnd);
443        }
444        if self.entries[index].version != h.version {
445            log::info!(
446                "HandleMap accessed with handle with wrong version {:?} (entry version is {})",
447                h,
448                self.entries[index].version
449            );
450            return Err(HandleError::StaleVersion);
451        }
452        // At this point, we know the handle version matches the entry version,
453        // but if someone created a specially invalid handle, they could have
454        // its version match the version they expect an unoccupied index to
455        // have.
456        //
457        // We don't use any unsafe, so the worse thing that can happen here is
458        // that we get confused and panic, but still that's not great, so we
459        // check for this explicitly.
460        //
461        // Note that `active` versions are always even, as they start at 1, and
462        // are incremented on both insertion and deletion.
463        //
464        // Anyway, this is just for sanity checking, we already check this in
465        // practice when we convert `u64`s into `Handle`s, which is the only
466        // way we ever use these in the real world.
467        if (h.version % 2) != 0 {
468            log::info!(
469                "HandleMap given handle with matching but illegal version: {:?}",
470                h,
471            );
472            return Err(HandleError::StaleVersion);
473        }
474        Ok(index)
475    }
476
477    /// Delete an item from the HandleMap.
478    pub fn delete(&mut self, h: Handle) -> Result<(), HandleError> {
479        self.remove(h).map(drop)
480    }
481
482    /// Remove an item from the HandleMap, returning the old value.
483    pub fn remove(&mut self, h: Handle) -> Result<T, HandleError> {
484        let index = self.check_handle(h)?;
485        let prev = {
486            // Scoped mutable borrow of entry.
487            let entry = &mut self.entries[index];
488            entry.version += 1;
489            let index = h.index;
490            let last_state =
491                std::mem::replace(&mut entry.state, EntryState::InFreeList(self.first_free));
492            self.num_entries -= 1;
493            self.first_free = index;
494
495            if let EntryState::Active(value) = last_state {
496                value
497            } else {
498                // This indicates either a bug in HandleMap or memory
499                // corruption. Abandon all hope.
500                unreachable!(
501                    "Handle {:?} passed validation but references unoccupied entry",
502                    h
503                );
504            }
505        };
506        self.debug_check_valid();
507        Ok(prev)
508    }
509
510    /// Get a reference to the item referenced by the handle, or return a
511    /// [`HandleError`] describing the problem.
512    pub fn get(&self, h: Handle) -> Result<&T, HandleError> {
513        let idx = self.check_handle(h)?;
514        let entry = &self.entries[idx];
515        // This should be caught by check_handle above, but we avoid panicking
516        // because we'd rather not poison any locks we don't have to poison
517        let item = entry
518            .state
519            .get_item()
520            .ok_or_else(|| HandleError::InvalidHandle)?;
521        Ok(item)
522    }
523
524    /// Get a mut reference to the item referenced by the handle, or return a
525    /// [`HandleError`] describing the problem.
526    pub fn get_mut(&mut self, h: Handle) -> Result<&mut T, HandleError> {
527        let idx = self.check_handle(h)?;
528        let entry = &mut self.entries[idx];
529        // This should be caught by check_handle above, but we avoid panicking
530        // because we'd rather not poison any locks we don't have to poison
531        let item = entry
532            .state
533            .get_item_mut()
534            .ok_or_else(|| HandleError::InvalidHandle)?;
535        Ok(item)
536    }
537}
538
539impl<T> Default for HandleMap<T> {
540    #[inline]
541    fn default() -> Self {
542        HandleMap::new()
543    }
544}
545
546impl<T> ops::Index<Handle> for HandleMap<T> {
547    type Output = T;
548    #[inline]
549    fn index(&self, h: Handle) -> &T {
550        self.get(h)
551            .expect("Indexed into HandleMap with invalid handle!")
552    }
553}
554
555// We don't implement IndexMut intentionally (implementing ops::Index is
556// dubious enough)
557
558/// A Handle we allow to be returned over the FFI by implementing [`IntoFfi`].
559/// This type is intentionally not `#[repr(C)]`, and getting the data out of the
560/// FFI is done using `Handle::from_u64`, or it's implemetation of `From<u64>`.
561///
562/// It consists of, at a minimum:
563///
564/// - A "map id" (used to ensure you're using it with the correct map)
565/// - a "version" (incremented when the value in the index changes, used to
566///   detect multiple frees, use after free, and ABA and ABA)
567/// - and a field indicating which index it goes into.
568///
569/// In practice, it may also contain extra information to help detect other
570/// errors (currently it stores a "magic value" used to detect invalid
571/// [`Handle`]s).
572///
573/// These fields may change but the following guarantees are made about the
574/// internal representation:
575///
576/// - This will always be representable in 64 bits.
577/// - The bits, when interpreted as a signed 64 bit integer, will be positive
578///   (that is to say, it will *actually* be representable in 63 bits, since
579///   this makes the most significant bit unavailable for the purposes of
580///   encoding). This guarantee makes things slightly less dubious when passing
581///   things to Java, gives us some extra validation ability, etc.
582#[derive(Copy, Clone, Debug, PartialEq)]
583pub struct Handle {
584    map_id: u16,
585    version: u16,
586    index: u16,
587}
588
589// We stuff this into the top 16 bits of the handle when u16 encoded to detect
590// various sorts of weirdness. It's the letters 'A' and 'S' as ASCII, but the
591// only important thing about it is that the most significant bit be unset.
592const HANDLE_MAGIC: u16 = 0x4153_u16;
593
594impl Handle {
595    /// Convert a `Handle` to a `u64`. You can also use `Into::into` directly.
596    /// Most uses of this will be automatic due to our [`IntoFfi`] implementation.
597    #[inline]
598    pub fn into_u64(self) -> u64 {
599        let map_id = u64::from(self.map_id);
600        let version = u64::from(self.version);
601        let index = u64::from(self.index);
602        // SOMEDAY: we could also use this as a sort of CRC if we were really paranoid.
603        // e.g. `magic = combine_to_u16(map_id, version, index)`.
604        let magic = u64::from(HANDLE_MAGIC);
605        (magic << 48) | (map_id << 32) | (index << 16) | version
606    }
607
608    /// Convert a `u64` to a `Handle`. Inverse of `into_u64`. We also implement
609    /// `From::from` (which will panic instead of returning Err).
610    ///
611    /// Returns [`HandleError::InvalidHandle`](HandleError) if the bits cannot
612    /// possibly represent a valid handle.
613    pub fn from_u64(v: u64) -> Result<Self, HandleError> {
614        if !Handle::is_valid(v) {
615            log::warn!("Illegal handle! {:x}", v);
616            if v == 0 {
617                Err(HandleError::NullHandle)
618            } else {
619                Err(HandleError::InvalidHandle)
620            }
621        } else {
622            let map_id = (v >> 32) as u16;
623            let index = (v >> 16) as u16;
624            let version = v as u16;
625            Ok(Self {
626                map_id,
627                version,
628                index,
629            })
630        }
631    }
632
633    /// Returns whether or not `v` makes a bit pattern that could represent an
634    /// encoded [`Handle`].
635    #[inline]
636    pub fn is_valid(v: u64) -> bool {
637        (v >> 48) == u64::from(HANDLE_MAGIC) &&
638        // The "bottom" field is the version. We increment it both when
639        // inserting and removing, and they're all initially 1. So, all valid
640        // handles that we returned should have an even version.
641        ((v & 1) == 0)
642    }
643}
644
645impl From<u64> for Handle {
646    fn from(u: u64) -> Self {
647        Handle::from_u64(u).expect("Illegal handle!")
648    }
649}
650
651impl From<Handle> for u64 {
652    #[inline]
653    fn from(h: Handle) -> u64 {
654        h.into_u64()
655    }
656}
657
658unsafe impl IntoFfi for Handle {
659    type Value = u64;
660    // Note: intentionally does not encode a valid handle for any map.
661    #[inline]
662    fn ffi_default() -> u64 {
663        0u64
664    }
665    #[inline]
666    fn into_ffi_value(self) -> u64 {
667        self.into_u64()
668    }
669}
670
671/// `ConcurrentHandleMap` is a relatively thin wrapper around
672/// `RwLock<HandleMap<Mutex<T>>>`. Due to the nested locking, it's not possible
673/// to implement the same API as [`HandleMap`], however it does implement an API
674/// that offers equivalent functionality, as well as several functions that
675/// greatly simplify FFI usage (see example below).
676///
677/// See the [module level documentation](index.html) for more info.
678///
679/// # Example
680///
681/// ```rust,no_run
682/// # #[macro_use] extern crate lazy_static;
683/// # extern crate ffi_support;
684/// # use ffi_support::*;
685/// # use std::sync::*;
686///
687/// // Somewhere...
688/// struct Thing { value: f64 }
689///
690/// lazy_static! {
691///     static ref ITEMS: ConcurrentHandleMap<Thing> = ConcurrentHandleMap::new();
692/// }
693///
694/// #[no_mangle]
695/// pub extern "C" fn mylib_new_thing(value: f64, err: &mut ExternError) -> u64 {
696///     // Most uses will be `ITEMS.insert_with_result`. Note that this already
697///     // calls `call_with_output` (or `call_with_result` if this were
698///     // `insert_with_result`) for you.
699///     ITEMS.insert_with_output(err, || Thing { value })
700/// }
701///
702/// #[no_mangle]
703/// pub extern "C" fn mylib_thing_value(h: u64, err: &mut ExternError) -> f64 {
704///     // Or `ITEMS.call_with_result` for the fallible functions.
705///     ITEMS.call_with_output(err, h, |thing| thing.value)
706/// }
707///
708/// #[no_mangle]
709/// pub extern "C" fn mylib_thing_set_value(h: u64, new_value: f64, err: &mut ExternError) {
710///     ITEMS.call_with_output_mut(err, h, |thing| {
711///         thing.value = new_value;
712///     })
713/// }
714///
715/// // Note: defines the following function:
716/// // pub extern "C" fn mylib_destroy_thing(h: u64, err: &mut ExternError)
717/// define_handle_map_deleter!(ITEMS, mylib_destroy_thing);
718/// ```
719pub struct ConcurrentHandleMap<T> {
720    /// The underlying map. Public so that more advanced use-cases
721    /// may use it as they please.
722    pub map: RwLock<HandleMap<Mutex<T>>>,
723}
724
725impl<T> ConcurrentHandleMap<T> {
726    /// Construct a new `ConcurrentHandleMap`.
727    pub fn new() -> Self {
728        Self {
729            map: RwLock::new(HandleMap::new()),
730        }
731    }
732
733    /// Get the number of entries in the `ConcurrentHandleMap`.
734    ///
735    /// This takes the map's `read` lock.
736    #[inline]
737    pub fn len(&self) -> usize {
738        let map = self.map.read().unwrap();
739        map.len()
740    }
741
742    /// Returns true if the `ConcurrentHandleMap` is empty.
743    ///
744    /// This takes the map's `read` lock.
745    #[inline]
746    pub fn is_empty(&self) -> bool {
747        self.len() == 0
748    }
749
750    /// Insert an item into the map, returning the newly allocated handle to the
751    /// item.
752    ///
753    /// # Locking
754    ///
755    /// Note that this requires taking the map's write lock, and so it will
756    /// block until all other threads have finished any read/write operations.
757    pub fn insert(&self, v: T) -> Handle {
758        // Fails if the lock is poisoned. Not clear what we should do here... We
759        // could always insert anyway (by matching on LockResult), but that
760        // seems... really quite dubious.
761        let mut map = self.map.write().unwrap();
762        map.insert(Mutex::new(v))
763    }
764
765    /// Remove an item from the map.
766    ///
767    /// # Locking
768    ///
769    /// Note that this requires taking the map's write lock, and so it will
770    /// block until all other threads have finished any read/write operations.
771    pub fn delete(&self, h: Handle) -> Result<(), HandleError> {
772        // We use `remove` and not delete (and use the inner block) to ensure
773        // that if `v`'s destructor panics, we aren't holding the write lock
774        // when it happens, so that the map itself doesn't get poisoned.
775        let v = {
776            let mut map = self.map.write().unwrap();
777            map.remove(h)
778        };
779        v.map(drop)
780    }
781
782    /// Convenient wrapper for `delete` which takes a `u64` that it will
783    /// convert to a handle.
784    ///
785    /// The main benefit (besides convenience) of this over the version
786    /// that takes a [`Handle`] is that it allows handling handle-related errors
787    /// in one place.
788    pub fn delete_u64(&self, h: u64) -> Result<(), HandleError> {
789        self.delete(Handle::from_u64(h)?)
790    }
791
792    /// Remove an item from the map, returning either the item,
793    /// or None if its guard mutex got poisoned at some point.
794    ///
795    /// # Locking
796    ///
797    /// Note that this requires taking the map's write lock, and so it will
798    /// block until all other threads have finished any read/write operations.
799    pub fn remove(&self, h: Handle) -> Result<Option<T>, HandleError> {
800        let mut map = self.map.write().unwrap();
801        let mutex = map.remove(h)?;
802        Ok(mutex.into_inner().ok())
803    }
804
805    /// Convenient wrapper for `remove` which takes a `u64` that it will
806    /// convert to a handle.
807    ///
808    /// The main benefit (besides convenience) of this over the version
809    /// that takes a [`Handle`] is that it allows handling handle-related errors
810    /// in one place.
811    pub fn remove_u64(&self, h: u64) -> Result<Option<T>, HandleError> {
812        self.remove(Handle::from_u64(h)?)
813    }
814
815    /// Call `callback` with a non-mutable reference to the item from the map,
816    /// after acquiring the necessary locks.
817    ///
818    /// # Locking
819    ///
820    /// Note that this requires taking both:
821    ///
822    /// - The map's read lock, and so it will block until all other threads have
823    ///   finished any write operations.
824    /// - The mutex on the slot the handle is mapped to.
825    ///
826    /// And so it will block if there are ongoing write operations, or if
827    /// another thread is reading from the same handle.
828    ///
829    /// # Panics
830    ///
831    /// This will panic if a previous `get()` or `get_mut()` call has panicked
832    /// inside it's callback. The solution to this
833    ///
834    /// (It may also panic if the handle map detects internal state corruption,
835    /// however this should not happen except for bugs in the handle map code).
836    pub fn get<F, E, R>(&self, h: Handle, callback: F) -> Result<R, E>
837    where
838        F: FnOnce(&T) -> Result<R, E>,
839        E: From<HandleError>,
840    {
841        self.get_mut(h, |v| callback(v))
842    }
843
844    /// Call `callback` with a mutable reference to the item from the map, after
845    /// acquiring the necessary locks.
846    ///
847    /// # Locking
848    ///
849    /// Note that this requires taking both:
850    ///
851    /// - The map's read lock, and so it will block until all other threads have
852    ///   finished any write operations.
853    /// - The mutex on the slot the handle is mapped to.
854    ///
855    /// And so it will block if there are ongoing write operations, or if
856    /// another thread is reading from the same handle.
857    ///
858    /// # Panics
859    ///
860    /// This will panic if a previous `get()` or `get_mut()` call has panicked
861    /// inside it's callback. The only solution to this is to remove and reinsert
862    /// said item.
863    ///
864    /// (It may also panic if the handle map detects internal state corruption,
865    /// however this should not happen except for bugs in the handle map code).
866    pub fn get_mut<F, E, R>(&self, h: Handle, callback: F) -> Result<R, E>
867    where
868        F: FnOnce(&mut T) -> Result<R, E>,
869        E: From<HandleError>,
870    {
871        // XXX figure out how to handle poison...
872        let map = self.map.read().unwrap();
873        let mtx = map.get(h)?;
874        let mut hm = mtx.lock().unwrap();
875        callback(&mut *hm)
876    }
877
878    /// Convenient wrapper for `get` which takes a `u64` that it will convert to
879    /// a handle.
880    ///
881    /// The other benefit (besides convenience) of this over the version
882    /// that takes a [`Handle`] is that it allows handling handle-related errors
883    /// in one place.
884    ///
885    /// # Locking
886    ///
887    /// Note that this requires taking both:
888    ///
889    /// - The map's read lock, and so it will block until all other threads have
890    ///   finished any write operations.
891    /// - The mutex on the slot the handle is mapped to.
892    ///
893    /// And so it will block if there are ongoing write operations, or if
894    /// another thread is reading from the same handle.
895    pub fn get_u64<F, E, R>(&self, u: u64, callback: F) -> Result<R, E>
896    where
897        F: FnOnce(&T) -> Result<R, E>,
898        E: From<HandleError>,
899    {
900        self.get(Handle::from_u64(u)?, callback)
901    }
902
903    /// Convenient wrapper for [`Self::get_mut`] which takes a `u64` that it will
904    /// convert to a handle.
905    ///
906    /// The main benefit (besides convenience) of this over the version
907    /// that takes a [`Handle`] is that it allows handling handle-related errors
908    /// in one place.
909    ///
910    /// # Locking
911    ///
912    /// Note that this requires taking both:
913    ///
914    /// - The map's read lock, and so it will block until all other threads have
915    ///   finished any write operations.
916    /// - The mutex on the slot the handle is mapped to.
917    ///
918    /// And so it will block if there are ongoing write operations, or if
919    /// another thread is reading from the same handle.
920    pub fn get_mut_u64<F, E, R>(&self, u: u64, callback: F) -> Result<R, E>
921    where
922        F: FnOnce(&mut T) -> Result<R, E>,
923        E: From<HandleError>,
924    {
925        self.get_mut(Handle::from_u64(u)?, callback)
926    }
927
928    /// Helper that performs both a
929    /// [`call_with_result`][crate::call_with_result] and
930    /// [`get`](ConcurrentHandleMap::get_mut).
931    pub fn call_with_result_mut<R, E, F>(
932        &self,
933        out_error: &mut ExternError,
934        h: u64,
935        callback: F,
936    ) -> R::Value
937    where
938        F: std::panic::UnwindSafe + FnOnce(&mut T) -> Result<R, E>,
939        ExternError: From<E>,
940        R: IntoFfi,
941    {
942        use crate::call_with_result;
943        call_with_result(out_error, || -> Result<_, ExternError> {
944            // We can't reuse get_mut here because it would require E:
945            // From<HandleError>, which is inconvenient...
946            let h = Handle::from_u64(h)?;
947            let map = self.map.read().unwrap();
948            let mtx = map.get(h)?;
949            let mut hm = mtx.lock().unwrap();
950            Ok(callback(&mut *hm)?)
951        })
952    }
953
954    /// Helper that performs both a
955    /// [`call_with_result`][crate::call_with_result] and
956    /// [`get`](ConcurrentHandleMap::get).
957    pub fn call_with_result<R, E, F>(
958        &self,
959        out_error: &mut ExternError,
960        h: u64,
961        callback: F,
962    ) -> R::Value
963    where
964        F: std::panic::UnwindSafe + FnOnce(&T) -> Result<R, E>,
965        ExternError: From<E>,
966        R: IntoFfi,
967    {
968        self.call_with_result_mut(out_error, h, |r| callback(r))
969    }
970
971    /// Helper that performs both a
972    /// [`call_with_output`][crate::call_with_output] and
973    /// [`get`](ConcurrentHandleMap::get).
974    pub fn call_with_output<R, F>(
975        &self,
976        out_error: &mut ExternError,
977        h: u64,
978        callback: F,
979    ) -> R::Value
980    where
981        F: std::panic::UnwindSafe + FnOnce(&T) -> R,
982        R: IntoFfi,
983    {
984        self.call_with_result(out_error, h, |r| -> Result<_, HandleError> {
985            Ok(callback(r))
986        })
987    }
988
989    /// Helper that performs both a
990    /// [`call_with_output`][crate::call_with_output] and
991    /// [`get_mut`](ConcurrentHandleMap::get).
992    pub fn call_with_output_mut<R, F>(
993        &self,
994        out_error: &mut ExternError,
995        h: u64,
996        callback: F,
997    ) -> R::Value
998    where
999        F: std::panic::UnwindSafe + FnOnce(&mut T) -> R,
1000        R: IntoFfi,
1001    {
1002        self.call_with_result_mut(out_error, h, |r| -> Result<_, HandleError> {
1003            Ok(callback(r))
1004        })
1005    }
1006
1007    /// Use `constructor` to create and insert a `T`, while inside a
1008    /// [`call_with_result`][crate::call_with_result] call (to handle panics and
1009    /// map errors onto an [`ExternError`][crate::ExternError]).
1010    pub fn insert_with_result<E, F>(&self, out_error: &mut ExternError, constructor: F) -> u64
1011    where
1012        F: std::panic::UnwindSafe + FnOnce() -> Result<T, E>,
1013        ExternError: From<E>,
1014    {
1015        use crate::call_with_result;
1016        call_with_result(out_error, || -> Result<_, ExternError> {
1017            // Note: it's important that we don't call the constructor while
1018            // we're holding the write lock, because we don't want to poison
1019            // the entire map if it panics!
1020            let to_insert = constructor()?;
1021            Ok(self.insert(to_insert))
1022        })
1023    }
1024
1025    /// Equivalent to
1026    /// [`insert_with_result`](ConcurrentHandleMap::insert_with_result) for the
1027    /// case where the constructor cannot produce an error.
1028    ///
1029    /// The name is somewhat dubious, since there's no `output`, but it's
1030    /// intended to make it clear that it contains a
1031    /// [`call_with_output`][crate::call_with_output] internally.
1032    pub fn insert_with_output<F>(&self, out_error: &mut ExternError, constructor: F) -> u64
1033    where
1034        F: std::panic::UnwindSafe + FnOnce() -> T,
1035    {
1036        // The Err type isn't important here beyond being convertable to ExternError
1037        self.insert_with_result(out_error, || -> Result<_, HandleError> {
1038            Ok(constructor())
1039        })
1040    }
1041}
1042
1043impl<T> Default for ConcurrentHandleMap<T> {
1044    #[inline]
1045    fn default() -> Self {
1046        Self::new()
1047    }
1048}
1049
1050// Returns the next map_id.
1051fn next_handle_map_id() -> u16 {
1052    let id = HANDLE_MAP_ID_COUNTER
1053        .fetch_add(1, Ordering::SeqCst)
1054        .wrapping_add(1);
1055    id as u16
1056}
1057
1058// Note: These IDs are only used to detect using a key against the wrong HandleMap.
1059// We ensure they're randomly initialized, to prevent using them across separately
1060// compiled .so files.
1061lazy_static::lazy_static! {
1062    // This should be `AtomicU16`, but those aren't stablilized yet.
1063    // Instead, we just cast to u16 on read.
1064    static ref HANDLE_MAP_ID_COUNTER: AtomicUsize = {
1065        // Abuse HashMap's RandomState to get a strong RNG without bringing in
1066        // the `rand` crate (OTOH maybe we should just bring in the rand crate?)
1067        use std::collections::hash_map::RandomState;
1068        use std::hash::{BuildHasher, Hasher};
1069        let init = RandomState::new().build_hasher().finish() as usize;
1070        AtomicUsize::new(init)
1071    };
1072}
1073
1074#[cfg(test)]
1075mod test {
1076    use super::*;
1077
1078    #[derive(PartialEq, Debug)]
1079    pub(super) struct Foobar(usize);
1080
1081    #[test]
1082    fn test_invalid_handle() {
1083        assert_eq!(Handle::from_u64(0), Err(HandleError::NullHandle));
1084        // Valid except `version` is odd
1085        assert_eq!(
1086            Handle::from_u64((u64::from(HANDLE_MAGIC) << 48) | 0x1234_0012_0001),
1087            Err(HandleError::InvalidHandle)
1088        );
1089
1090        assert_eq!(
1091            Handle::from_u64((u64::from(HANDLE_MAGIC) << 48) | 0x1234_0012_0002),
1092            Ok(Handle {
1093                version: 0x0002,
1094                index: 0x0012,
1095                map_id: 0x1234,
1096            })
1097        );
1098    }
1099
1100    #[test]
1101    fn test_correct_value_single() {
1102        let mut map = HandleMap::new();
1103        let handle = map.insert(Foobar(1234));
1104        assert_eq!(map.get(handle).unwrap(), &Foobar(1234));
1105        map.delete(handle).unwrap();
1106        assert_eq!(map.get(handle), Err(HandleError::StaleVersion));
1107    }
1108
1109    #[test]
1110    fn test_correct_value_multiple() {
1111        let mut map = HandleMap::new();
1112        let handle1 = map.insert(Foobar(1234));
1113        let handle2 = map.insert(Foobar(4321));
1114        assert_eq!(map.get(handle1).unwrap(), &Foobar(1234));
1115        assert_eq!(map.get(handle2).unwrap(), &Foobar(4321));
1116        map.delete(handle1).unwrap();
1117        assert_eq!(map.get(handle1), Err(HandleError::StaleVersion));
1118        assert_eq!(map.get(handle2).unwrap(), &Foobar(4321));
1119    }
1120
1121    #[test]
1122    fn test_wrong_map() {
1123        let mut map1 = HandleMap::new();
1124        let mut map2 = HandleMap::new();
1125
1126        let handle1 = map1.insert(Foobar(1234));
1127        let handle2 = map2.insert(Foobar(1234));
1128
1129        assert_eq!(map1.get(handle1).unwrap(), &Foobar(1234));
1130        assert_eq!(map2.get(handle2).unwrap(), &Foobar(1234));
1131
1132        assert_eq!(map1.get(handle2), Err(HandleError::WrongMap));
1133        assert_eq!(map2.get(handle1), Err(HandleError::WrongMap));
1134    }
1135
1136    #[test]
1137    fn test_bad_index() {
1138        let map: HandleMap<Foobar> = HandleMap::new();
1139        assert_eq!(
1140            map.get(Handle {
1141                map_id: map.id,
1142                version: 2,
1143                index: 100
1144            }),
1145            Err(HandleError::IndexPastEnd)
1146        );
1147    }
1148
1149    #[test]
1150    fn test_resizing() {
1151        let mut map = HandleMap::new();
1152        let mut handles = vec![];
1153        for i in 0..1000 {
1154            handles.push(map.insert(Foobar(i)))
1155        }
1156        for (i, &h) in handles.iter().enumerate() {
1157            assert_eq!(map.get(h).unwrap(), &Foobar(i));
1158            assert_eq!(map.remove(h).unwrap(), Foobar(i));
1159        }
1160        let mut handles2 = vec![];
1161        for i in 1000..2000 {
1162            // Not really related to this test, but it's convenient to check this here.
1163            let h = map.insert(Foobar(i));
1164            let hu = h.into_u64();
1165            assert_eq!(Handle::from_u64(hu).unwrap(), h);
1166            handles2.push(hu);
1167        }
1168
1169        for (i, (&h0, h1u)) in handles.iter().zip(handles2).enumerate() {
1170            // It's still a stale version, even though the slot is occupied again.
1171            assert_eq!(map.get(h0), Err(HandleError::StaleVersion));
1172            let h1 = Handle::from_u64(h1u).unwrap();
1173            assert_eq!(map.get(h1).unwrap(), &Foobar(i + 1000));
1174        }
1175    }
1176
1177    /// Tests that check our behavior when panicing.
1178    ///
1179    /// Naturally these require panic=unwind, which means we can't run them when
1180    /// generating coverage (well, `-Zprofile`-based coverage can't -- although
1181    /// ptrace-based coverage like tarpaulin can), and so we turn them off.
1182    ///
1183    /// (For clarity, `cfg(coverage)` is not a standard thing. We add it in
1184    /// `automation/emit_coverage_info.sh`, and you can force it by adding
1185    /// "--cfg coverage" to your RUSTFLAGS manually if you need to do so).
1186    #[cfg(not(coverage))]
1187    mod panic_tests {
1188        use super::*;
1189
1190        struct PanicOnDrop(());
1191        impl Drop for PanicOnDrop {
1192            fn drop(&mut self) {
1193                panic!("intentional panic (drop)");
1194            }
1195        }
1196
1197        #[test]
1198        fn test_panicking_drop() {
1199            let map = ConcurrentHandleMap::new();
1200            let h = map.insert(PanicOnDrop(())).into_u64();
1201            let mut e = ExternError::success();
1202            crate::call_with_result(&mut e, || map.delete_u64(h));
1203            assert_eq!(e.get_code(), crate::ErrorCode::PANIC);
1204            let _ = unsafe { e.get_and_consume_message() };
1205            assert!(!map.map.is_poisoned());
1206            let inner = map.map.read().unwrap();
1207            inner.assert_valid();
1208            assert_eq!(inner.len(), 0);
1209        }
1210
1211        #[test]
1212        fn test_panicking_call_with() {
1213            let map = ConcurrentHandleMap::new();
1214            let h = map.insert(Foobar(0)).into_u64();
1215            let mut e = ExternError::success();
1216            map.call_with_output(&mut e, h, |_thing| {
1217                panic!("intentional panic (call_with_output)");
1218            });
1219
1220            assert_eq!(e.get_code(), crate::ErrorCode::PANIC);
1221            let _ = unsafe { e.get_and_consume_message() };
1222
1223            {
1224                assert!(!map.map.is_poisoned());
1225                let inner = map.map.read().unwrap();
1226                inner.assert_valid();
1227                assert_eq!(inner.len(), 1);
1228                let mut seen = false;
1229                for e in &inner.entries {
1230                    if let EntryState::Active(v) = &e.state {
1231                        assert!(!seen);
1232                        assert!(v.is_poisoned());
1233                        seen = true;
1234                    }
1235                }
1236            }
1237            assert!(map.delete_u64(h).is_ok());
1238            assert!(!map.map.is_poisoned());
1239            let inner = map.map.read().unwrap();
1240            inner.assert_valid();
1241            assert_eq!(inner.len(), 0);
1242        }
1243
1244        #[test]
1245        fn test_panicking_insert_with() {
1246            let map = ConcurrentHandleMap::new();
1247            let mut e = ExternError::success();
1248            let res = map.insert_with_output(&mut e, || {
1249                panic!("intentional panic (insert_with_output)");
1250            });
1251
1252            assert_eq!(e.get_code(), crate::ErrorCode::PANIC);
1253            let _ = unsafe { e.get_and_consume_message() };
1254
1255            assert_eq!(res, 0);
1256
1257            assert!(!map.map.is_poisoned());
1258            let inner = map.map.read().unwrap();
1259            inner.assert_valid();
1260            assert_eq!(inner.len(), 0);
1261        }
1262    }
1263}