feat(tower)!: predicates refactoring

hit-box · Jul 19, 2023 · 5d94f26 · 5d94f26
1 parent 5f547b9
commit 5d94f26
Show file tree

Hide file tree

Showing 18 changed files with 352 additions and 341 deletions.
diff --git a/examples/examples/axum.rs b/examples/examples/axum.rs
@@ -51,7 +51,7 @@ async fn main() {
         )
         .layer(
             ServiceBuilder::new()
-                .layer(Cache::builder().backend(inmemory).build())
+                // .layer(Cache::builder().backend(inmemory).build())
                 .layer(Cache::builder().backend(backend).build()),
         );
 

diff --git a/examples/examples/tower.rs b/examples/examples/tower.rs
@@ -12,6 +12,12 @@ async fn handle(_: Request<Body>) -> Result<Response<Body>, Infallible> {
 
 #[tokio::main]
 async fn main() {
+    let subscriber = tracing_subscriber::fmt()
+        .pretty()
+        .with_env_filter("debug,hitbox=trace")
+        .finish();
+    tracing::subscriber::set_global_default(subscriber).unwrap();
+
     let inmemory = StrettoBackendBuilder::new(12960, 1e6 as i64)
         .finalize()
         .unwrap();

diff --git a/hitbox-backend/src/predicates.rs b/hitbox-backend/src/predicates.rs
@@ -1,53 +1,53 @@
+use std::sync::Arc;
+
 use async_trait::async_trait;
 
 pub enum PredicateResult<S> {
     Cacheable(S),
     NonCacheable(S),
 }
 
-impl<S> PredicateResult<S>
-where
-    S: Send + 'static,
-{
-    pub async fn chain<P: Predicate<S> + ?Sized>(self, predicate: &Box<P>) -> PredicateResult<S> {
-        dbg!("PredicateResult::chain");
-        match self {
-            PredicateResult::NonCacheable(subject) => predicate.check(subject).await,
-            PredicateResult::Cacheable(subject) => PredicateResult::Cacheable(subject),
-        }
-    }
-}
+// impl<S> PredicateResult<S>
+// where
+//     S: Send + 'static,
+// {
+//     pub async fn chain<P: Predicate + ?Sized>(self, predicate: &Box<P>) -> PredicateResult<S> {
+//         dbg!("PredicateResult::chain");
+//         match self {
+//             PredicateResult::NonCacheable(subject) => predicate.check(subject).await,
+//             PredicateResult::Cacheable(subject) => PredicateResult::Cacheable(subject),
+//         }
+//     }
+// }
 
 #[async_trait]
-pub trait Predicate<S>: Send {
-    async fn check(&self, subject: S) -> PredicateResult<S>;
-}
-
-pub struct NeutralPredicate;
-
-impl NeutralPredicate {
-    pub fn new() -> Self {
-        NeutralPredicate
-    }
+pub trait Predicate {
+    type Subject;
+    async fn check(&self, subject: Self::Subject) -> PredicateResult<Self::Subject>;
 }
 
 #[async_trait]
-impl<S> Predicate<S> for NeutralPredicate
+impl<T> Predicate for Box<T>
 where
-    S: Send + 'static,
+    T: Predicate + ?Sized + Sync,
+    T::Subject: Send,
 {
-    async fn check(&self, subject: S) -> PredicateResult<S> {
-        PredicateResult::Cacheable(subject)
+    type Subject = T::Subject;
+
+    async fn check(&self, subject: T::Subject) -> PredicateResult<T::Subject> {
+        self.as_ref().check(subject).await
     }
 }
 
 #[async_trait]
-impl<S, T> Predicate<S> for Box<T>
+impl<T> Predicate for Arc<T>
 where
-    T: Predicate<S> + Sync,
-    S: Send + 'static,
+    T: Predicate + Send + Sync + ?Sized,
+    T::Subject: Send,
 {
-    async fn check(&self, subject: S) -> PredicateResult<S> {
+    type Subject = T::Subject;
+
+    async fn check(&self, subject: T::Subject) -> PredicateResult<T::Subject> {
         self.as_ref().check(subject).await
     }
 }

diff --git a/hitbox-backend/src/response.rs b/hitbox-backend/src/response.rs
@@ -34,17 +34,16 @@ where
 {
     type Cached;
 
-    async fn cache_policy(self, predicates: impl Predicate<Self> + Send) -> CachePolicy<Self> {
-        unimplemented!();
-        // let predicates_result = stream::iter(predicates)
-        //     .fold(PredicateResult::Cacheable(self), PredicateResult::chain)
-        //     .await;
-        // match predicates_result {
-        //     PredicateResult::Cacheable(res) => {
-        //         CachePolicy::Cacheable(CachedValue::new(res.into_cached().await, Utc::now()))
-        //     }
-        //     PredicateResult::NonCacheable(res) => CachePolicy::NonCacheable(res),
-        // }
+    async fn cache_policy<P>(self, predicates: P) -> CachePolicy<Self>
+    where
+        P: Predicate<Subject = Self> + Send + Sync,
+    {
+        match predicates.check(self).await {
+            PredicateResult::Cacheable(res) => {
+                CachePolicy::Cacheable(CachedValue::new(res.into_cached().await, Utc::now()))
+            }
+            PredicateResult::NonCacheable(res) => CachePolicy::NonCacheable(res),
+        }
     }
 
     async fn into_cached(self) -> Self::Cached;

diff --git a/hitbox-http/src/predicates/body.rs b/hitbox-http/src/predicates/body.rs
@@ -13,34 +13,33 @@ pub struct Body<P> {
     inner: P,
 }
 
-pub trait BodyPredicate<ReqBody>: Sized {
+pub trait BodyPredicate: Sized {
     fn body(self) -> Body<Self>;
 }
 
-impl<P, ReqBody> BodyPredicate<ReqBody> for P
+impl<P> BodyPredicate for P
 where
-    P: Predicate<CacheableHttpRequest<ReqBody>>,
+    P: Predicate,
 {
     fn body(self) -> Body<Self> {
         Body { inner: self }
     }
 }
 
 #[async_trait]
-impl<P, ReqBody> Predicate<CacheableHttpRequest<ReqBody>> for Body<P>
+impl<P, ReqBody> Predicate for Body<P>
 where
     ReqBody: HttpBody + Send + 'static,
-    P: Predicate<CacheableHttpRequest<ReqBody>> + Send + Sync,
+    P: Predicate<Subject = CacheableHttpRequest<ReqBody>> + Send + Sync,
 
     // debug bounds
     ReqBody::Error: Debug,
     ReqBody::Data: Send,
     ReqBody: FromBytes,
 {
-    async fn check(
-        &self,
-        request: CacheableHttpRequest<ReqBody>,
-    ) -> PredicateResult<CacheableHttpRequest<ReqBody>> {
+    type Subject = P::Subject;
+
+    async fn check(&self, request: Self::Subject) -> PredicateResult<Self::Subject> {
         match self.inner.check(request).await {
             PredicateResult::Cacheable(request) => {
                 let (parts, body) = request.into_parts();

diff --git a/hitbox-http/src/predicates/header.rs b/hitbox-http/src/predicates/header.rs
@@ -2,34 +2,56 @@ use crate::CacheableHttpRequest;
 use async_trait::async_trait;
 use hitbox::predicates::{Operation, Predicate, PredicateResult};
 
-pub struct Header {
+pub struct Header<P> {
     pub name: String,
     pub value: String,
     pub operation: Operation,
+    inner: P,
+}
+
+pub trait HeaderPredicate: Sized {
+    fn header(self, name: String, value: String) -> Header<Self>;
+}
+
+impl<P> HeaderPredicate for P
+where
+    P: Predicate,
+{
+    fn header(self, name: String, value: String) -> Header<P> {
+        Header {
+            name,
+            value,
+            operation: Operation::Eq,
+            inner: self,
+        }
+    }
 }
 
 #[async_trait]
-impl<ReqBody> Predicate<CacheableHttpRequest<ReqBody>> for Header
+impl<P, ReqBody> Predicate for Header<P>
 where
     ReqBody: Send + 'static,
+    P: Predicate<Subject = CacheableHttpRequest<ReqBody>> + Send + Sync,
 {
-    async fn check(
-        &self,
-        request: CacheableHttpRequest<ReqBody>,
-    ) -> PredicateResult<CacheableHttpRequest<ReqBody>> {
+    type Subject = P::Subject;
+
+    async fn check(&self, request: Self::Subject) -> PredicateResult<Self::Subject> {
         dbg!("HeaderPredicate::check");
-        match self.operation {
-            Operation::Eq => match request.parts().headers.get(self.name.as_str()) {
-                Some(header_value) => {
-                    if self.value.as_str() == header_value {
-                        PredicateResult::Cacheable(request)
-                    } else {
-                        PredicateResult::NonCacheable(request)
+        match self.inner.check(request).await {
+            PredicateResult::Cacheable(request) => match self.operation {
+                Operation::Eq => match request.parts().headers.get(self.name.as_str()) {
+                    Some(header_value) => {
+                        if self.value.as_str() == header_value {
+                            PredicateResult::Cacheable(request)
+                        } else {
+                            PredicateResult::NonCacheable(request)
+                        }
                     }
-                }
-                None => PredicateResult::NonCacheable(request),
+                    None => PredicateResult::NonCacheable(request),
+                },
+                _ => unimplemented!(),
             },
-            _ => unimplemented!(),
+            PredicateResult::NonCacheable(request) => PredicateResult::NonCacheable(request),
         }
     }
 }
diff --git a/hitbox-http/src/predicates/mod.rs b/hitbox-http/src/predicates/mod.rs
@@ -1,30 +1,58 @@
+use std::{
+    marker::PhantomData,
+    sync::{atomic::AtomicPtr, Arc, RwLock},
+};
+
 use async_trait::async_trait;
 use hitbox_backend::predicates::{Predicate, PredicateResult};
 
-use crate::CacheableHttpRequest;
+use crate::{CacheableHttpRequest, CacheableHttpResponse};
 
 pub mod body;
 pub mod header;
 pub mod path;
 pub mod query;
 
-pub struct NeutralPredicate;
+pub struct NeutralPredicate<ReqBody> {
+    _req: PhantomData<AtomicPtr<Box<ReqBody>>>, // FIX: NOT HEHE
+}
 
-impl NeutralPredicate {
+impl<ReqBody> NeutralPredicate<ReqBody> {
     pub fn new() -> Self {
-        NeutralPredicate
+        NeutralPredicate { _req: PhantomData }
     }
 }
 
 #[async_trait]
-impl<ReqBody> Predicate<CacheableHttpRequest<ReqBody>> for NeutralPredicate
+impl<ReqBody> Predicate for NeutralPredicate<ReqBody>
 where
     ReqBody: Send + 'static,
 {
-    async fn check(
-        &self,
-        subject: CacheableHttpRequest<ReqBody>,
-    ) -> PredicateResult<CacheableHttpRequest<ReqBody>> {
+    type Subject = CacheableHttpRequest<ReqBody>;
+
+    async fn check(&self, subject: Self::Subject) -> PredicateResult<Self::Subject> {
+        PredicateResult::Cacheable(subject)
+    }
+}
+
+pub struct NeutralResponsePredicate<ResBody> {
+    _res: PhantomData<fn() -> ResBody>, // FIX: HEHE
+}
+
+impl<ResBody> NeutralResponsePredicate<ResBody> {
+    pub fn new() -> Self {
+        NeutralResponsePredicate { _res: PhantomData }
+    }
+}
+
+#[async_trait]
+impl<ResBody> Predicate for NeutralResponsePredicate<ResBody>
+where
+    ResBody: Send + 'static,
+{
+    type Subject = CacheableHttpResponse<ResBody>;
+
+    async fn check(&self, subject: Self::Subject) -> PredicateResult<Self::Subject> {
         PredicateResult::Cacheable(subject)
     }
 }
diff --git a/hitbox-http/src/predicates/path.rs b/hitbox-http/src/predicates/path.rs
@@ -10,19 +10,13 @@ pub struct Path<P> {
     inner: P,
 }
 
-// impl<P, ReqBody> Path<P, ReqBody> {
-//     pub fn new(value: String) -> Self {
-//         Self(ResourceDef::new(value))
-//     }
-// }
-
-pub trait PathPredicate<ReqBody>: Sized {
+pub trait PathPredicate: Sized {
     fn path(self, resource: ResourceDef) -> Path<Self>;
 }
 
-impl<P, ReqBody> PathPredicate<ReqBody> for P
+impl<P> PathPredicate for P
 where
-    P: Predicate<CacheableHttpRequest<ReqBody>>,
+    P: Predicate,
 {
     fn path(self, resource: ResourceDef) -> Path<Self> {
         Path {
@@ -33,23 +27,25 @@ where
 }
 
 #[async_trait]
-impl<P, ReqBody> Predicate<CacheableHttpRequest<ReqBody>> for Path<P>
+impl<P, ReqBody> Predicate for Path<P>
 where
-    P: Predicate<CacheableHttpRequest<ReqBody>> + Sync,
+    P: Predicate<Subject = CacheableHttpRequest<ReqBody>> + Send + Sync,
     ReqBody: Send + 'static,
 {
-    async fn check(
-        &self,
-        request: CacheableHttpRequest<ReqBody>,
-    ) -> PredicateResult<CacheableHttpRequest<ReqBody>> {
+    type Subject = P::Subject;
+
+    async fn check(&self, request: Self::Subject) -> PredicateResult<Self::Subject> {
         match self.inner.check(request).await {
-            PredicateResult::Cacheable(request) => {
-                if self.resource.is_match(request.parts().uri.path()) {
-                    PredicateResult::Cacheable(request)
-                } else {
-                    PredicateResult::NonCacheable(request)
+            PredicateResult::Cacheable(request) => match self.inner.check(request).await {
+                PredicateResult::Cacheable(request) => {
+                    if self.resource.is_match(request.parts().uri.path()) {
+                        PredicateResult::Cacheable(request)
+                    } else {
+                        PredicateResult::NonCacheable(request)
+                    }
                 }
-            }
+                PredicateResult::NonCacheable(request) => PredicateResult::NonCacheable(request),
+            },
             PredicateResult::NonCacheable(request) => PredicateResult::NonCacheable(request),
         }
     }