异步特质方法与轮询扩展特质

Rust 的特质系统不允许在特质中定义异步方法。通过返回boxed futures，async-trait 包解决了这个问题。

对于大多数情况，async-trait 工作得很好，然而一些底层特质不能每次被调用时都进行堆分配，因此其他包如 tokio 开始使用 "扩展特质"。它们基本上将 Read::read 转换为 AsyncRead::poll_read，它使用轮询，并提供了一个新的方法 AsyncReadExt::read，该方法返回一个内部使用 poll_read 的 future（《AsyncReadExt》为所有《AsyncRead》自动实现）。由于 poll_read 是同步的，特质中不包含任何异步方法，可以用 Rust 编译。所有实现者只需实现 poll_read 而不是 read。

手动编写扩展特质可能会非常繁琐，因此可以使用 async_trait_ext 宏来自动处理所有这些模板代码。

动态特质对象

对于像以下这样的特质：

#[async_trait_ext]
trait AsyncRead {
    async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Result<usize>;
}

async_trait_ext 默认生成一个结构体，如下所示，以实现 AsyncRead::read 的 future。此结构体包含 Self，因此需要确定大小。

struct AsyncReadRead<'a, SelfTy: AsyncRead>(&'a mut SelfTy, &'a mut [u8]);

然而，也可以通过使用 async_trait_ext(dynamic) 选择动态类型。

生成以下内容，无需确定大小：

#[async_trait_ext(dynamic)]
trait AsyncRead {
    async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Result<usize>;
}

struct AsyncReadRead<'a>(&'a mut dyn AsyncRead, &'a mut [u8]);

提供的函数

有时为特质提供一些函数是有用的。使用 provided 属性标记函数会将它们移动到扩展特质中。

遗憾的是，需要 nightly 版本的 type_alias_impl_trait 功能来命名异步块的类型。因此，provided 属性只能与 async-trait-ext 的 provided 功能（默认未启用）一起使用。

示例

#![feature(type_alias_impl_trait)]

#[async_trait_ext(dynamic)]
trait AsyncRead {
    async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Result<usize>;
    
    #[async_fn(provided)]
    async fn read_u8<'a>(&'a mut self) -> Result<u8> {
        let mut buf = [0];
        self.read(&mut buf).await?;
        Ok(buf[0])
    }
}

示例

非动态

#[async_trait_ext]
pub trait Lock {
    async fn lock(&self) -> Result<LockGuard>;
}

展开为

pub trait Lock {
    fn poll_lock(&self, ctx: &mut ::core::task::Context) -> ::core::task::Poll<Result<LockGuard>>;
}

/// the future returned by [`LockExt::lock`]
pub struct LockLock<'__default_lifetime, __Self: Lock>(
    &'__default_lifetime __Self,
    ::core::marker::PhantomData<fn(__Self)>,
    ::core::marker::PhantomData<&'__default_lifetime ()>,
);

impl<'__default_lifetime, __Self: Lock> ::core::future::Future
    for LockLock<'__default_lifetime, __Self>
{
    type Output = Result<LockGuard>;
    
    fn poll(
        mut self: ::core::pin::Pin<&mut Self>,
        cx: &mut ::core::task::Context,
    ) -> ::core::task::Poll<Self::Output> {
        let this = &mut *self;
        <__Self as Lock>::poll_lock(this.0.into(), cx)
    }
}

pub trait LockExt: Lock + ::core::marker::Sized {
    fn lock(&self) -> LockLock<'_, Self>;
}

impl<__IMPL: Lock> LockExt for __IMPL {
    fn lock(&self) -> LockLock<'_, Self> {
        LockLock(
            self,
            ::core::marker::PhantomData,
            ::core::marker::PhantomData,
        )
    }
}

动态

#[async_trait_ext(dynamic)]
pub trait Write {
    async fn write<'a>(&'a mut self, buf: &'a [u8]) -> Result<usize>;
}

展开为

pub trait Write {
    fn poll_write<'a>(
        &'a mut self,
        buf: &'a [u8],
        ctx: &mut ::core::task::Context,
    ) -> ::core::task::Poll<Result<usize>>;
}

/// the future returned by [`WriteExt::write`]
pub struct WriteWrite<'a>(
    &'a mut dyn Write,
    &'a [u8],
    ::core::marker::PhantomData<&'a ()>,
);

impl<'a> ::core::future::Future for WriteWrite<'a> {
    type Output = Result<usize>;

    fn poll(
        mut self: ::core::pin::Pin<&mut Self>,
        cx: &mut ::core::task::Context,
    ) -> ::core::task::Poll<Self::Output> {
        let this = &mut *self;
        Write::poll_write(this.0.into(), this.1.into(), cx)
    }
}

pub trait WriteExt: Write {
    fn write<'a>(&'a mut self, buf: &'a [u8]) -> WriteWrite<'a>;
}

impl<__IMPL: Write> WriteExt for __IMPL {
    fn write<'a>(&'a mut self, buf: &'a [u8]) -> WriteWrite<'a> {
        WriteWrite(self, buf, ::core::marker::PhantomData)
    }
}

提供的方法

#![feature(type_alias_impl_trait)]

#[async_trait_ext]
pub trait ReadStatic {
    async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Result<usize>;

    #[async_fn(provided)]
    async fn read_until<'a>(&'a mut self, byte: u8, mut buf: &'a mut [u8]) -> Result<usize> {
        let mut b = [0];
        let mut bytes_read = 0;

        while !buf.is_empty() {
            match self.read(&mut b).await? {
                1 if b[0] != byte => {
                    bytes_read += 1;
                    buf[0] = b[0];
                    buf = &mut buf[1..];
                }
                _ => break,
            }
        }

        Ok(bytes_read)
    }
}

展开为

#![feature(type_alias_impl_trait)]

pub trait ReadStatic {
    fn poll_read<'a>(
        &'a mut self,
        buf: &'a mut [u8],
        ctx: &mut ::core::task::Context,
    ) -> ::core::task::Poll<Result<usize>>;
}

/// the future returned by [`ReadStaticExt::read`]
pub struct ReadStaticRead<'a, __Self: ReadStatic>(
    &'a mut __Self,
    &'a mut [u8],
    ::core::marker::PhantomData<fn(__Self)>,
    ::core::marker::PhantomData<&'a ()>,
);

impl<'a, __Self: ReadStatic> ::core::future::Future for ReadStaticRead<'a, __Self> {
    type Output = Result<usize>;

    fn poll(
        mut self: ::core::pin::Pin<&mut Self>,
        cx: &mut ::core::task::Context,
    ) -> ::core::task::Poll<Self::Output> {
        let this = &mut *self;
        <__Self as ReadStatic>::poll_read(this.0.into(), this.1.into(), cx)
    }
}

/// the future returned by [`ReadStaticExt::read_until`]
pub type ReadStaticReadUntil<'a, __Self: ReadStatic> =
    impl ::core::future::Future<Output = Result<usize>>;

pub trait ReadStaticExt: ReadStatic + ::core::marker::Sized {
    fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> ReadStaticRead<'a, Self>;
    fn read_until<'a>(&'a mut self, byte: u8, buf: &'a mut [u8]) -> ReadStaticReadUntil<'a, Self>;
}

impl<__IMPL: ReadStatic> ReadStaticExt for __IMPL {
    fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> ReadStaticRead<'a, Self> {
        ReadStaticRead(
            self,
            buf,
            ::core::marker::PhantomData,
            ::core::marker::PhantomData,
        )
    }

    fn read_until<'a>(
        &'a mut self,
        byte: u8,
        mut buf: &'a mut [u8],
    ) -> ReadStaticReadUntil<'a, Self> {
        async move {
            {
                let mut b = [0];
                let mut bytes_read = 0;
                while !buf.is_empty() {
                    match self.read(&mut b).await? {
                        1 if b[0] != byte => {
                            bytes_read += 1;
                            buf[0] = b[0];
                            buf = &mut buf[1..];
                        }
                        _ => break,
                    }
                }
                Ok(bytes_read)
            }
        }
    }
}

提供的方法 + 动态

#![feature(type_alias_impl_trait)]

#[async_trait_ext(dynamic)]
pub trait ReadDynamic {
    async fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Result<usize>;

    #[async_fn(provided)]
    async fn read_until<'a>(&'a mut self, byte: u8, mut buf: &'a mut [u8]) -> Result<usize> {
        let mut b = [0];
        let mut bytes_read = 0;

        while !buf.is_empty() {
            match self.read(&mut b).await? {
                1 if b[0] != byte => {
                    bytes_read += 1;
                    buf[0] = b[0];
                    buf = &mut buf[1..];
                }
                _ => break,
            }
        }

        Ok(bytes_read)
    }
}

展开为

#![feature(type_alias_impl_trait)]

pub trait ReadDynamic {
    fn poll_read<'a>(
        &'a mut self,
        buf: &'a mut [u8],
        ctx: &mut ::core::task::Context,
    ) -> ::core::task::Poll<Result<usize>>;
}

/// the future returned by [`ReadDynamicExt::read`]
pub struct ReadDynamicRead<'a>(
    &'a mut dyn ReadDynamic,
    &'a mut [u8],
    ::core::marker::PhantomData<&'a ()>,
);

impl<'a> ::core::future::Future for ReadDynamicRead<'a> {
    type Output = Result<usize>;
    fn poll(
        mut self: ::core::pin::Pin<&mut Self>,
        cx: &mut ::core::task::Context,
    ) -> ::core::task::Poll<Self::Output> {
        let this = &mut *self;
        ReadDynamic::poll_read(this.0.into(), this.1.into(), cx)
    }
}

/// the future returned by [`ReadDynamicExt::read_until`]
pub type ReadDynamicReadUntil<'a> = impl ::core::future::Future<Output = Result<usize>>;

pub trait ReadDynamicExt: ReadDynamic {
    fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> ReadDynamicRead<'a>;
    fn read_until<'a>(&'a mut self, byte: u8, buf: &'a mut [u8]) -> ReadDynamicReadUntil<'a>;
}

impl<__IMPL: ReadDynamic> ReadDynamicExt for __IMPL {
    fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> ReadDynamicRead<'a> {
        ReadDynamicRead(self, buf, ::core::marker::PhantomData)
    }

    fn read_until<'a>(&'a mut self, byte: u8, buf: &'a mut [u8]) -> ReadDynamicReadUntil<'a> {
        async fn fn_impl<'a>(
            this: &mut dyn ReadDynamicExt,
            byte: u8,
            mut buf: &'a mut [u8],
        ) -> Result<usize> {
            {
                let mut b = [0];
                let mut bytes_read = 0;
                while !buf.is_empty() {
                    match this.read(&mut b).await? {
                        1 if b[0] != byte => {
                            bytes_read += 1;
                            buf[0] = b[0];
                            buf = &mut buf[1..];
                        }
                        _ => break,
                    }
                }
                Ok(bytes_read)
            }
        }
        fn_impl(self, byte, buf)
    }
}

测试中还包含了一些示例，说明了如何使用这些特性。