Move blob fee tracking into execute_operation.

linera-base/src/identifiers.rs

@@ -768,6 +768,23 @@ impl ModuleId {
     }
 }
 
+impl<Abi, Parameters, InstantiationArgument> ModuleId<Abi, Parameters, InstantiationArgument> {
+    /// Returns the ID of the contract bytecode blob.
+    pub fn contract_blob_id(&self) -> BlobId {
+        BlobId::new(self.contract_blob_hash, BlobType::ContractBytecode)
+    }
+
+    /// Returns the ID of the service bytecode blob.
+    pub fn service_blob_id(&self) -> BlobId {
+        BlobId::new(self.service_blob_hash, BlobType::ServiceBytecode)
+    }
+
+    /// Returns the contract and service blob ID.
+    pub fn bytecode_blob_ids(&self) -> [BlobId; 2] {
+        [self.contract_blob_id(), self.service_blob_id()]
+    }
+}
+
 impl<Abi, Parameters, InstantiationArgument> ModuleId<Abi, Parameters, InstantiationArgument> {
     /// Forgets the ABI of a module ID (if any).
     pub fn forget_abi(self) -> ModuleId {

linera-chain/src/chain.rs

@@ -740,12 +740,12 @@ where
             tracker: ResourceTracker::default(),
             account: block.authenticated_signer,
         };
+        let blob_sizes = published_blobs
+            .iter()
+            .map(|blob| (blob.id(), blob.bytes().len()))
+            .collect::<BTreeMap<_, _>>();
         ensure!(
-            block.published_blob_ids()
-                == published_blobs
-                    .iter()
-                    .map(|blob| blob.id())
-                    .collect::<BTreeSet<_>>(),
+            blob_sizes.keys().eq(&block.published_blob_ids()),
             ChainError::InternalError("published_blobs mismatch".to_string())
         );
         resource_controller
@@ -758,17 +758,10 @@ where
             .track_executed_block_size_sequence_extension(0, block.operations.len())
             .with_execution_context(ChainExecutionContext::Block)?;
         for blob in published_blobs {
-            let blob_type = blob.content().blob_type();
-            if blob_type == BlobType::Data
-                || blob_type == BlobType::ContractBytecode
-                || blob_type == BlobType::ServiceBytecode
-            {
-                resource_controller
-                    .with_state(&mut self.execution_state.system)
-                    .await?
-                    .track_blob_published(blob.content())
-                    .with_execution_context(ChainExecutionContext::Block)?;
-            }
+            resource_controller
+                .policy
+                .check_blob_size(blob.content())
+                .with_execution_context(ChainExecutionContext::Block)?;
             self.execution_state.system.used_blobs.insert(&blob.id())?;
         }
 
@@ -843,6 +836,7 @@ where
                         context,
                         local_time,
                         operation.clone(),
+                        &blob_sizes,
                         &mut txn_tracker,
                         &mut resource_controller,
                     ))
@@ -893,11 +887,15 @@ where
                 ))
                 .with_execution_context(chain_execution_context)?;
             for blob in &txn_outcome.blobs {
+                resource_controller
+                    .policy
+                    .check_blob_size(blob.content())
+                    .with_execution_context(chain_execution_context)?;
                 if blob.content().blob_type() == BlobType::Data {
                     resource_controller
                         .with_state(&mut self.execution_state.system)
                         .await?
-                        .track_blob_published(blob.content())
+                        .track_blob_published(blob.bytes().len())
                         .with_execution_context(chain_execution_context)?;
                 }
             }

linera-core/src/unit_tests/wasm_client_tests.rs

@@ -20,7 +20,7 @@ use async_graphql::Request;
 use counter::CounterAbi;
 use linera_base::{
     data_types::{Amount, Bytecode, Event, OracleResponse},
-    identifiers::{AccountOwner, ApplicationId, BlobId, BlobType, StreamId, StreamName},
+    identifiers::{AccountOwner, ApplicationId, StreamId, StreamName},
     ownership::{ChainOwnership, TimeoutConfig},
     vm::VmRuntime,
 };
@@ -895,10 +895,7 @@ async fn test_memory_fuel_limit(wasm_runtime: WasmRuntime) -> anyhow::Result<()>
     let module_id = module_id.with_abi::<counter::CounterAbi, (), u64>();
     let mut blobs = publisher
         .storage_client()
-        .read_blobs(&[
-            BlobId::new(module_id.contract_blob_hash, BlobType::ContractBytecode),
-            BlobId::new(module_id.service_blob_hash, BlobType::ServiceBytecode),
-        ])
+        .read_blobs(&module_id.bytecode_blob_ids())
         .await?
         .into_iter()
         .flatten();

linera-core/src/unit_tests/wasm_worker_tests.rs

@@ -287,6 +287,7 @@ where
                 application_id,
                 bytes: user_operation,
             },
+            &BTreeMap::new(),
             &mut TransactionTracker::new(
                 0,
                 0,

linera-execution/src/execution.rs

@@ -1,12 +1,12 @@
 // Copyright (c) Zefchain Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::{mem, vec};
+use std::{collections::BTreeMap, mem, vec};
 
 use futures::{FutureExt, StreamExt};
 use linera_base::{
     data_types::{Amount, BlockHeight, Timestamp},
-    identifiers::{Account, AccountOwner, BlobType, ChainId, Destination},
+    identifiers::{Account, AccountOwner, BlobId, BlobType, ChainId, Destination},
 };
 use linera_views::{
     context::Context,
@@ -259,6 +259,7 @@ where
         context: OperationContext,
         local_time: Timestamp,
         operation: Operation,
+        blob_sizes: &BTreeMap<BlobId, usize>,
         txn_tracker: &mut TransactionTracker,
         resource_controller: &mut ResourceController<Option<AccountOwner>>,
     ) -> Result<(), ExecutionError> {
@@ -267,7 +268,7 @@ where
             Operation::System(op) => {
                 let new_application = self
                     .system
-                    .execute_operation(context, *op, txn_tracker, resource_controller)
+                    .execute_operation(context, *op, blob_sizes, txn_tracker, resource_controller)
                     .await?;
                 if let Some((application_id, argument)) = new_application {
                     let user_action = UserAction::Instantiate(context, argument);

linera-execution/src/resources.rs

@@ -7,7 +7,7 @@ use std::{sync::Arc, time::Duration};
 
 use custom_debug_derive::Debug;
 use linera_base::{
-    data_types::{Amount, ArithmeticError, BlobContent},
+    data_types::{Amount, ArithmeticError},
     ensure,
     identifiers::AccountOwner,
 };
@@ -292,9 +292,8 @@ where
     }
 
     /// Tracks a number of blob bytes published.
-    pub fn track_blob_published(&mut self, content: &BlobContent) -> Result<(), ExecutionError> {
-        self.policy.check_blob_size(content)?;
-        let size = content.bytes().len() as u64;
+    pub fn track_blob_published(&mut self, size: usize) -> Result<(), ExecutionError> {
+        let size = size as u64;
         {
             let tracker = self.tracker.as_mut();
             tracker.blob_bytes_published = tracker

linera-execution/src/system.rs

@@ -310,6 +310,7 @@ where
         &mut self,
         context: OperationContext,
         operation: SystemOperation,
+        blob_sizes: &BTreeMap<BlobId, usize>,
         txn_tracker: &mut TransactionTracker,
         resource_controller: &mut ResourceController<Option<AccountOwner>>,
     ) -> Result<Option<(ApplicationId, Vec<u8>)>, ExecutionError> {
@@ -407,14 +408,19 @@ where
                 }
             }
             PublishModule { module_id } => {
-                self.blob_published(&BlobId::new(
-                    module_id.contract_blob_hash,
-                    BlobType::ContractBytecode,
-                ))?;
-                self.blob_published(&BlobId::new(
-                    module_id.service_blob_hash,
-                    BlobType::ServiceBytecode,
-                ))?;
+                let blob_ids = module_id.bytecode_blob_ids();
+                {
+                    let mut controller = resource_controller.with_state(self).await?;
+                    for blob_id in blob_ids {
+                        let size = *blob_sizes
+                            .get(&blob_id)
+                            .ok_or_else(|| ExecutionError::BlobsNotFound(vec![blob_id]))?;
+                        controller.track_blob_published(size)?;
+                    }
+                }
+                for blob_id in blob_ids {
+                    self.blob_published(&blob_id)?;
+                }
             }
             CreateApplication {
                 module_id,
@@ -440,7 +446,15 @@ where
                 new_application = Some((app_id, instantiation_argument));
             }
             PublishDataBlob { blob_hash } => {
-                self.blob_published(&BlobId::new(blob_hash, BlobType::Data))?;
+                let blob_id = BlobId::new(blob_hash, BlobType::Data);
+                let size = *blob_sizes
+                    .get(&blob_id)
+                    .ok_or_else(|| ExecutionError::BlobsNotFound(vec![blob_id]))?;
+                resource_controller
+                    .with_state(self)
+                    .await?
+                    .track_blob_published(size)?;
+                self.blob_published(&blob_id)?;
             }
             ReadBlob { blob_id } => {
                 let content = self.read_blob_content(blob_id).await?;
@@ -918,35 +932,20 @@ where
         &mut self,
         module_id: &ModuleId,
     ) -> Result<(BlobId, BlobId), ExecutionError> {
-        let contract_bytecode_blob_id =
-            BlobId::new(module_id.contract_blob_hash, BlobType::ContractBytecode);
+        let blob_ids = module_id.bytecode_blob_ids();
 
         let mut missing_blobs = Vec::new();
-        if !self
-            .context()
-            .extra()
-            .contains_blob(contract_bytecode_blob_id)
-            .await?
-        {
-            missing_blobs.push(contract_bytecode_blob_id);
-        }
-
-        let service_bytecode_blob_id =
-            BlobId::new(module_id.service_blob_hash, BlobType::ServiceBytecode);
-        if !self
-            .context()
-            .extra()
-            .contains_blob(service_bytecode_blob_id)
-            .await?
-        {
-            missing_blobs.push(service_bytecode_blob_id);
+        for blob_id in blob_ids {
+            if !self.context().extra().contains_blob(blob_id).await? {
+                missing_blobs.push(blob_id);
+            }
         }
 
         ensure!(
             missing_blobs.is_empty(),
             ExecutionError::BlobsNotFound(missing_blobs)
         );
 
-        Ok((contract_bytecode_blob_id, service_bytecode_blob_id))
+        Ok((blob_ids[0], blob_ids[1]))
     }
 }

linera-execution/src/unit_tests/system_tests.rs

@@ -79,6 +79,7 @@ async fn application_message_index() -> anyhow::Result<()> {
         .execute_operation(
             context,
             operation,
+            &BTreeMap::new(),
             &mut txn_tracker,
             &mut ResourceController::default(),
         )
@@ -112,6 +113,7 @@ async fn open_chain_message_index() {
         .execute_operation(
             context,
             operation,
+            &BTreeMap::new(),
             &mut txn_tracker,
             &mut ResourceController::default(),
         )

linera-execution/tests/contract_runtime_apis.rs

@@ -97,6 +97,7 @@ async fn test_transfer_system_api(
         context,
         Timestamp::from(0),
         operation,
+        &BTreeMap::new(),
         &mut tracker,
         &mut controller,
     )
@@ -184,6 +185,7 @@ async fn test_unauthorized_transfer_system_api(
             context,
             Timestamp::from(0),
             operation,
+            &BTreeMap::new(),
             &mut TransactionTracker::new(
                 0,
                 0,
@@ -287,6 +289,7 @@ async fn test_claim_system_api(
             context,
             Timestamp::from(0),
             operation,
+            &BTreeMap::new(),
             &mut tracker,
             &mut controller,
         )
@@ -440,6 +443,7 @@ async fn test_unauthorized_claims(
             context,
             Timestamp::from(0),
             operation,
+            &BTreeMap::new(),
             &mut tracker,
             &mut controller,
         )
@@ -494,6 +498,7 @@ async fn test_read_chain_balance_system_api(chain_balance: Amount) {
         context,
         Timestamp::from(0),
         operation,
+        &BTreeMap::new(),
         &mut TransactionTracker::new(
             0,
             0,
@@ -545,6 +550,7 @@ async fn test_read_owner_balance_system_api(
         context,
         Timestamp::from(0),
         operation,
+        &BTreeMap::new(),
         &mut TransactionTracker::new(0, 0, Some(blob_oracle_responses(blobs.iter()))),
         &mut controller,
     )
@@ -586,6 +592,7 @@ async fn test_read_owner_balance_returns_zero_for_missing_accounts(missing_accou
         context,
         Timestamp::from(0),
         operation,
+        &BTreeMap::new(),
         &mut TransactionTracker::new(0, 0, Some(blob_oracle_responses(blobs.iter()))),
         &mut controller,
     )
@@ -630,6 +637,7 @@ async fn test_read_owner_balances_system_api(
         context,
         Timestamp::from(0),
         operation,
+        &BTreeMap::new(),
         &mut TransactionTracker::new(0, 0, Some(blob_oracle_responses(blobs.iter()))),
         &mut controller,
     )
@@ -674,6 +682,7 @@ async fn test_read_balance_owners_system_api(
         context,
         Timestamp::from(0),
         operation,
+        &BTreeMap::new(),
         &mut TransactionTracker::new(0, 0, Some(blob_oracle_responses(blobs.iter()))),
         &mut controller,
     )
@@ -907,6 +916,7 @@ async fn test_query_service(authorized_apps: Option<Vec<()>>) -> Result<(), Exec
         context,
         Timestamp::from(0),
         operation,
+        &BTreeMap::new(),
         &mut TransactionTracker::new(
             0,
             0,
@@ -983,6 +993,7 @@ async fn test_perform_http_request(authorized_apps: Option<Vec<()>>) -> Result<(
         context,
         Timestamp::from(0),
         operation,
+        &BTreeMap::new(),
         &mut TransactionTracker::new(
             0,
             0,

linera-execution/tests/revm.rs

@@ -3,7 +3,7 @@
 
 #![cfg(with_revm)]
 
-use std::sync::Arc;
+use std::{collections::BTreeMap, sync::Arc};
 
 use alloy_sol_types::{sol, SolCall, SolValue};
 use linera_base::{
@@ -131,6 +131,7 @@ async fn test_fuel_for_counter_revm_application() -> anyhow::Result<()> {
             operation_context,
             Timestamp::from(0),
             operation,
+            &BTreeMap::new(),
             &mut txn_tracker,
             &mut controller,
         )