rwf2/Rocket
1
0
Fork 0
mirror of https://github.com/rwf2/Rocket.git synced 2025-01-18 07:29:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Swap 'Rocket' manually without using 'replace_with'.

Browse Source
This commit is contained in:
Sergio Benitez 2020-06-28 14:01:06 -07:00
parent 3ced188f7d
commit 9277ddafdf
2 changed files with 19 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
core/lib/src/ext.rs
View File

@ -2,7 +2,7 @@ use std::io::{self, Cursor};
use std::pin::Pin;
use std::task::{Poll, Context};
use futures::{ready, Future, future::BoxFuture, stream::Stream};
use futures::{ready, future::BoxFuture, stream::Stream};
use tokio::io::{AsyncRead, AsyncReadExt as _};
use crate::http::hyper;
@ -111,53 +111,3 @@ impl AsyncRead for AsyncReadBody {
}
}
}
// The code below was adapted from the `replace_with` crate and reproduced here
// under the rights granted by the MIT license. The code is copyright the
// `replace_with` developers. See LICENSE-MIT for the full text.
struct OnUnwind<F: FnOnce()>(std::mem::ManuallyDrop<F>);
impl<F: FnOnce()> Drop for OnUnwind<F> {
#[inline(always)]
fn drop(&mut self) {
(unsafe { std::ptr::read(&*self.0) })();
}
}
#[inline(always)]
pub async fn async_on_unwind<F, Fut, T, P: FnOnce()>(f: F, p: P) -> T
where F: FnOnce() -> Fut, Fut: Future<Output = T>,
{
let x = OnUnwind(std::mem::ManuallyDrop::new(p));
let t = f().await;
let _ = unsafe { std::ptr::read(&*x.0) };
std::mem::forget(x);
t
}
#[inline]
pub async fn async_replace_with_or_else<T, Fut, F, D>(dest: &mut T, d: D, f: F)
where Fut: Future<Output = T>,
F: FnOnce(T) -> Fut,
D: FnOnce() -> T,
{
unsafe {
let old = std::ptr::read(dest);
let new = async_on_unwind(
|| async move { f(old).await },
|| std::ptr::write(dest, d()),
).await;
std::ptr::write(dest, new);
}
}
#[inline]
pub async fn async_replace_with<T, Fut, F>(dest: &mut T, f: F)
where Fut: Future<Output = T>,
F: FnOnce(T) -> Fut,
{
async_replace_with_or_else(dest, || std::process::abort(), f).await
}

20
core/lib/src/rocket.rs
View File

@ -24,7 +24,7 @@ use crate::outcome::Outcome;
use crate::error::{LaunchError, LaunchErrorKind};
use crate::fairing::{Fairing, Fairings};
use crate::logger::PaintExt;
use crate::ext::{AsyncReadExt, async_replace_with};
use crate::ext::AsyncReadExt;
use crate::shutdown::{ShutdownHandle, ShutdownHandleManaged};
use crate::http::{Method, Status, Header};
@ -109,6 +109,21 @@ impl Rocket {
self
}
// Create a "dummy" instance of `Rocket` to use while mem-swapping `self`.
fn dummy() -> Rocket {
Rocket {
manifest: vec![],
config: Config::development(),
router: Router::new(),
default_catchers: HashMap::new(),
catchers: HashMap::new(),
managed_state: Container::new(),
fairings: Fairings::new(),
shutdown_handle: ShutdownHandle(mpsc::channel(1).0),
shutdown_receiver: None,
}
}
// Instead of requiring the user to individually `await` each call to
// `attach()`, some operations are queued in `self.pending`. Functions that
// want to provide read access to any data from the Cargo, such as
@ -133,7 +148,8 @@ impl Rocket {
PreLaunchOp::Mount(base, routes) => self._mount(base, routes),
PreLaunchOp::Register(catchers) => self._register(catchers),
PreLaunchOp::Attach(fairing) => {
async_replace_with(self, |rocket| rocket._attach(fairing)).await;
let rocket = mem::replace(self, Rocket::dummy());
*self = rocket._attach(fairing).await;
self.manifest.append(&mut manifest);
manifest = mem::replace(&mut self.manifest, vec![]);
}