Revert "Revert "make storage egress and ingress usage component aware (#29846)" (#29883)" (#29917)

GitOrigin-RevId: dab054fa2ecc8da07854bab1234d479c17ea9f04

Author: Sujay Jayakar

Cargo.lock

@@ -483,7 +483,10 @@ impl<RT: Runtime> ExportWorker<RT> {
                         content_type,
                         file_storage_entry.sha256.clone(),
                     )
-                    .track_storage_egress_size(file_stream.content_length as u64);
+                    .track_storage_egress_size(
+                        component_path.clone(),
+                        file_stream.content_length as u64,
+                    );
                 zip_snapshot_upload
                     .stream_full_file(path, file_stream.stream)
                     .await?;
@@ -1218,7 +1221,13 @@ mod tests {
 
         let usage = usage.gather_user_stats();
         assert!(usage.database_egress_size.is_empty());
-        assert_eq!(usage.storage_egress_size, 3);
+        assert_eq!(
+            *usage
+                .storage_egress_size
+                .get(&ComponentPath::test_user())
+                .unwrap(),
+            3
+        );
 
         Ok(())
     }

crates/application/src/export_worker.rs

@@ -2096,7 +2096,7 @@ async fn import_storage_table<RT: Runtime>(
                 content_type,
                 entry.sha256,
             )
-            .track_storage_ingress_size(file_size);
+            .track_storage_ingress_size(component_path.clone(), file_size);
         num_files += 1;
         if let Some(import_id) = import_id {
             best_effort_update_progress_message(
@@ -3517,8 +3517,11 @@ a
         .await?;
 
         let stats = usage.gather_user_stats();
-        assert!(stats.database_ingress_size[&(component_path, table_name.to_string())] > 0);
-        assert_eq!(stats.storage_ingress_size, 9);
+        assert!(stats.database_ingress_size[&(component_path.clone(), table_name.to_string())] > 0);
+        assert_eq!(
+            *stats.storage_ingress_size.get(&component_path).unwrap(),
+            9u64
+        );
 
         Ok(())
     }

crates/application/src/snapshot_import.rs

@@ -77,6 +77,7 @@ pub enum UsageEvent {
     /// import/export).
     StorageBandwidth {
         id: String,
+        component_path: Option<String>,
         ingress: u64,
         egress: u64,
     },

crates/events/src/usage.rs

@@ -243,8 +243,9 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
         let stream = storage_get_stream.stream;
         let content_length = ContentLength(storage_get_stream.content_length as u64);
 
+        let component_path = ComponentPath::TODO();
         let call_tracker = usage_tracker.track_storage_call(
-            ComponentPath::TODO(),
+            component_path.clone(),
             "get range",
             storage_id,
             content_type.clone(),
@@ -255,11 +256,12 @@ impl<RT: Runtime> TransactionalFileStorage<RT> {
             content_length,
             content_range,
             content_type,
-            stream: Self::track_stream_usage(stream, get_file_type, call_tracker),
+            stream: Self::track_stream_usage(component_path, stream, get_file_type, call_tracker),
         })
     }
 
     fn track_stream_usage(
+        component_path: ComponentPath,
         stream: BoxStream<'static, futures::io::Result<bytes::Bytes>>,
         get_file_type: GetFileType,
         storage_call_tracker: Box<dyn StorageCallTracker>,
@@ -291,7 +293,8 @@ 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(bytes_size);
+                        storage_call_tracker
+                            .track_storage_egress_size(component_path.clone(), bytes_size);
                     }
                     bytes
                 }),
@@ -429,8 +432,14 @@ impl<RT: Runtime> FileStorage<RT> {
             .await?;
 
         usage_tracker
-            .track_storage_call(component_path, "store", storage_id, content_type, sha256)
-            .track_storage_ingress_size(size as u64);
+            .track_storage_call(
+                component_path.clone(),
+                "store",
+                storage_id,
+                content_type,
+                sha256,
+            )
+            .track_storage_ingress_size(component_path, size as u64);
         Ok(virtual_id)
     }
 }

crates/file_storage/src/core.rs

@@ -80,15 +80,16 @@ impl<RT: Runtime> TaskExecutor<RT> {
             .storage_store_file_entry(self.identity.clone(), self.component_id(), entry)
             .await?;
 
+        let component_path = ComponentPath::TODO();
         self.usage_tracker
             .track_storage_call(
-                ComponentPath::TODO(),
+                component_path.clone(),
                 "store",
                 storage_id,
                 content_type,
                 sha256,
             )
-            .track_storage_ingress_size(size as u64);
+            .track_storage_ingress_size(component_path, size as u64);
 
         Ok(storage_doc_id)
     }

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

@@ -4,8 +4,12 @@ package usage;
 
 message FunctionUsageStats {
     repeated CounterWithTag storage_calls = 1;
+    // TODO(ENG-7342): Remove after services are pushed with by_component version
     optional uint64 storage_ingress_size = 2;
+    repeated CounterWithComponent storage_ingress_size_by_component = 8;
+    // TODO(ENG-7342): Remove after services are pushed with by_component version
     optional uint64 storage_egress_size = 3;
+    repeated CounterWithComponent storage_egress_size_by_component = 9;
     repeated CounterWithTag database_ingress_size = 4;
     repeated CounterWithTag database_egress_size = 5;
     repeated CounterWithTag vector_ingress_size = 6;
@@ -17,3 +21,8 @@ message CounterWithTag {
     optional string table_name = 1;
     optional uint64 count = 2;
 }
+
+message CounterWithComponent {
+    optional string component_path = 1;
+    optional uint64 count = 2;
+}

crates/pb/protos/usage.proto

@@ -13,6 +13,7 @@ anyhow = { workspace = true }
 common = { path = "../common" }
 events = { path = "../events" }
 headers = { workspace = true }
+maplit = { workspace = true }
 metrics = { path = "../metrics" }
 parking_lot = { workspace = true, features = ["hardware-lock-elision"] }
 pb = { path = "../pb" }