Add migration script to prune wasm codes

x/wasm/keeper/keeper.go

@@ -205,6 +205,11 @@ func (k Keeper) storeCodeInfo(ctx sdk.Context, codeID uint64, codeInfo types.Cod
 	store.Set(types.GetCodeKey(codeID), k.cdc.MustMarshal(&codeInfo))
 }
 
+func (k Keeper) deleteCodeInfo(ctx sdk.Context, codeID uint64) {
+	store := ctx.KVStore(k.storeKey)
+	store.Delete(types.GetCodeKey(codeID))
+}
+
 func (k Keeper) importCode(ctx sdk.Context, codeID uint64, codeInfo types.CodeInfo, wasmCode []byte) error {
 	if ioutils.IsGzip(wasmCode) {
 		var err error
@@ -1226,3 +1231,80 @@ func (h DefaultWasmVMContractResponseHandler) Handle(ctx sdk.Context, contractAd
 	}
 	return result, nil
 }
+
+// GetByteCodeByChecksum queries the wasm code by checksum
+func (k Keeper) GetByteCodeByChecksum(ctx sdk.Context, checksum wasmvm.Checksum) ([]byte, error) {
+	return k.wasmVM.GetCode(checksum)
+}
+
+// iteratePinnedCodes iterates over all pinned code ids in ascending order
+func (k Keeper) iteratePinnedCodes(ctx sdk.Context, cb func(codeID uint64) bool) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.PinnedCodeIndexPrefix)
+	iter := store.Iterator(nil, nil)
+	defer iter.Close()
+
+	for ; iter.Valid(); iter.Next() {
+		codeID := sdk.BigEndianToUint64(iter.Key())
+		if cb(codeID) {
+			return
+		}
+	}
+}
+
+// PruneWasmCodes deletes code info for unpinned codes <= than maxCodeID
+func (k Keeper) PruneWasmCodes(ctx sdk.Context, maxCodeID uint64) error {
+	usedCodeIDs := make(map[uint64]struct{})
+	usedChecksums := make(map[string]struct{})
+
+	// collect code ids used by contracts
+	k.IterateContractInfo(ctx, func(_ sdk.AccAddress, info types.ContractInfo) bool {
+		usedCodeIDs[info.CodeID] = struct{}{}
+		return false
+	})
+
+	// collect pinned code ids
+	k.iteratePinnedCodes(ctx, func(codeID uint64) bool {
+		usedCodeIDs[codeID] = struct{}{}
+		return false
+	})
+
+	// check if instances are used only by unpinned code ids <= maxCodeID
+	k.IterateCodeInfos(ctx, func(codeID uint64, info types.CodeInfo) bool {
+		if codeID > maxCodeID { // keep all
+			usedChecksums[string(info.CodeHash)] = struct{}{}
+			return false
+		}
+		if _, ok := usedCodeIDs[codeID]; ok {
+			usedChecksums[string(info.CodeHash)] = struct{}{}
+		}
+		return false
+	})
+
+	var (
+		deletedCodeInfoCounter int
+		deletedWasmFileCounter int
+	)
+	// delete all unpinned code ids <= maxCodeID and all the
+	// instances used only by unpinned code ids <= maxCodeID
+	k.IterateCodeInfos(ctx, func(codeID uint64, info types.CodeInfo) bool {
+		if codeID > maxCodeID {
+			return true
+		}
+		if _, ok := usedCodeIDs[codeID]; !ok {
+			k.deleteCodeInfo(ctx, codeID)
+			deletedCodeInfoCounter++
+			if _, ok := usedChecksums[string(info.CodeHash)]; !ok {
+				if err := k.wasmVM.RemoveCode(info.CodeHash); err != nil {
+					k.Logger(ctx).Error("failed to delete wasm file on disk", "checksum", info.CodeHash)
+				} else {
+					deletedWasmFileCounter++
+				}
+			}
+		}
+		return false
+	})
+	k.Logger(ctx).Info("executed prune wasm code",
+		"total wasm files", deletedWasmFileCounter, "total code infos", deletedCodeInfoCounter)
+
+	return nil
+}

x/wasm/keeper/migrations.go

@@ -7,6 +7,7 @@ import (
 	v1 "github.com/CosmWasm/wasmd/x/wasm/migrations/v1"
 	v2 "github.com/CosmWasm/wasmd/x/wasm/migrations/v2"
 	v3 "github.com/CosmWasm/wasmd/x/wasm/migrations/v3"
+	v4 "github.com/CosmWasm/wasmd/x/wasm/migrations/v4"
 )
 
 // Migrator is a struct for handling in-place store migrations.
@@ -36,3 +37,9 @@ func (m Migrator) Migrate2to3(ctx sdk.Context) error {
 func (m Migrator) Migrate3to4(ctx sdk.Context) error {
 	return v3.NewMigrator(m.keeper, m.keeper.storeCodeInfo).Migrate3to4(ctx, m.keeper.storeKey, m.keeper.cdc)
 }
+
+// Migrate4to5 migrates the x/wasm module state from the consensus
+// version 4 to version 5.
+func (m Migrator) Migrate4to5(ctx sdk.Context) error {
+	return v4.NewMigrator(m.keeper).Migrate4to5(ctx)
+}

x/wasm/keeper/wasmtesting/mock_engine.go

@@ -36,6 +36,7 @@ type MockWasmer struct {
 	IBCPacketTimeoutFn   func(codeID wasmvm.Checksum, env wasmvmtypes.Env, msg wasmvmtypes.IBCPacketTimeoutMsg, store wasmvm.KVStore, goapi wasmvm.GoAPI, querier wasmvm.Querier, gasMeter wasmvm.GasMeter, gasLimit uint64, deserCost wasmvmtypes.UFraction) (*wasmvmtypes.IBCBasicResponse, uint64, error)
 	PinFn                func(checksum wasmvm.Checksum) error
 	UnpinFn              func(checksum wasmvm.Checksum) error
+	RemoveCodeFn         func(checksum wasmvm.Checksum) error
 	GetMetricsFn         func() (*wasmvmtypes.Metrics, error)
 }
 
@@ -177,6 +178,13 @@ func (m *MockWasmer) Unpin(checksum wasmvm.Checksum) error {
 	return m.UnpinFn(checksum)
 }
 
+func (m *MockWasmer) RemoveCode(checksum wasmvm.Checksum) error {
+	if m.RemoveCodeFn == nil {
+		panic("not supposed to be called!")
+	}
+	return m.RemoveCodeFn(checksum)
+}
+
 func (m *MockWasmer) GetMetrics() (*wasmvmtypes.Metrics, error) {
 	if m.GetMetricsFn == nil {
 		panic("not expected to be called")

x/wasm/migrations/v4/store.go

@@ -0,0 +1,27 @@
+package v4
+
+import (
+	"math"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+// Keeper abstract keeper
+type wasmKeeper interface {
+	PruneWasmCodes(ctx sdk.Context, maxCodeID uint64) error
+}
+
+// Migrator is a struct for handling in-place store migrations.
+type Migrator struct {
+	keeper wasmKeeper
+}
+
+// NewMigrator returns a new Migrator.
+func NewMigrator(k wasmKeeper) Migrator {
+	return Migrator{keeper: k}
+}
+
+// Migrate4to5 migrates from version 4 to 5.
+func (m Migrator) Migrate4to5(ctx sdk.Context) error {
+	return m.keeper.PruneWasmCodes(ctx, math.MaxUint64)
+}

x/wasm/migrations/v4/store_test.go

@@ -0,0 +1,9 @@
+package v4_test
+
+import (
+	"testing"
+)
+
+func TestMigrate4To5(t *testing.T) {
+	t.Skip("TODO")
+}

x/wasm/types/wasmer_engine.go

@@ -254,6 +254,9 @@ type WasmerEngine interface {
 	// Unpin is idempotent.
 	Unpin(checksum wasmvm.Checksum) error
 
+	// RemoveCode removes the wasm code referenced by checksum.
+	RemoveCode(checksum wasmvm.Checksum) error
+
 	// GetMetrics some internal metrics for monitoring purposes.
 	GetMetrics() (*wasmvmtypes.Metrics, error)
 }