Log sha256 on storage serving calls (#29596)

GitOrigin-RevId: ee27aeabe7c74874ae0dae7098f16064739993ff
get-convex · Sep 5, 2024 · c3c1be0 · c3c1be0
1 parent bc8009c
commit c3c1be0
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 6 deletions.
diff --git a/crates/application/src/export_worker.rs b/crates/application/src/export_worker.rs
@@ -519,6 +519,7 @@ impl<RT: Runtime> ExportWorker<RT> {
                         "snapshot_export",
                         file_storage_entry.storage_id.clone(),
                         content_type,
+                        file_storage_entry.sha256.clone(),
                     )
                     .track_storage_egress_size(file_stream.content_length as u64);
                 zip_snapshot_upload

diff --git a/crates/application/src/snapshot_import.rs b/crates/application/src/snapshot_import.rs
@@ -2089,7 +2089,12 @@ async fn import_storage_table<RT: Runtime>(
             .map(|ct| ct.parse())
             .transpose()?;
         usage
-            .track_storage_call("snapshot_import", entry.storage_id, content_type)
+            .track_storage_call(
+                "snapshot_import",
+                entry.storage_id,
+                content_type,
+                entry.sha256,
+            )
             .track_storage_ingress_size(file_size);
         num_files += 1;
         if let Some(import_id) = import_id {

diff --git a/crates/events/src/usage.rs b/crates/events/src/usage.rs
@@ -67,6 +67,7 @@ pub enum UsageEvent {
         storage_id: String,
         call: String,
         content_type: Option<String>,
+        sha256: String,
     },
     /// Bandwidth from a storage call outside of a user function (e.g. snapshot
     /// import/export).

diff --git a/crates/file_storage/src/core.rs b/crates/file_storage/src/core.rs
@@ -224,7 +224,7 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
         let FileStorageEntry {
             storage_id,
             storage_key,
-            sha256: _,
+            sha256,
             size,
             content_type,
         } = file;
@@ -241,7 +241,7 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
         let content_length = ContentLength(storage_get_stream.content_length as u64);
 
         let call_tracker =
-            usage_tracker.track_storage_call("get range", storage_id, content_type.clone());
+            usage_tracker.track_storage_call("get range", storage_id, content_type.clone(), sha256);
 
         Ok(FileRangeStream {
             content_length,
@@ -406,6 +406,7 @@ impl<RT: Runtime> FileStorage<RT> {
             .as_ref()
             .map(|ct| ct.parse())
             .transpose()?;
+        let sha256 = entry.sha256.clone();
 
         // Start/Complete transaction after the slow upload process
         // to avoid OCC risk.
@@ -419,7 +420,7 @@ impl<RT: Runtime> FileStorage<RT> {
             .await?;
 
         usage_tracker
-            .track_storage_call("store", storage_id, content_type)
+            .track_storage_call("store", storage_id, content_type, sha256)
             .track_storage_ingress_size(size as u64);
         Ok(virtual_id)
     }

diff --git a/crates/isolate/src/environment/action/storage.rs b/crates/isolate/src/environment/action/storage.rs
@@ -73,13 +73,14 @@ impl<RT: Runtime> TaskExecutor<RT> {
             .await?;
         let storage_id = entry.storage_id.clone();
         let size = entry.size;
+        let sha256 = entry.sha256.clone();
         let storage_doc_id = self
             .action_callbacks
             .storage_store_file_entry(self.identity.clone(), self.component_id(), entry)
             .await?;
 
         self.usage_tracker
-            .track_storage_call("store", storage_id, content_type)
+            .track_storage_call("store", storage_id, content_type, sha256)
             .track_storage_ingress_size(size as u64);
 
         Ok(storage_doc_id)

diff --git a/crates/usage_tracking/src/lib.rs b/crates/usage_tracking/src/lib.rs
@@ -27,7 +27,10 @@ use pb::usage::{
     CounterWithTag as CounterWithTagProto,
     FunctionUsageStats as FunctionUsageStatsProto,
 };
-use value::heap_size::WithHeapSize;
+use value::{
+    heap_size::WithHeapSize,
+    sha256::Sha256Digest,
+};
 
 mod metrics;
 
@@ -232,6 +235,7 @@ pub trait StorageUsageTracker: Send + Sync {
         storage_api: &'static str,
         storage_id: StorageUuid,
         content_type: Option<ContentType>,
+        sha256: Sha256Digest,
     ) -> Box<dyn StorageCallTracker>;
 }
 
@@ -280,6 +284,7 @@ impl StorageUsageTracker for UsageCounter {
         storage_api: &'static str,
         storage_id: StorageUuid,
         content_type: Option<ContentType>,
+        sha256: Sha256Digest,
     ) -> Box<dyn StorageCallTracker> {
         let execution_id = ExecutionId::new();
         metrics::storage::log_storage_call();
@@ -290,6 +295,7 @@ impl StorageUsageTracker for UsageCounter {
             storage_id: storage_id.to_string(),
             call: storage_api.to_string(),
             content_type: content_type.map(|c| c.to_string()),
+            sha256: sha256.as_hex(),
         }]);
 
         Box::new(IndependentStorageCallTracker::new(
@@ -467,6 +473,7 @@ impl StorageUsageTracker for FunctionUsageTracker {
         storage_api: &'static str,
         _storage_id: StorageUuid,
         _content_type: Option<ContentType>,
+        _sha256: Sha256Digest,
     ) -> Box<dyn StorageCallTracker> {
         let mut state = self.state.lock();
         metrics::storage::log_storage_call();