breakdown file bandwidth by tag (#30752)

GitOrigin-RevId: afd71770e3d201b3949303863626f5279945731f

Author: rrwang7

crates/application/src/export_worker.rs

@@ -487,6 +487,7 @@ impl<RT: Runtime> ExportWorker<RT> {
                     )
                     .track_storage_egress_size(
                         component_path.clone(),
+                        "snapshot_export".to_string(),
                         file_stream.content_length as u64,
                     );
                 zip_snapshot_upload

crates/application/src/snapshot_import.rs

@@ -2154,7 +2154,11 @@ async fn import_storage_table<RT: Runtime>(
                 content_type,
                 entry.sha256,
             )
-            .track_storage_ingress_size(component_path.clone(), file_size);
+            .track_storage_ingress_size(
+                component_path.clone(),
+                "snapshot_import".to_string(),
+                file_size,
+            );
         num_files += 1;
         if let Some(import_id) = import_id {
             best_effort_update_progress_message(

crates/events/src/usage.rs

@@ -78,6 +78,7 @@ pub enum UsageEvent {
     StorageBandwidth {
         id: String,
         component_path: Option<String>,
+        tag: String,
         ingress: u64,
         egress: u64,
     },

crates/file_storage/src/core.rs

@@ -302,8 +302,11 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
                     if let Ok(ref bytes) = bytes {
                         let bytes_size = bytes.len() as u64;
                         log_get_file_chunk_size(bytes_size, get_file_type);
-                        storage_call_tracker
-                            .track_storage_egress_size(component_path.clone(), bytes_size);
+                        storage_call_tracker.track_storage_egress_size(
+                            component_path.clone(),
+                            "get_range".to_string(),
+                            bytes_size,
+                        );
                     }
                     bytes
                 }),
@@ -448,7 +451,7 @@ impl<RT: Runtime> FileStorage<RT> {
                 content_type,
                 sha256,
             )
-            .track_storage_ingress_size(component_path, size as u64);
+            .track_storage_ingress_size(component_path, "store".to_string(), size as u64);
         Ok(virtual_id)
     }
 }

crates/isolate/src/environment/action/storage.rs

@@ -92,7 +92,7 @@ impl<RT: Runtime> TaskExecutor<RT> {
                 content_type,
                 sha256,
             )
-            .track_storage_ingress_size(component_path, size as u64);
+            .track_storage_ingress_size(component_path, "store".to_string(), size as u64);
 
         Ok(storage_doc_id)
     }

crates/usage_tracking/src/lib.rs

@@ -269,8 +269,18 @@ pub trait StorageUsageTracker: Send + Sync {
 }
 
 pub trait StorageCallTracker: Send + Sync {
-    fn track_storage_ingress_size(&self, component_path: ComponentPath, ingress_size: u64);
-    fn track_storage_egress_size(&self, component_path: ComponentPath, egress_size: u64);
+    fn track_storage_ingress_size(
+        &self,
+        component_path: ComponentPath,
+        tag: String,
+        ingress_size: u64,
+    );
+    fn track_storage_egress_size(
+        &self,
+        component_path: ComponentPath,
+        tag: String,
+        egress_size: u64,
+    );
 }
 
 struct IndependentStorageCallTracker {
@@ -288,21 +298,33 @@ impl IndependentStorageCallTracker {
 }
 
 impl StorageCallTracker for IndependentStorageCallTracker {
-    fn track_storage_ingress_size(&self, component_path: ComponentPath, ingress_size: u64) {
+    fn track_storage_ingress_size(
+        &self,
+        component_path: ComponentPath,
+        tag: String,
+        ingress_size: u64,
+    ) {
         metrics::storage::log_storage_ingress_size(ingress_size);
         self.usage_logger.record(vec![UsageEvent::StorageBandwidth {
             id: self.execution_id.to_string(),
             component_path: component_path.serialize(),
+            tag,
             ingress: ingress_size,
             egress: 0,
         }]);
     }
 
-    fn track_storage_egress_size(&self, component_path: ComponentPath, egress_size: u64) {
+    fn track_storage_egress_size(
+        &self,
+        component_path: ComponentPath,
+        tag: String,
+        egress_size: u64,
+    ) {
         metrics::storage::log_storage_egress_size(egress_size);
         self.usage_logger.record(vec![UsageEvent::StorageBandwidth {
             id: self.execution_id.to_string(),
             component_path: component_path.serialize(),
+            tag,
             ingress: 0,
             egress: egress_size,
         }]);
@@ -492,16 +514,30 @@ impl FunctionUsageTracker {
 // For UDFs, we track storage at the per UDF level, no finer. So we can just
 // aggregate over the entire UDF and not worry about sending usage events or
 // creating unique execution ids.
+// Note: If we want finer-grained breakdown of file bandwidth, we can thread the
+// tag through FunctionUsageStats. For now we're just interested in the
+// breakdown of file bandwidth from functions vs external sources like snapshot
+// export/cloud backups.
 impl StorageCallTracker for FunctionUsageTracker {
-    fn track_storage_ingress_size(&self, component_path: ComponentPath, ingress_size: u64) {
+    fn track_storage_ingress_size(
+        &self,
+        component_path: ComponentPath,
+        _tag: String,
+        ingress_size: u64,
+    ) {
         let mut state = self.state.lock();
         metrics::storage::log_storage_ingress_size(ingress_size);
         state
             .storage_ingress_size
             .mutate_entry_or_default(component_path, |count| *count += ingress_size);
     }
 
-    fn track_storage_egress_size(&self, component_path: ComponentPath, egress_size: u64) {
+    fn track_storage_egress_size(
+        &self,
+        component_path: ComponentPath,
+        _tag: String,
+        egress_size: u64,
+    ) {
         let mut state = self.state.lock();
         metrics::storage::log_storage_egress_size(egress_size);
         state