src/test/ui/drop/dynamic-drop-async.rs

   1 // Test that values are not leaked in async functions, even in the cases where:
   2 // * Dropping one of the values panics while running the future.
   3 // * The future is dropped at one of its suspend points.
   4 // * Dropping one of the values panics while dropping the future.
   5
   6 // run-pass
   7 // edition:2018
   8 // ignore-wasm32-bare compiled with panic=abort by default
   9
  10 #![allow(unused_assignments)]
  11 #![allow(unused_variables)]
  12 #![feature(slice_patterns)]
  13 #![feature(async_await)]
  14
  15 use std::{
  16     cell::{Cell, RefCell},
  17     future::Future,
  18     marker::Unpin,
  19     panic,
  20     pin::Pin,
  21     ptr,
  22     rc::Rc,
  23     task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
  24     usize,
  25 };
  26
  27 struct InjectedFailure;
  28
  29 struct Defer<T> {
  30     ready: bool,
  31     value: Option<T>,
  32 }
  33
  34 impl<T: Unpin> Future for Defer<T> {
  35     type Output = T;
  36     fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
  37         if self.ready {
  38             Poll::Ready(self.value.take().unwrap())
  39         } else {
  40             self.ready = true;
  41             Poll::Pending
  42         }
  43     }
  44 }
  45
  46 /// Allocator tracks the creation and destruction of `Ptr`s.
  47 /// The `failing_op`-th operation will panic.
  48 struct Allocator {
  49     data: RefCell<Vec<bool>>,
  50     failing_op: usize,
  51     cur_ops: Cell<usize>,
  52 }
  53
  54 impl panic::UnwindSafe for Allocator {}
  55 impl panic::RefUnwindSafe for Allocator {}
  56
  57 impl Drop for Allocator {
  58     fn drop(&mut self) {
  59         let data = self.data.borrow();
  60         if data.iter().any(|d| *d) {
  61             panic!("missing free: {:?}", data);
  62         }
  63     }
  64 }
  65
  66 impl Allocator {
  67     fn new(failing_op: usize) -> Self {
  68         Allocator { failing_op, cur_ops: Cell::new(0), data: RefCell::new(vec![]) }
  69     }
  70     fn alloc(&self) -> impl Future<Output = Ptr<'_>> + '_ {
  71         self.fallible_operation();
  72
  73         let mut data = self.data.borrow_mut();
  74
  75         let addr = data.len();
  76         data.push(true);
  77         Defer { ready: false, value: Some(Ptr(addr, self)) }
  78     }
  79     fn fallible_operation(&self) {
  80         self.cur_ops.set(self.cur_ops.get() + 1);
  81
  82         if self.cur_ops.get() == self.failing_op {
  83             panic!(InjectedFailure);
  84         }
  85     }
  86 }
  87
  88 // Type that tracks whether it was dropped and can panic when it's created or
  89 // destroyed.
  90 struct Ptr<'a>(usize, &'a Allocator);
  91 impl<'a> Drop for Ptr<'a> {
  92     fn drop(&mut self) {
  93         match self.1.data.borrow_mut()[self.0] {
  94             false => panic!("double free at index {:?}", self.0),
  95             ref mut d => *d = false,
  96         }
  97
  98         self.1.fallible_operation();
  99     }
 100 }
 101
 102 async fn dynamic_init(a: Rc<Allocator>, c: bool) {
 103     let _x;
 104     if c {
 105         _x = Some(a.alloc().await);
 106     }
 107 }
 108
 109 async fn dynamic_drop(a: Rc<Allocator>, c: bool) {
 110     let x = a.alloc().await;
 111     if c {
 112         Some(x)
 113     } else {
 114         None
 115     };
 116 }
 117
 118 struct TwoPtrs<'a>(Ptr<'a>, Ptr<'a>);
 119 async fn struct_dynamic_drop(a: Rc<Allocator>, c0: bool, c1: bool, c: bool) {
 120     for i in 0..2 {
 121         let x;
 122         let y;
 123         if (c0 && i == 0) || (c1 && i == 1) {
 124             x = (a.alloc().await, a.alloc().await, a.alloc().await);
 125             y = TwoPtrs(a.alloc().await, a.alloc().await);
 126             if c {
 127                 drop(x.1);
 128                 a.alloc().await;
 129                 drop(y.0);
 130                 a.alloc().await;
 131             }
 132         }
 133     }
 134 }
 135
 136 async fn field_assignment(a: Rc<Allocator>, c0: bool) {
 137     let mut x = (TwoPtrs(a.alloc().await, a.alloc().await), a.alloc().await);
 138
 139     x.1 = a.alloc().await;
 140     x.1 = a.alloc().await;
 141
 142     let f = (x.0).0;
 143     a.alloc().await;
 144     if c0 {
 145         (x.0).0 = f;
 146     }
 147     a.alloc().await;
 148 }
 149
 150 async fn assignment(a: Rc<Allocator>, c0: bool, c1: bool) {
 151     let mut _v = a.alloc().await;
 152     let mut _w = a.alloc().await;
 153     if c0 {
 154         drop(_v);
 155     }
 156     _v = _w;
 157     if c1 {
 158         _w = a.alloc().await;
 159     }
 160 }
 161
 162 async fn array_simple(a: Rc<Allocator>) {
 163     let _x = [a.alloc().await, a.alloc().await, a.alloc().await, a.alloc().await];
 164 }
 165
 166 async fn vec_simple(a: Rc<Allocator>) {
 167     let _x = vec![a.alloc().await, a.alloc().await, a.alloc().await, a.alloc().await];
 168 }
 169
 170 async fn mixed_drop_and_nondrop(a: Rc<Allocator>) {
 171     // check that destructor panics handle drop
 172     // and non-drop blocks in the same scope correctly.
 173     //
 174     // Surprisingly enough, this used to not work.
 175     let (x, y, z);
 176     x = a.alloc().await;
 177     y = 5;
 178     z = a.alloc().await;
 179 }
 180
 181 #[allow(unreachable_code)]
 182 async fn vec_unreachable(a: Rc<Allocator>) {
 183     let _x = vec![a.alloc().await, a.alloc().await, a.alloc().await, return];
 184 }
 185
 186 async fn slice_pattern_one_of(a: Rc<Allocator>, i: usize) {
 187     let array = [a.alloc().await, a.alloc().await, a.alloc().await, a.alloc().await];
 188     let _x = match i {
 189         0 => {
 190             let [a, ..] = array;
 191             a
 192         }
 193         1 => {
 194             let [_, a, ..] = array;
 195             a
 196         }
 197         2 => {
 198             let [_, _, a, _] = array;
 199             a
 200         }
 201         3 => {
 202             let [_, _, _, a] = array;
 203             a
 204         }
 205         _ => panic!("unmatched"),
 206     };
 207     a.alloc().await;
 208 }
 209
 210 async fn subslice_pattern_from_end_with_drop(a: Rc<Allocator>, arg: bool, arg2: bool) {
 211     let arr = [a.alloc().await, a.alloc().await, a.alloc().await, a.alloc().await, a.alloc().await];
 212     if arg2 {
 213         drop(arr);
 214         return;
 215     }
 216
 217     if arg {
 218         let [.., _x, _] = arr;
 219     } else {
 220         let [_, _y..] = arr;
 221     }
 222     a.alloc().await;
 223 }
 224
 225 async fn subslice_pattern_reassign(a: Rc<Allocator>) {
 226     let mut ar = [a.alloc().await, a.alloc().await, a.alloc().await];
 227     let [_, _, _x] = ar;
 228     ar = [a.alloc().await, a.alloc().await, a.alloc().await];
 229     let [_, _y..] = ar;
 230     a.alloc().await;
 231 }
 232
 233 fn run_test<F, G>(cx: &mut Context<'_>, ref f: F)
 234 where
 235     F: Fn(Rc<Allocator>) -> G,
 236     G: Future<Output = ()>,
 237 {
 238     for polls in 0.. {
 239         // Run without any panics to find which operations happen after the
 240         // penultimate `poll`.
 241         let first_alloc = Rc::new(Allocator::new(usize::MAX));
 242         let mut fut = Box::pin(f(first_alloc.clone()));
 243         let mut ops_before_last_poll = 0;
 244         let mut completed = false;
 245         for _ in 0..polls {
 246             ops_before_last_poll = first_alloc.cur_ops.get();
 247             if let Poll::Ready(()) = fut.as_mut().poll(cx) {
 248                 completed = true;
 249             }
 250         }
 251         drop(fut);
 252
 253         // Start at `ops_before_last_poll` so that we will always be able to
 254         // `poll` the expected number of times.
 255         for failing_op in ops_before_last_poll..first_alloc.cur_ops.get() {
 256             let alloc = Rc::new(Allocator::new(failing_op + 1));
 257             let f = &f;
 258             let cx = &mut *cx;
 259             let result = panic::catch_unwind(panic::AssertUnwindSafe(move || {
 260                 let mut fut = Box::pin(f(alloc));
 261                 for _ in 0..polls {
 262                     let _ = fut.as_mut().poll(cx);
 263                 }
 264                 drop(fut);
 265             }));
 266             match result {
 267                 Ok(..) => panic!("test executed more ops on first call"),
 268                 Err(e) => {
 269                     if e.downcast_ref::<InjectedFailure>().is_none() {
 270                         panic::resume_unwind(e);
 271                     }
 272                 }
 273             }
 274         }
 275
 276         if completed {
 277             break;
 278         }
 279     }
 280 }
 281
 282 fn clone_waker(data: *const ()) -> RawWaker {
 283     RawWaker::new(data, &RawWakerVTable::new(clone_waker, drop, drop, drop))
 284 }
 285
 286 fn main() {
 287     let waker = unsafe { Waker::from_raw(clone_waker(ptr::null())) };
 288     let context = &mut Context::from_waker(&waker);
 289
 290     run_test(context, |a| dynamic_init(a, false));
 291     run_test(context, |a| dynamic_init(a, true));
 292     run_test(context, |a| dynamic_drop(a, false));
 293     run_test(context, |a| dynamic_drop(a, true));
 294
 295     run_test(context, |a| assignment(a, false, false));
 296     run_test(context, |a| assignment(a, false, true));
 297     run_test(context, |a| assignment(a, true, false));
 298     run_test(context, |a| assignment(a, true, true));
 299
 300     run_test(context, |a| array_simple(a));
 301     run_test(context, |a| vec_simple(a));
 302     run_test(context, |a| vec_unreachable(a));
 303
 304     run_test(context, |a| struct_dynamic_drop(a, false, false, false));
 305     run_test(context, |a| struct_dynamic_drop(a, false, false, true));
 306     run_test(context, |a| struct_dynamic_drop(a, false, true, false));
 307     run_test(context, |a| struct_dynamic_drop(a, false, true, true));
 308     run_test(context, |a| struct_dynamic_drop(a, true, false, false));
 309     run_test(context, |a| struct_dynamic_drop(a, true, false, true));
 310     run_test(context, |a| struct_dynamic_drop(a, true, true, false));
 311     run_test(context, |a| struct_dynamic_drop(a, true, true, true));
 312
 313     run_test(context, |a| field_assignment(a, false));
 314     run_test(context, |a| field_assignment(a, true));
 315
 316     run_test(context, |a| mixed_drop_and_nondrop(a));
 317
 318     run_test(context, |a| slice_pattern_one_of(a, 0));
 319     run_test(context, |a| slice_pattern_one_of(a, 1));
 320     run_test(context, |a| slice_pattern_one_of(a, 2));
 321     run_test(context, |a| slice_pattern_one_of(a, 3));
 322
 323     run_test(context, |a| subslice_pattern_from_end_with_drop(a, true, true));
 324     run_test(context, |a| subslice_pattern_from_end_with_drop(a, true, false));
 325     run_test(context, |a| subslice_pattern_from_end_with_drop(a, false, true));
 326     run_test(context, |a| subslice_pattern_from_end_with_drop(a, false, false));
 327     run_test(context, |a| subslice_pattern_reassign(a));
 328 }