From 5b3651910a77f9e30b1f0b1ce8e5b426308d2568 Mon Sep 17 00:00:00 2001
From: Xuanwo <github@xuanwo.io>
Date: Tue, 27 Aug 2024 21:08:05 +0800
Subject: [PATCH] PoC: Integration to tower

Signed-off-by: Xuanwo <github@xuanwo.io>
---
 Cargo.toml                |  2 +-
 backon-tower/Cargo.toml   | 15 ++++++++++++++
 backon-tower/src/layer.rs | 18 ++++++++++++++++
 backon-tower/src/lib.rs   |  6 ++++++
 backon-tower/src/retry.rs | 43 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100644 backon-tower/Cargo.toml
 create mode 100644 backon-tower/src/layer.rs
 create mode 100644 backon-tower/src/lib.rs
 create mode 100644 backon-tower/src/retry.rs

diff --git a/Cargo.toml b/Cargo.toml
index 56b1302..f037133 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,5 +1,5 @@
 [workspace]
-members = ["backon"]
+members = ["backon", "backon-*"]
 resolver = "2"
 
 [workspace.package]
diff --git a/backon-tower/Cargo.toml b/backon-tower/Cargo.toml
new file mode 100644
index 0000000..632b06d
--- /dev/null
+++ b/backon-tower/Cargo.toml
@@ -0,0 +1,15 @@
+[package]
+description = "RetryLayer for Tower based on BackON."
+documentation = "https://docs.rs/backon-tower"
+name = "backon-tower"
+rust-version = "1.70"
+version = "0.1.0"
+
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+backon = "1"
+tower = { git = "https://github.com/tower-rs/tower", rev = "7155101" }
+# tower = "0.5"
diff --git a/backon-tower/src/layer.rs b/backon-tower/src/layer.rs
new file mode 100644
index 0000000..650f2a0
--- /dev/null
+++ b/backon-tower/src/layer.rs
@@ -0,0 +1,18 @@
+use crate::Retry;
+use backon::BackoffBuilder;
+use tower::Layer;
+
+pub struct RetryLayer<B: BackoffBuilder + Clone> {
+    backoff: B,
+}
+
+impl<B, S> Layer<S> for RetryLayer<B>
+where
+    B: BackoffBuilder + Clone,
+{
+    type Service = Retry<B, S>;
+
+    fn layer(&self, service: S) -> Self::Service {
+        Retry::new(self.backoff.clone(), service)
+    }
+}
diff --git a/backon-tower/src/lib.rs b/backon-tower/src/lib.rs
new file mode 100644
index 0000000..b3c076d
--- /dev/null
+++ b/backon-tower/src/lib.rs
@@ -0,0 +1,6 @@
+#![feature(impl_trait_in_assoc_type)]
+
+mod layer;
+pub use layer::RetryLayer;
+mod retry;
+pub use self::retry::Retry;
diff --git a/backon-tower/src/retry.rs b/backon-tower/src/retry.rs
new file mode 100644
index 0000000..13430bb
--- /dev/null
+++ b/backon-tower/src/retry.rs
@@ -0,0 +1,43 @@
+use backon::{BackoffBuilder, Retryable};
+use std::future::Future;
+use tower::Service;
+
+pub struct Retry<B: BackoffBuilder + Clone, S> {
+    backoff: B,
+    service: S,
+}
+
+impl<B: BackoffBuilder + Clone, S> Retry<B, S> {
+    pub fn new(backoff: B, service: S) -> Self {
+        Retry { backoff, service }
+    }
+}
+
+impl<B, S, Request> Service<Request> for Retry<B, S>
+where
+    Request: Clone,
+    B: BackoffBuilder + Clone,
+    S: Service<Request> + Clone,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = impl Future<Output = Result<S::Response, S::Error>>;
+
+    fn poll_ready(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), Self::Error>> {
+        // NOTE: the Future::poll impl for ResponseFuture assumes that Retry::poll_ready is
+        // equivalent to Ready.service.poll_ready. If this ever changes, that code must be updated
+        // as well.
+        self.service.poll_ready(cx)
+    }
+
+    fn call(&mut self, request: Request) -> Self::Future {
+        let req = request.clone();
+        let mut srv = self.service.clone();
+        let backoff = self.backoff.clone();
+
+        (move || srv.call(req.clone())).retry(backoff.clone())
+    }
+}