vendor/rustc-rayon/tests/sort-panic-safe.rs

   1 use rand::distributions::Uniform;
   2 use rand::{thread_rng, Rng};
   3 use rayon::prelude::*;
   4 use std::cell::Cell;
   5 use std::cmp::{self, Ordering};
   6 use std::panic;
   7 use std::sync::atomic::AtomicUsize;
   8 use std::sync::atomic::Ordering::Relaxed;
   9 use std::thread;
  10
  11 static VERSIONS: AtomicUsize = AtomicUsize::new(0);
  12
  13 lazy_static::lazy_static! {
  14     static ref DROP_COUNTS: Vec<AtomicUsize> = (0..20_000).map(|_| AtomicUsize::new(0)).collect();
  15 }
  16
  17 #[derive(Clone, Eq)]
  18 struct DropCounter {
  19     x: u32,
  20     id: usize,
  21     version: Cell<usize>,
  22 }
  23
  24 impl PartialEq for DropCounter {
  25     fn eq(&self, other: &Self) -> bool {
  26         self.partial_cmp(other) == Some(Ordering::Equal)
  27     }
  28 }
  29
  30 impl PartialOrd for DropCounter {
  31     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
  32         self.version.set(self.version.get() + 1);
  33         other.version.set(other.version.get() + 1);
  34         VERSIONS.fetch_add(2, Relaxed);
  35         self.x.partial_cmp(&other.x)
  36     }
  37 }
  38
  39 impl Ord for DropCounter {
  40     fn cmp(&self, other: &Self) -> Ordering {
  41         self.partial_cmp(other).unwrap()
  42     }
  43 }
  44
  45 impl Drop for DropCounter {
  46     fn drop(&mut self) {
  47         DROP_COUNTS[self.id].fetch_add(1, Relaxed);
  48         VERSIONS.fetch_sub(self.version.get(), Relaxed);
  49     }
  50 }
  51
  52 macro_rules! test {
  53     ($input:ident, $func:ident) => {
  54         let len = $input.len();
  55
  56         // Work out the total number of comparisons required to sort
  57         // this array...
  58         let count = AtomicUsize::new(0);
  59         $input.to_owned().$func(|a, b| {
  60             count.fetch_add(1, Relaxed);
  61             a.cmp(b)
  62         });
  63
  64         let mut panic_countdown = count.load(Relaxed);
  65         let step = if len <= 100 {
  66             1
  67         } else {
  68             cmp::max(1, panic_countdown / 10)
  69         };
  70
  71         // ... and then panic after each `step` comparisons.
  72         loop {
  73             // Refresh the counters.
  74             VERSIONS.store(0, Relaxed);
  75             for i in 0..len {
  76                 DROP_COUNTS[i].store(0, Relaxed);
  77             }
  78
  79             let v = $input.to_owned();
  80             let _ = thread::spawn(move || {
  81                 let mut v = v;
  82                 let panic_countdown = AtomicUsize::new(panic_countdown);
  83                 v.$func(|a, b| {
  84                     if panic_countdown.fetch_sub(1, Relaxed) == 1 {
  85                         SILENCE_PANIC.with(|s| s.set(true));
  86                         panic!();
  87                     }
  88                     a.cmp(b)
  89                 })
  90             })
  91             .join();
  92
  93             // Check that the number of things dropped is exactly
  94             // what we expect (i.e. the contents of `v`).
  95             for (i, c) in DROP_COUNTS.iter().enumerate().take(len) {
  96                 let count = c.load(Relaxed);
  97                 assert!(
  98                     count == 1,
  99                     "found drop count == {} for i == {}, len == {}",
 100                     count,
 101                     i,
 102                     len
 103                 );
 104             }
 105
 106             // Check that the most recent versions of values were dropped.
 107             assert_eq!(VERSIONS.load(Relaxed), 0);
 108
 109             if panic_countdown < step {
 110                 break;
 111             }
 112             panic_countdown -= step;
 113         }
 114     };
 115 }
 116
 117 thread_local!(static SILENCE_PANIC: Cell<bool> = Cell::new(false));
 118
 119 #[test]
 120 fn sort_panic_safe() {
 121     let prev = panic::take_hook();
 122     panic::set_hook(Box::new(move |info| {
 123         if !SILENCE_PANIC.with(Cell::get) {
 124             prev(info);
 125         }
 126     }));
 127
 128     for &len in &[1, 2, 3, 4, 5, 10, 20, 100, 500, 5_000, 20_000] {
 129         let len_dist = Uniform::new(0, len);
 130         for &modulus in &[5, 30, 1_000, 20_000] {
 131             for &has_runs in &[false, true] {
 132                 let mut rng = thread_rng();
 133                 let mut input = (0..len)
 134                     .map(|id| DropCounter {
 135                         x: rng.gen_range(0..modulus),
 136                         id,
 137                         version: Cell::new(0),
 138                     })
 139                     .collect::<Vec<_>>();
 140
 141                 if has_runs {
 142                     for c in &mut input {
 143                         c.x = c.id as u32;
 144                     }
 145
 146                     for _ in 0..5 {
 147                         let a = rng.sample(&len_dist);
 148                         let b = rng.sample(&len_dist);
 149                         if a < b {
 150                             input[a..b].reverse();
 151                         } else {
 152                             input.swap(a, b);
 153                         }
 154                     }
 155                 }
 156
 157                 test!(input, par_sort_by);
 158                 test!(input, par_sort_unstable_by);
 159             }
 160         }
 161     }
 162 }