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, string(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, string(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, string(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, string(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,8 +648,8 @@ 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, discount := k.checkDiscountEligibility(ctx, string(codeInfo.CodeHash), k.IsPinnedCode(ctx, contractInfo.CodeID))
+	replyCosts := k.gasRegister.ReplyCosts(discount, reply)
 	ctx.GasMeter().ConsumeGas(replyCosts, "Loading CosmWasm module: reply")
 
 	env := types.NewEnv(ctx, contractAddress)
@@ -834,7 +842,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, string(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 +1167,23 @@ func (k Keeper) IsPinnedCode(ctx context.Context, codeID uint64) bool {
 	return ok
 }
 
+func (k Keeper) checkDiscountEligibility(ctx sdk.Context, checksum string, 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/keeper/relay_test.go

@@ -427,7 +427,7 @@ func TestOnRecvPacket(t *testing.T) {
 		},
 		"submessage reply can overwrite ack data": {
 			contractAddr:   example.Contract,
-			expContractGas: types.DefaultInstanceCostDiscount + myContractGas + storageCosts,
+			expContractGas: types.DefaultInstanceCost + myContractGas + storageCosts + 1_000,
 			contractResp: &wasmvmtypes.IBCReceiveResult{
 				Ok: &wasmvmtypes.IBCReceiveResponse{
 					Acknowledgement: []byte("myAck"),

x/wasm/keeper/submsg_test.go

@@ -243,14 +243,14 @@ func TestDispatchSubMsgErrorHandling(t *testing.T) {
 		"send tokens": {
 			submsgID:         5,
 			msg:              validBankSend,
-			resultAssertions: []assertion{assertReturnedEvents(0), assertGasUsed(110_000, 112_000)},
+			resultAssertions: []assertion{assertReturnedEvents(0), assertGasUsed(110_000, 113_000)},
 		},
 		"not enough tokens": {
 			submsgID:    6,
 			msg:         invalidBankSend,
 			subMsgError: true,
 			// uses less gas than the send tokens (cost of bank transfer)
-			resultAssertions: []assertion{assertGasUsed(78_000, 81_000), assertErrorString("codespace: sdk, code: 5")},
+			resultAssertions: []assertion{assertGasUsed(78_000, 82_000), assertErrorString("codespace: sdk, code: 5")},
 		},
 		"out of gas panic with no gas limit": {
 			submsgID:        7,
@@ -263,23 +263,23 @@ func TestDispatchSubMsgErrorHandling(t *testing.T) {
 			msg:      validBankSend,
 			gasLimit: &subGasLimit,
 			// uses same gas as call without limit (note we do not charge the 40k on reply)
-			resultAssertions: []assertion{assertReturnedEvents(0), assertGasUsed(110_000, 112_000)},
+			resultAssertions: []assertion{assertReturnedEvents(0), assertGasUsed(110_000, 113_000)},
 		},
 		"not enough tokens with limit": {
 			submsgID:    16,
 			msg:         invalidBankSend,
 			subMsgError: true,
 			gasLimit:    &subGasLimit,
 			// uses same gas as call without limit (note we do not charge the 40k on reply)
-			resultAssertions: []assertion{assertGasUsed(78_000, 81_000), assertErrorString("codespace: sdk, code: 5")},
+			resultAssertions: []assertion{assertGasUsed(78_000, 82_000), assertErrorString("codespace: sdk, code: 5")},
 		},
 		"out of gas caught with gas limit": {
 			submsgID:    17,
 			msg:         infiniteLoop,
 			subMsgError: true,
 			gasLimit:    &subGasLimit,
 			// uses all the subGasLimit, plus the 52k or so for the main contract
-			resultAssertions: []assertion{assertGasUsed(subGasLimit+75_000, subGasLimit+77_000), assertErrorString("codespace: sdk, code: 11")},
+			resultAssertions: []assertion{assertGasUsed(subGasLimit+75_000, subGasLimit+78_000), assertErrorString("codespace: sdk, code: 11")},
 		},
 		"instantiate contract gets address in data and events": {
 			submsgID:         21,

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, c TxContracts) sdk.Context {
+	if c == nil {
+		panic("tx contracts must not be nil")
+	}
+	return ctx.WithValue(contextKeyTxContracts, c)
+}
+
+// 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

@@ -429,3 +429,18 @@
 	}
 	return []string{}
 }
+
+type TxContracts map[string]struct{}
+
+func NewTxContracts() TxContracts {
+	return make(TxContracts, 0)
+}
+
+func (a TxContracts) AddContract(checksum string) {
+	a[checksum] = struct{}{}
+}
+
+func (a TxContracts) Exists(checksum string) bool {
+	_, ok := a[checksum]
+	return ok
+}