use core::cell::Cell;
use core::fmt;
use core::future::Future;
use core::marker::PhantomData;
use core::pin::Pin;
use core::task::{Context, Poll};
use std::sync::Arc;
use pi_result::{InteropResultExt, RawResult};
use crate::bounded_blocking;
use crate::{BufferFailure, DetachableWriteBuffer, MmapRange};
struct MmapDetachOnDrop<T> {
task: Option<async_global_executor::Task<T>>,
}
impl<T> MmapDetachOnDrop<T> {
fn new(task: async_global_executor::Task<T>) -> Self {
Self { task: Some(task) }
}
}
impl<T> Unpin for MmapDetachOnDrop<T> {}
impl<T> Future for MmapDetachOnDrop<T> {
type Output = T;
fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<T> {
let this = self.get_mut();
let result = Pin::new(
this.task.as_mut().expect("mmap task must exist while polling"),
)
.poll(context);
if result.is_ready() {
this.task.take();
}
result
}
}
impl<T> Drop for MmapDetachOnDrop<T> {
fn drop(&mut self) {
if let Some(task) = self.task.take() {
task.detach();
}
}
}
struct MappingCleanup {
release: Option<Box<dyn FnOnce() + Send + Sync + 'static>>,
}
impl MappingCleanup {
fn new(release: Box<dyn FnOnce() + Send + Sync + 'static>) -> Self {
Self {
release: Some(release),
}
}
}
impl Drop for MappingCleanup {
fn drop(&mut self) {
if let Some(release) = self.release.take() {
release();
}
}
}
struct ReadMmapState {
range: MmapRange,
mapping: memmap2::Mmap,
cleanup: MappingCleanup,
}
pub struct ReadMmapHandle {
state: Arc<ReadMmapState>,
}
impl ReadMmapHandle {
pub(crate) fn from_mapping(
range: MmapRange,
mapping: memmap2::Mmap,
release: Box<dyn FnOnce() + Send + Sync + 'static>,
) -> Self {
debug_assert_eq!(range.len(), mapping.len() as u64);
Self {
state: Arc::new(ReadMmapState {
range,
mapping,
cleanup: MappingCleanup::new(release),
}),
}
}
#[must_use]
pub fn range(&self) -> &MmapRange {
&self.state.range
}
#[must_use]
#[allow(clippy::len_without_is_empty)]
pub fn len(&self) -> usize {
self.state.mapping.len()
}
#[must_use]
pub fn as_bytes(&self) -> &[u8] {
&self.state.mapping
}
}
impl AsRef<[u8]> for ReadMmapHandle {
fn as_ref(&self) -> &[u8] {
self.as_bytes()
}
}
impl DetachableWriteBuffer for ReadMmapHandle {
type Detached = ReadMmapHandle;
type Recovery = ();
fn try_detach(self) -> RawResult<(Self::Detached, Self::Recovery), BufferFailure<Self>> {
Ok((self, ()))
}
fn recover_from_detached(
detached: Self::Detached,
_recovery: Self::Recovery,
) -> Self {
detached
}
}
impl<'a> DetachableWriteBuffer for &'a ReadMmapHandle {
type Detached = ReadMmapHandle;
type Recovery = &'a ReadMmapHandle;
fn try_detach(self) -> RawResult<(Self::Detached, Self::Recovery), BufferFailure<Self>> {
let detached = ReadMmapHandle {
state: Arc::clone(&self.state),
};
Ok((detached, self))
}
fn recover_from_detached(
detached: Self::Detached,
recovery: Self::Recovery,
) -> Self {
drop(detached);
recovery
}
}
impl Clone for ReadMmapHandle {
fn clone(&self) -> Self {
Self {
state: Arc::clone(&self.state),
}
}
}
impl fmt::Debug for ReadMmapHandle {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ReadMmapHandle")
.field("range", self.range())
.field("len", &self.len())
.finish()
}
}
impl fmt::Display for ReadMmapHandle {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(formatter, "read-only mmap {}", self.range())
}
}
struct ReadWriteMmapState {
range: MmapRange,
mapping: memmap2::MmapRaw,
cleanup: MappingCleanup,
}
pub struct ReadWriteMmapHandle {
state: Arc<ReadWriteMmapState>,
not_sync: PhantomData<Cell<()>>,
}
impl ReadWriteMmapHandle {
pub(crate) fn from_mapping(
range: MmapRange,
mapping: memmap2::MmapRaw,
release: Box<dyn FnOnce() + Send + Sync + 'static>,
) -> Self {
debug_assert_eq!(range.len(), mapping.len() as u64);
Self {
state: Arc::new(ReadWriteMmapState {
range,
mapping,
cleanup: MappingCleanup::new(release),
}),
not_sync: PhantomData,
}
}
#[must_use]
pub fn range(&self) -> &MmapRange {
&self.state.range
}
#[must_use]
#[allow(clippy::len_without_is_empty)]
pub fn len(&self) -> usize {
self.state.mapping.len()
}
#[must_use]
pub fn as_bytes(&self) -> &[u8] {
unsafe {
core::slice::from_raw_parts(
self.state.mapping.as_ptr(),
self.state.mapping.len(),
)
}
}
#[must_use]
pub fn as_bytes_mut(&mut self) -> &mut [u8] {
unsafe {
core::slice::from_raw_parts_mut(
self.state.mapping.as_mut_ptr(),
self.state.mapping.len(),
)
}
}
pub fn flush(
&mut self,
) -> impl Future<Output = pi_result::Result<()>> + Send + '_ {
let state = Arc::clone(&self.state);
async move {
let task = async_global_executor::spawn(async move {
bounded_blocking::unblock(move || state.mapping.flush())
.await?
.into_classified_error()
});
MmapDetachOnDrop::new(task).await
}
}
}
impl AsMut<[u8]> for ReadWriteMmapHandle {
fn as_mut(&mut self) -> &mut [u8] {
self.as_bytes_mut()
}
}
impl fmt::Debug for ReadWriteMmapHandle {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ReadWriteMmapHandle")
.field("range", self.range())
.field("len", &self.len())
.finish()
}
}
impl fmt::Display for ReadWriteMmapHandle {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(formatter, "read-write mmap {}", self.range())
}
}