From 60d77ab9955dfd6fd17a43a4fd10fd0b28c2b0f3 Mon Sep 17 00:00:00 2001
From: Abdulla Abdurakhmanov <me@abdolence.dev>
Date: Sat, 7 Oct 2023 14:56:27 +0200
Subject: [PATCH] List/Query cache support for /firestore/response_time tracing
 attribute (#131)

---
 src/db/get.rs   | 11 +++++++++++
 src/db/list.rs  | 11 +++++++++++
 src/db/query.rs | 12 ++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/src/db/get.rs b/src/db/get.rs
index fe58f85..e0fbd76 100644
--- a/src/db/get.rs
+++ b/src/db/get.rs
@@ -813,14 +813,25 @@ impl FirestoreDb {
                 "/firestore/collection_name" = collection_id,
                 "/firestore/ids_count" = full_doc_ids.len(),
                 "/firestore/cache_result" = field::Empty,
+                "/firestore/response_time" = field::Empty
             );
 
+            let begin_query_utc: DateTime<Utc> = Utc::now();
+
             let cached_stream: BoxStream<FirestoreResult<(String, Option<FirestoreDocument>)>> =
                 cache.get_docs_by_paths(full_doc_ids).await?;
 
             let cached_vec: Vec<(String, Option<FirestoreDocument>)> =
                 cached_stream.try_collect::<Vec<_>>().await?;
 
+            let end_query_utc: DateTime<Utc> = Utc::now();
+            let query_duration = end_query_utc.signed_duration_since(begin_query_utc);
+
+            span.record(
+                "/firestore/response_time",
+                query_duration.num_milliseconds(),
+            );
+
             if cached_vec.len() == full_doc_ids.len()
                 || matches!(
                     self.session_params.cache_mode,
diff --git a/src/db/list.rs b/src/db/list.rs
index 0348c19..eedfac1 100644
--- a/src/db/list.rs
+++ b/src/db/list.rs
@@ -495,8 +495,11 @@ impl FirestoreDb {
                 "Firestore List Cached",
                 "/firestore/collection_name" = params.collection_id,
                 "/firestore/cache_result" = field::Empty,
+                "/firestore/response_time" = field::Empty
             );
 
+            let begin_query_utc: DateTime<Utc> = Utc::now();
+
             let collection_path = if let Some(parent) = params.parent.as_ref() {
                 format!("{}/{}", parent, params.collection_id.as_str())
             } else {
@@ -509,6 +512,14 @@ impl FirestoreDb {
 
             let cached_result = cache.list_all_docs(&collection_path).await?;
 
+            let end_query_utc: DateTime<Utc> = Utc::now();
+            let query_duration = end_query_utc.signed_duration_since(begin_query_utc);
+
+            span.record(
+                "/firestore/response_time",
+                query_duration.num_milliseconds(),
+            );
+
             match cached_result {
                 FirestoreCachedValue::UseCached(stream) => {
                     span.record("/firestore/cache_result", "hit");
diff --git a/src/db/query.rs b/src/db/query.rs
index 1541a66..f1a04cf 100644
--- a/src/db/query.rs
+++ b/src/db/query.rs
@@ -168,8 +168,11 @@ impl FirestoreDb {
                         "Firestore Query Cached",
                         "/firestore/collection_name" = collection_id.as_str(),
                         "/firestore/cache_result" = field::Empty,
+                        "/firestore/response_time" = field::Empty
                     );
 
+                    let begin_query_utc: DateTime<Utc> = Utc::now();
+
                     let collection_path = if let Some(parent) = params.parent.as_ref() {
                         format!("{}/{}", parent, collection_id)
                     } else {
@@ -177,6 +180,15 @@ impl FirestoreDb {
                     };
 
                     let result = cache.query_docs(&collection_path, params).await?;
+
+                    let end_query_utc: DateTime<Utc> = Utc::now();
+                    let query_duration = end_query_utc.signed_duration_since(begin_query_utc);
+
+                    span.record(
+                        "/firestore/response_time",
+                        query_duration.num_milliseconds(),
+                    );
+
                     match result {
                         FirestoreCachedValue::UseCached(stream) => {
                             span.record("/firestore/cache_result", "hit");