From 26d029656b9c93607f6425708fccaeeae5432254 Mon Sep 17 00:00:00 2001 From: MrCroxx Date: Tue, 24 Sep 2024 06:57:21 +0000 Subject: [PATCH] refactor: replace manual future enum with auto_enum Signed-off-by: MrCroxx --- foyer-storage/Cargo.toml | 1 + foyer-storage/src/engine.rs | 66 ++++++----------------------- foyer-storage/src/storage/either.rs | 56 ++++-------------------- 3 files changed, 22 insertions(+), 101 deletions(-) diff --git a/foyer-storage/Cargo.toml b/foyer-storage/Cargo.toml index e14de4a2..dcb1b2ca 100644 --- a/foyer-storage/Cargo.toml +++ b/foyer-storage/Cargo.toml @@ -18,6 +18,7 @@ anyhow = "1.0" # TODO(MrCroxx): use `array_chunks` after `#![feature(array_chunks)]` is stable. array-util = "1" async-channel = "2" +auto_enums = { version = "0.8", features = ["futures03"] } bincode = "1" bitflags = "2.3.1" bytes = "1" diff --git a/foyer-storage/src/engine.rs b/foyer-storage/src/engine.rs index 33629b3e..eb16abe7 100644 --- a/foyer-storage/src/engine.rs +++ b/foyer-storage/src/engine.rs @@ -12,15 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::{ - fmt::Debug, - marker::PhantomData, - pin::Pin, - sync::Arc, - task::{Context, Poll}, -}; +use std::{fmt::Debug, marker::PhantomData, sync::Arc}; use ahash::RandomState; +use auto_enums::auto_enum; use foyer_common::code::{HashBuilder, StorageKey, StorageValue}; use foyer_memory::CacheEntry; use futures::Future; @@ -96,45 +91,6 @@ where } } -enum StoreFuture { - Noop(F1), - Large(F2), - Small(F3), - Combined(F4), -} - -impl StoreFuture { - pub fn as_pin_mut(self: Pin<&mut Self>) -> StoreFuture, Pin<&mut F2>, Pin<&mut F3>, Pin<&mut F4>> { - unsafe { - match *Pin::get_unchecked_mut(self) { - StoreFuture::Noop(ref mut inner) => StoreFuture::Noop(Pin::new_unchecked(inner)), - StoreFuture::Large(ref mut inner) => StoreFuture::Large(Pin::new_unchecked(inner)), - StoreFuture::Small(ref mut inner) => StoreFuture::Small(Pin::new_unchecked(inner)), - StoreFuture::Combined(ref mut inner) => StoreFuture::Combined(Pin::new_unchecked(inner)), - } - } - } -} - -impl Future for StoreFuture -where - F1: Future, - F2: Future, - F3: Future, - F4: Future, -{ - type Output = F1::Output; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match self.as_pin_mut() { - StoreFuture::Noop(future) => future.poll(cx), - StoreFuture::Large(future) => future.poll(cx), - StoreFuture::Small(future) => future.poll(cx), - StoreFuture::Combined(future) => future.poll(cx), - } - } -} - #[expect(clippy::type_complexity)] pub enum EngineConfig where @@ -251,12 +207,13 @@ where } } + #[auto_enum(Future)] fn load(&self, hash: u64) -> impl Future>> + Send + 'static { match self { - Engine::Noop(storage) => StoreFuture::Noop(storage.load(hash)), - Engine::Large(storage) => StoreFuture::Large(storage.load(hash)), - Engine::Small(storage) => StoreFuture::Small(storage.load(hash)), - Engine::Combined(storage) => StoreFuture::Combined(storage.load(hash)), + Engine::Noop(storage) => storage.load(hash), + Engine::Large(storage) => storage.load(hash), + Engine::Small(storage) => storage.load(hash), + Engine::Combined(storage) => storage.load(hash), } } @@ -296,12 +253,13 @@ where } } + #[auto_enum(Future)] fn wait(&self) -> impl Future + Send + 'static { match self { - Engine::Noop(storage) => StoreFuture::Noop(storage.wait()), - Engine::Large(storage) => StoreFuture::Large(storage.wait()), - Engine::Small(storage) => StoreFuture::Small(storage.wait()), - Engine::Combined(storage) => StoreFuture::Combined(storage.wait()), + Engine::Noop(storage) => storage.wait(), + Engine::Large(storage) => storage.wait(), + Engine::Small(storage) => storage.wait(), + Engine::Combined(storage) => storage.wait(), } } } diff --git a/foyer-storage/src/storage/either.rs b/foyer-storage/src/storage/either.rs index 97d945ec..c6c34c16 100644 --- a/foyer-storage/src/storage/either.rs +++ b/foyer-storage/src/storage/either.rs @@ -12,13 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::{ - fmt::Debug, - pin::Pin, - sync::Arc, - task::{Context, Poll}, -}; +use std::{fmt::Debug, sync::Arc}; +use auto_enums::auto_enum; use foyer_common::code::{HashBuilder, StorageKey, StorageValue}; use foyer_memory::CacheEntry; use futures::{ @@ -28,41 +24,6 @@ use futures::{ use crate::{error::Result, storage::Storage, DeviceStats}; -enum OrderFuture { - LeftFirst(F1), - RightFirst(F2), - Parallel(F3), -} - -impl OrderFuture { - pub fn as_pin_mut(self: Pin<&mut Self>) -> OrderFuture, Pin<&mut F2>, Pin<&mut F3>> { - unsafe { - match *Pin::get_unchecked_mut(self) { - OrderFuture::LeftFirst(ref mut inner) => OrderFuture::LeftFirst(Pin::new_unchecked(inner)), - OrderFuture::RightFirst(ref mut inner) => OrderFuture::RightFirst(Pin::new_unchecked(inner)), - OrderFuture::Parallel(ref mut inner) => OrderFuture::Parallel(Pin::new_unchecked(inner)), - } - } - } -} - -impl Future for OrderFuture -where - F1: Future, - F2: Future, - F3: Future, -{ - type Output = F1::Output; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match self.as_pin_mut() { - OrderFuture::LeftFirst(future) => future.poll(cx), - OrderFuture::RightFirst(future) => future.poll(cx), - OrderFuture::Parallel(future) => future.poll(cx), - } - } -} - /// Order of ops. #[derive(Debug, Clone, Copy)] pub enum Order { @@ -226,24 +187,25 @@ where } } + #[auto_enum(Future)] fn load(&self, hash: u64) -> impl Future>> + Send + 'static { let fleft = self.left.load(hash); let fright = self.right.load(hash); match self.load_order { // FIXME(MrCroxx): false-positive on hash collision. - Order::LeftFirst => OrderFuture::LeftFirst(fleft.then(|res| match res { + Order::LeftFirst => fleft.then(|res| match res { Ok(Some(kv)) => ready(Ok(Some(kv))).left_future(), Err(e) => ready(Err(e)).left_future(), Ok(None) => fright.right_future(), - })), + }), // FIXME(MrCroxx): false-positive on hash collision. - Order::RightFirst => OrderFuture::RightFirst(fright.then(|res| match res { + Order::RightFirst => fright.then(|res| match res { Ok(Some(kv)) => ready(Ok(Some(kv))).left_future(), Err(e) => ready(Err(e)).left_future(), Ok(None) => fleft.right_future(), - })), + }), Order::Parallel => { - OrderFuture::Parallel(async move { + async move { pin_mut!(fleft); pin_mut!(fright); // Returns a 4-way `Either` by nesting `Either` in `Either`. @@ -261,7 +223,7 @@ where }, }) .await - }) + } } } }