Handle discount in setup costs for in memory cache contracts

app/ante.go

@@ -55,6 +55,7 @@ func NewAnteHandler(options HandlerOptions) (sdk.AnteHandler, error) {
 		wasmkeeper.NewLimitSimulationGasDecorator(options.WasmConfig.SimulationGasLimit), // after setup context to enforce limits early
 		wasmkeeper.NewCountTXDecorator(options.TXCounterStoreService),
 		wasmkeeper.NewGasRegisterDecorator(options.WasmKeeper.GetGasRegister()),
+		wasmkeeper.NewTxContractsDecorator(),
 		circuitante.NewCircuitBreakerDecorator(options.CircuitKeeper),
 		ante.NewExtensionOptionsDecorator(options.ExtensionOptionChecker),
 		ante.NewValidateBasicDecorator(),

x/wasm/keeper/ante.go

@@ -120,3 +120,16 @@
 func (g GasRegisterDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (sdk.Context, error) {
 	return next(types.WithGasRegister(ctx, g.gasRegister), tx, simulate)
 }
+
+type TxContractsDecorator struct{}
+
+// NewTxContractsDecorator constructor.
+func NewTxContractsDecorator() *TxContractsDecorator {
+	return &TxContractsDecorator{}
+}
+
+// AnteHandle initializes a new TxContracts object to the context.
+func (d TxContractsDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (sdk.Context, error) {
+	txContracts := types.NewTxContracts()
+	return next(types.WithTxContracts(ctx, txContracts), tx, simulate)
+}

x/wasm/keeper/keeper.go

@@ -252,13 +252,17 @@ func (k Keeper) instantiate(
 		return nil, nil, types.ErrEmpty.Wrap("creator")
 	}
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(sdkCtx, codeID), len(initMsg))
-	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: instantiate")
 
 	codeInfo := k.GetCodeInfo(ctx, codeID)
 	if codeInfo == nil {
 		return nil, nil, types.ErrNoSuchCodeFn(codeID).Wrapf("code id %d", codeID)
 	}
+
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, codeInfo.CodeHash, k.IsPinnedCode(sdkCtx, codeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(initMsg))
+
+	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: instantiate")
+
 	if !authPolicy.CanInstantiateContract(codeInfo.InstantiateConfig, creator) {
 		return nil, nil, errorsmod.Wrap(sdkerrors.ErrUnauthorized, "can not instantiate")
 	}
@@ -395,7 +399,9 @@ func (k Keeper) execute(ctx context.Context, contractAddress, caller sdk.AccAddr
 		return nil, err
 	}
 
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(ctx, contractInfo.CodeID), len(msg))
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, codeInfo.CodeHash, k.IsPinnedCode(ctx, contractInfo.CodeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(msg))
+
 	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: execute")
 
 	// add more funds
@@ -555,7 +561,8 @@ func (k Keeper) callMigrateEntrypoint(
 	msg []byte,
 	newCodeID uint64,
 ) (*wasmvmtypes.Response, error) {
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(sdkCtx, newCodeID), len(msg))
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, newChecksum, k.IsPinnedCode(sdkCtx, newCodeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(msg))
 	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: migrate")
 
 	env := types.NewEnv(sdkCtx, contractAddress)
@@ -596,9 +603,10 @@ func (k Keeper) Sudo(ctx context.Context, contractAddress sdk.AccAddress, msg []
 	if err != nil {
 		return nil, err
 	}
-
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(ctx, contractInfo.CodeID), len(msg))
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, codeInfo.CodeHash, k.IsPinnedCode(ctx, contractInfo.CodeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(msg))
+
 	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: sudo")
 
 	env := types.NewEnv(sdkCtx, contractAddress)
@@ -640,7 +648,6 @@ func (k Keeper) reply(ctx sdk.Context, contractAddress sdk.AccAddress, reply was
 		return nil, err
 	}
 
-	// always consider this pinned
 	replyCosts := k.gasRegister.ReplyCosts(true, reply)
 	ctx.GasMeter().ConsumeGas(replyCosts, "Loading CosmWasm module: reply")
 
@@ -834,7 +841,8 @@ func (k Keeper) QuerySmart(ctx context.Context, contractAddr sdk.AccAddress, req
 		return nil, err
 	}
 
-	setupCost := k.gasRegister.SetupContractCost(k.IsPinnedCode(sdkCtx, contractInfo.CodeID), len(req))
+	sdkCtx, discount := k.checkDiscountEligibility(sdkCtx, codeInfo.CodeHash, k.IsPinnedCode(ctx, contractInfo.CodeID))
+	setupCost := k.gasRegister.SetupContractCost(discount, len(req))
 	sdkCtx.GasMeter().ConsumeGas(setupCost, "Loading CosmWasm module: query")
 
 	// prepare querier
@@ -1158,6 +1166,23 @@ func (k Keeper) IsPinnedCode(ctx context.Context, codeID uint64) bool {
 	return ok
 }
 
+func (k Keeper) checkDiscountEligibility(ctx sdk.Context, checksum []byte, isPinned bool) (sdk.Context, bool) {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	if isPinned {
+		return sdkCtx, true
+	}
+
+	txContracts, ok := types.TxContractsFromContext(ctx)
+	if !ok {
+		txContracts = types.NewTxContracts()
+	} else if txContracts.Exists(checksum) {
+		return sdkCtx, true
+	}
+
+	txContracts.AddContract(checksum)
+	return types.WithTxContracts(sdkCtx, txContracts), false
+}
+
 // InitializePinnedCodes updates wasmvm to pin to cache all contracts marked as pinned
 func (k Keeper) InitializePinnedCodes(ctx context.Context) error {
 	store := prefix.NewStore(runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx)), types.PinnedCodeIndexPrefix)

x/wasm/keeper/recurse_test.go

@@ -56,9 +56,11 @@ func initRecurseContract(t *testing.T) (contract sdk.AccAddress, ctx sdk.Context
 
 func TestGasCostOnQuery(t *testing.T) {
 	const (
-		GasNoWork uint64 = 63_968
+		GasNoWork           uint64 = 63_968
+		GasNoWorkDiscounted uint64 = 5_968
 		// Note: about 100 SDK gas (10k CosmWasm gas) for each round of sha256
-		GasWork50 uint64 = 64_207 // this is a little shy of 50k gas - to keep an eye on the limit
+		GasWork50           uint64 = 64_207 // this is a little shy of 50k gas - to keep an eye on the limit
+		GasWork50Discounted uint64 = 6_207
 
 		GasReturnUnhashed uint64 = 55
 		GasReturnHashed   uint64 = 46
@@ -86,23 +88,23 @@ func TestGasCostOnQuery(t *testing.T) {
 			msg: Recurse{
 				Depth: 1,
 			},
-			expectedGas: 2*GasNoWork + GasReturnUnhashed,
+			expectedGas: GasNoWork + GasNoWorkDiscounted + GasReturnUnhashed,
 		},
 		"recursion 1, some work": {
 			gasLimit: 400_000,
 			msg: Recurse{
 				Depth: 1,
 				Work:  50,
 			},
-			expectedGas: 2*GasWork50 + GasReturnHashed,
+			expectedGas: GasWork50 + GasWork50Discounted + GasReturnHashed,
 		},
 		"recursion 4, some work": {
 			gasLimit: 400_000,
 			msg: Recurse{
 				Depth: 4,
 				Work:  50,
 			},
-			expectedGas: 5*GasWork50 + 4*GasReturnHashed,
+			expectedGas: GasWork50 + 4*(GasWork50Discounted+GasReturnHashed),
 		},
 	}
 
@@ -172,8 +174,7 @@ func TestGasOnExternalQuery(t *testing.T) {
 			expOutOfGas: true,
 		},
 		"recursion 4, external gas limit": {
-			// this uses 244708 gas but give less
-			gasLimit: 4 * GasWork50,
+			gasLimit: GasWork50,
 			msg: Recurse{
 				Depth: 4,
 				Work:  50,
@@ -212,6 +213,9 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 	const (
 		// Note: about 100 SDK gas (10k CosmWasm gas) for each round of sha256
 		GasWork2k uint64 = 76_264 // = SetupContractCost + x // we have 6x gas used in cpu than in the instance
+
+		GasWork2kDiscounted uint64 = 18_264
+
 		// This is overhead for calling into a sub-contract
 		GasReturnHashed uint64 = 48
 	)
@@ -240,13 +244,13 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 				Work:  2000,
 			},
 			expectQueriesFromContract: 5,
-			expectedGas:               GasWork2k + 5*(GasWork2k+GasReturnHashed),
+			expectedGas:               GasWork2k + 5*(GasWork2kDiscounted+GasReturnHashed),
 		},
 		// this is where we expect an error...
 		// it has enough gas to run 5 times and die on the 6th (5th time dispatching to sub-contract)
 		// however, if we don't charge the cpu gas before sub-dispatching, we can recurse over 20 times
 		"deep recursion, should die on 5th level": {
-			gasLimit: 400_000,
+			gasLimit: 150_000,
 			msg: Recurse{
 				Depth: 50,
 				Work:  2000,
@@ -262,8 +266,8 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 			},
 			expectQueriesFromContract: 10,
 			expectOutOfGas:            false,
-			expectError:               "query wasm contract failed",          // Error we get from the contract instance doing the failing query, not wasmd
-			expectedGas:               10*(GasWork2k+GasReturnHashed) + 3279, // lots of additional gas for long error message
+			expectError:               "query wasm contract failed",                                                 // Error we get from the contract instance doing the failing query, not wasmd
+			expectedGas:               GasWork2k + GasReturnHashed + 9*(GasWork2kDiscounted+GasReturnHashed) + 3279, // lots of additional gas for long error message
 		},
 	}
 

x/wasm/types/context.go

@@ -20,6 +20,9 @@
 	contextKeyGasRegister = iota
 
 	contextKeyCallDepth contextKey = iota
+
+	// contracts in the current tx
+	contextKeyTxContracts contextKey = iota
 )
 
 // WithTXCounter stores a transaction counter value in the context
@@ -81,3 +84,17 @@
 	val, ok := ctx.Value(contextKeyGasRegister).(GasRegister)
 	return val, ok
 }
+
+// WithTxContracts stores the tx contracts into the context returned
+func WithTxContracts(ctx sdk.Context, tc TxContracts) sdk.Context {
+	if tc.GetContracts() == nil {
+		panic("tx contracts must not be nil")
+	}
+	return ctx.WithValue(contextKeyTxContracts, tc)
+}
+
+// TxContractsFromContext reads the tx contracts from the context
+func TxContractsFromContext(ctx context.Context) (TxContracts, bool) {
+	val, ok := ctx.Value(contextKeyTxContracts).(TxContracts)
+	return val, ok
+}

x/wasm/types/types.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"encoding/hex"
 	"fmt"
 	"reflect"
 
@@ -429,3 +430,29 @@
 	}
 	return []string{}
 }
+
+type txContracts map[string]struct{}
+
+type TxContracts struct {
+	contracts txContracts
+}
+
+func NewTxContracts() TxContracts {
+	c := make(txContracts, 0)
+	return TxContracts{contracts: c}
+}
+
+func (tc TxContracts) AddContract(checksum []byte) {
+	hexHash := hex.EncodeToString(checksum)
+	tc.contracts[hexHash] = struct{}{}
+}
+
+func (tc TxContracts) Exists(checksum []byte) bool {
+	hexHash := hex.EncodeToString(checksum)
+	_, ok := tc.contracts[hexHash]
+	return ok
+}
+
+func (tc TxContracts) GetContracts() txContracts {
+	return tc.contracts
+}