contracts: Remove weight pre charging (#8976)

* Remove pre-charging for code size

* Remove pre charging when reading values of fixed size

* Add new versions of API functions that leave out parameters

* Update CHANGELOG.md

* Apply suggestions from code review

Co-authored-by: Alexander Popiak <alexander.popiak@parity.io>

* Add v1 for seal_set_rent_allowance

* Remove unneeded trait bound

Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

Co-authored-by: Alexander Popiak <alexander.popiak@parity.io>
Co-authored-by: Guillaume Thiolliere <gui.thiolliere@gmail.com>

substrate/frame/contracts/src/wasm/code_cache.rs

@@ -81,14 +81,16 @@ where
 }
 
 /// Increment the refcount of a code in-storage by one.
-pub fn increment_refcount<T: Config>(code_hash: CodeHash<T>) -> Result<u32, DispatchError>
+pub fn increment_refcount<T: Config>(code_hash: CodeHash<T>, gas_meter: &mut GasMeter<T>)
+	-> Result<(), DispatchError>
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
 {
+	gas_meter.charge(CodeToken::UpdateRefcount(estimate_code_size::<T>(&code_hash)?))?;
 	<CodeStorage<T>>::mutate(code_hash, |existing| {
 		if let Some(module) = existing {
 			increment_64(&mut module.refcount);
-			Ok(module.original_code_len)
+			Ok(())
 		} else {
 			Err(Error::<T>::CodeNotFound.into())
 		}
@@ -96,23 +98,24 @@ where
 }
 
 /// Decrement the refcount of a code in-storage by one and remove the code when it drops to zero.
-pub fn decrement_refcount<T: Config>(code_hash: CodeHash<T>) -> u32
+pub fn decrement_refcount<T: Config>(code_hash: CodeHash<T>, gas_meter: &mut GasMeter<T>)
+	-> Result<(), DispatchError>
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
 {
+	if let Ok(len) = estimate_code_size::<T>(&code_hash) {
+		gas_meter.charge(CodeToken::UpdateRefcount(len))?;
+	}
 	<CodeStorage<T>>::mutate_exists(code_hash, |existing| {
 		if let Some(module) = existing {
-			let code_len = module.original_code_len;
 			module.refcount = module.refcount.saturating_sub(1);
 			if module.refcount == 0 {
 				*existing = None;
 				finish_removal::<T>(code_hash);
 			}
-			code_len
-		} else {
-			0
 		}
-	})
+	});
+	Ok(())
 }
 
 /// Load code with the given code hash.
@@ -120,13 +123,24 @@ where
 /// If the module was instrumented with a lower version of schedule than
 /// the current one given as an argument, then this function will perform
 /// re-instrumentation and update the cache in the storage.
+///
+/// # Note
+///
+/// If `reinstrument` is set it is assumed that the load is performed in the context of
+/// a contract call: This means we charge the size based cased for loading the contract.
 pub fn load<T: Config>(
 	code_hash: CodeHash<T>,
-	reinstrument: Option<(&Schedule<T>, &mut GasMeter<T>)>,
+	mut reinstrument: Option<(&Schedule<T>, &mut GasMeter<T>)>,
 ) -> Result<PrefabWasmModule<T>, DispatchError>
 where
 	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
 {
+	// The reinstrument case coincides with the cases where we need to charge extra
+	// based upon the code size: On-chain execution.
+	if let Some((_, gas_meter)) = &mut reinstrument {
+		gas_meter.charge(CodeToken::Load(estimate_code_size::<T>(&code_hash)?))?;
+	}
+
 	let mut prefab_module = <CodeStorage<T>>::get(code_hash)
 		.ok_or_else(|| Error::<T>::CodeNotFound)?;
 	prefab_module.code_hash = code_hash;
@@ -135,7 +149,7 @@ where
 		if prefab_module.instruction_weights_version < schedule.instruction_weights.version {
 			// The instruction weights have changed.
 			// We need to re-instrument the code with the new instruction weights.
-			gas_meter.charge(InstrumentToken(prefab_module.original_code_len))?;
+			gas_meter.charge(CodeToken::Instrument(prefab_module.original_code_len))?;
 			private::reinstrument(&mut prefab_module, schedule)?;
 		}
 	}
@@ -185,14 +199,50 @@ fn increment_64(refcount: &mut u64) {
 	");
 }
 
-/// Token to be supplied to the gas meter which charges the weight needed for reinstrumenting
-/// a contract of the specified size in bytes.
+/// Get the size of the instrumented code stored at `code_hash` without loading it.
+///
+/// The returned value is slightly too large because it also contains the fields apart from
+/// `code` which are located inside [`PrefabWasmModule`]. However, those are negligible when
+/// compared to the code size. Additionally, charging too much weight is completely safe.
+fn estimate_code_size<T: Config>(code_hash: &CodeHash<T>) -> Result<u32, DispatchError>
+where
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+{
+	let key = <CodeStorage<T>>::hashed_key_for(code_hash);
+	let mut data = [0u8; 0];
+	let len = sp_io::storage::read(&key, &mut data, 0).ok_or_else(|| Error::<T>::CodeNotFound)?;
+	Ok(len)
+}
+
+/// Costs for operations that are related to code handling.
 #[cfg_attr(test, derive(Debug, PartialEq, Eq))]
 #[derive(Clone, Copy)]
-struct InstrumentToken(u32);
+enum CodeToken {
+	/// Weight for instrumenting a contract contract of the supplied size in bytes.
+	Instrument(u32),
+	/// Weight for loading a contract per kilobyte.
+	Load(u32),
+	/// Weight for changing the refcount of a contract per kilobyte.
+	UpdateRefcount(u32),
+}
 
-impl<T: Config> Token<T> for InstrumentToken {
+impl<T> Token<T> for CodeToken
+where
+	T: Config,
+	T::AccountId: UncheckedFrom<T::Hash> + AsRef<[u8]>
+{
 	fn weight(&self) -> Weight {
-		T::WeightInfo::instrument(self.0 / 1024)
+		use self::CodeToken::*;
+		// In case of `Load` and `UpdateRefcount` we already covered the general costs of
+		// accessing the storage but still need to account for the actual size of the
+		// contract code. This is why we substract `T::*::(0)`. We need to do this at this
+		// point because when charging the general weight we do not know the size of
+		// the contract.
+		match *self {
+			Instrument(len) => T::WeightInfo::instrument(len / 1024),
+			Load(len) => T::WeightInfo::code_load(len / 1024).saturating_sub(T::WeightInfo::code_load(0)),
+			UpdateRefcount(len) =>
+				T::WeightInfo::code_refcount(len / 1024).saturating_sub(T::WeightInfo::code_refcount(0)),
+		}
 	}
 }