From 1e6b528e4b8042a80842226d40b7da2e4175f0d7 Mon Sep 17 00:00:00 2001
From: ibraheemdev <ibrah1440@gmail.com>
Date: Sun, 27 Jun 2021 21:33:59 -0400
Subject: [PATCH] add lazy data extractor

---
 Cargo.toml  |  1 +
 src/data.rs | 81 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 src/web.rs  |  2 +-
 3 files changed, 82 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 7556bd8d7..3933e6b3f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -101,6 +101,7 @@ smallvec = "1.6"
 socket2 = "0.4.0"
 time = { version = "0.2.23", default-features = false, features = ["std"] }
 url = "2.1"
+tokio = { version = "1.7.1", features = ["sync"] }
 
 [dev-dependencies]
 actix-test = { version = "0.1.0-beta.3", features = ["openssl", "rustls"] }
diff --git a/src/data.rs b/src/data.rs
index 174faba37..c1533106e 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -1,9 +1,11 @@
-use std::{any::type_name, ops::Deref, sync::Arc};
+use std::{any::type_name, cell::Cell, fmt, future::Future, ops::Deref, rc::Rc, sync::Arc};
 
 use actix_http::Extensions;
 use actix_utils::future::{err, ok, Ready};
 use futures_core::future::LocalBoxFuture;
+use futures_util::FutureExt;
 use serde::Serialize;
+use tokio::sync::OnceCell;
 
 use crate::{
     dev::Payload, error::ErrorInternalServerError, extract::FromRequest, request::HttpRequest,
@@ -149,6 +151,83 @@ impl<T: ?Sized + 'static> DataFactory for Data<T> {
     }
 }
 
+/// A lazy extractor for thread-local data.
+pub struct LazyData<T> {
+    inner: Rc<LazyDataInner<T>>,
+}
+
+pub struct LazyDataInner<T> {
+    cell: OnceCell<T>,
+    fut: Cell<Option<LocalBoxFuture<'static, T>>>,
+}
+
+impl<T> Clone for LazyData<T> {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+        }
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for LazyData<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Lazy")
+            .field("cell", &self.inner.cell)
+            .field("fut", &"..")
+            .finish()
+    }
+}
+
+impl<T> LazyData<T> {
+    pub fn new<F, Fut>(init: F) -> LazyData<T>
+    where
+        F: FnOnce() -> Fut,
+        Fut: Future<Output = T> + 'static,
+    {
+        Self {
+            inner: Rc::new(LazyDataInner {
+                cell: OnceCell::new(),
+                fut: Cell::new(Some(init().boxed_local())),
+            }),
+        }
+    }
+
+    pub async fn get(&self) -> &T {
+        self.inner
+            .cell
+            .get_or_init(|| async move {
+                match self.inner.fut.take() {
+                    Some(fut) => fut.await,
+                    None => panic!("LazyData instance has previously been poisoned"),
+                }
+            })
+            .await
+    }
+}
+
+impl<T: 'static> FromRequest for LazyData<T> {
+    type Config = ();
+    type Error = Error;
+    type Future = Ready<Result<Self, Error>>;
+
+    #[inline]
+    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
+        if let Some(lazy) = req.app_data::<LazyData<T>>() {
+            ok(lazy.clone())
+        } else {
+            log::debug!(
+                "Failed to construct App-level LazyData extractor. \
+                 Request path: {:?} (type: {})",
+                req.path(),
+                type_name::<T>(),
+            );
+            err(ErrorInternalServerError(
+                "App data is not configured, to configure use App::data()",
+            ))
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/web.rs b/src/web.rs
index 40ac46275..1fc9b5096 100644
--- a/src/web.rs
+++ b/src/web.rs
@@ -17,7 +17,7 @@ use crate::scope::Scope;
 use crate::service::WebService;
 
 pub use crate::config::ServiceConfig;
-pub use crate::data::Data;
+pub use crate::data::{Data, LazyData};
 pub use crate::request::HttpRequest;
 pub use crate::request_data::ReqData;
 pub use crate::types::*;