List/Query cache support for /firestore/response_time tracing attribute

abdolence · Oct 7, 2023 · b25dc95 · b25dc95
1 parent 992df9b
commit b25dc95
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 0 deletions.
diff --git 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
@@ -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
@@ -168,15 +168,27 @@ 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 {
                         format!("{}/{}", self.get_documents_path(), collection_id.as_str())
                     };
 
                     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");