Add rebuilding of wasmstore as a part of node initialization

arbos/programs/programs.go

@@ -247,6 +247,10 @@ func (p Programs) CallProgram(
 
 func getWasm(statedb vm.StateDB, program common.Address) ([]byte, error) {
 	prefixedWasm := statedb.GetCode(program)
+	return getWasmFromContractCode(prefixedWasm)
+}
+
+func getWasmFromContractCode(prefixedWasm []byte) ([]byte, error) {
 	if prefixedWasm == nil {
 		return nil, ProgramNotWasmError()
 	}

arbos/programs/wasmstorehelper.go

@@ -0,0 +1,80 @@
+// Copyright 2022-2024, Offchain Labs, Inc.
+// For license information, see https://github.com/nitro/blob/master/LICENSE
+
+//go:build !wasm
+// +build !wasm
+
+package programs
+
+import (
+	"fmt"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/log"
+)
+
+// SaveActiveProgramToWasmStore is used to save active stylus programs to wasm store during rebuilding
+func (p Programs) SaveActiveProgramToWasmStore(statedb *state.StateDB, codeHash common.Hash, code []byte, time uint64, debugMode bool, rebuildingStartBlockTime uint64) error {
+	params, err := p.Params()
+	if err != nil {
+		return err
+	}
+
+	program, err := p.getActiveProgram(codeHash, time, params)
+	if err != nil {
+		// The program is not active so return early
+		log.Info("program is not active, getActiveProgram returned error, hence do not include in rebuilding", "err", err)
+		return nil
+	}
+
+	// It might happen that node crashed some time after rebuilding commenced and before it completed, hence when rebuilding
+	// resumes after node is restarted the latest diskdb derived from statedb might now have codehashes that were activated
+	// during the last rebuilding session. In such cases we don't need to fetch moduleshashes but instead return early
+	// since they would already be added to the wasm store
+	currentHoursSince := hoursSinceArbitrum(rebuildingStartBlockTime)
+	if currentHoursSince < program.activatedAt {
+		return nil
+	}
+
+	moduleHash, err := p.moduleHashes.Get(codeHash)
+	if err != nil {
+		return err
+	}
+
+	// If already in wasm store then return early
+	localAsm, err := statedb.TryGetActivatedAsm(moduleHash)
+	if err == nil && len(localAsm) > 0 {
+		return nil
+	}
+
+	wasm, err := getWasmFromContractCode(code)
+	if err != nil {
+		log.Error("Failed to reactivate program while rebuilding wasm store: getWasmFromContractCode", "expected moduleHash", moduleHash, "err", err)
+		return fmt.Errorf("failed to reactivate program while rebuilding wasm store: %w", err)
+	}
+
+	unlimitedGas := uint64(0xffffffffffff)
+	// We know program is activated, so it must be in correct version and not use too much memory
+	// Empty program address is supplied because we dont have access to this during rebuilding of wasm store
+	info, asm, module, err := activateProgramInternal(statedb, common.Address{}, codeHash, wasm, params.PageLimit, program.version, debugMode, &unlimitedGas)
+	if err != nil {
+		log.Error("failed to reactivate program while rebuilding wasm store", "expected moduleHash", moduleHash, "err", err)
+		return fmt.Errorf("failed to reactivate program while rebuilding wasm store: %w", err)
+	}
+
+	if info.moduleHash != moduleHash {
+		log.Error("failed to reactivate program while rebuilding wasm store", "expected moduleHash", moduleHash, "got", info.moduleHash)
+		return fmt.Errorf("failed to reactivate program while rebuilding wasm store, expected ModuleHash: %v", moduleHash)
+	}
+
+	batch := statedb.Database().WasmStore().NewBatch()
+	rawdb.WriteActivation(batch, moduleHash, asm, module)
+	if err := batch.Write(); err != nil {
+		log.Error("failed writing re-activation to state while rebuilding wasm store", "err", err)
+		return err
+	}
+
+	return nil
+}

cmd/nitro/init.go

@@ -319,7 +319,38 @@ func openInitializeChainDb(ctx context.Context, stack *node.Node, config *NodeCo
 						return chainDb, l2BlockChain, fmt.Errorf("failed to recreate missing states: %w", err)
 					}
 				}
-
+				latestBlock := l2BlockChain.CurrentBlock()
+				if latestBlock.Number.Uint64() <= chainConfig.ArbitrumChainParams.GenesisBlockNum ||
+					types.DeserializeHeaderExtraInformation(latestBlock).ArbOSFormatVersion < params.ArbosVersion_Stylus {
+					// If there is only genesis block or no blocks in the blockchain, set Rebuilding of wasm store to Done
+					// If Stylus upgrade hasn't yet happened, skipping rebuilding of wasm store
+					log.Info("setting rebuilding of wasm store to done")
+					if err = gethexec.WriteToKeyValueStore(wasmDb, gethexec.RebuildingPositionKey, gethexec.RebuildingDone); err != nil {
+						return nil, nil, fmt.Errorf("unable to set rebuilding status of wasm store to done: %w", err)
+					}
+				} else {
+					position, err := gethexec.ReadFromKeyValueStore[common.Hash](wasmDb, gethexec.RebuildingPositionKey)
+					if err != nil {
+						log.Info("unable to get codehash position in rebuilding of wasm store, its possible it isnt initialized yet, so initializing it and starting rebuilding", "err", err)
+						if err := gethexec.WriteToKeyValueStore(wasmDb, gethexec.RebuildingPositionKey, common.Hash{}); err != nil {
+							return nil, nil, fmt.Errorf("unable to initialize codehash position in rebuilding of wasm store to beginning: %w", err)
+						}
+					}
+					if position != gethexec.RebuildingDone {
+						startBlockHash, err := gethexec.ReadFromKeyValueStore[common.Hash](wasmDb, gethexec.RebuildingStartBlockHashKey)
+						if err != nil {
+							log.Info("unable to get start block hash in rebuilding of wasm store, its possible it isnt initialized yet, so initializing it to latest block hash", "err", err)
+							if err := gethexec.WriteToKeyValueStore(wasmDb, gethexec.RebuildingStartBlockHashKey, latestBlock.Hash()); err != nil {
+								return nil, nil, fmt.Errorf("unable to initialize start block hash in rebuilding of wasm store to latest block hash: %w", err)
+							}
+							startBlockHash = latestBlock.Hash()
+						}
+						log.Info("starting or continuing rebuilding of wasm store", "codeHash", position, "startBlockHash", startBlockHash)
+						if err := gethexec.RebuildWasmStore(ctx, wasmDb, l2BlockChain, position, startBlockHash); err != nil {
+							return nil, nil, fmt.Errorf("error rebuilding of wasm store: %w", err)
+						}
+					}
+				}
 				return chainDb, l2BlockChain, nil
 			}
 			readOnlyDb.Close()
@@ -359,6 +390,13 @@ func openInitializeChainDb(ctx context.Context, stack *node.Node, config *NodeCo
 	}
 	chainDb := rawdb.WrapDatabaseWithWasm(chainData, wasmDb, 1)
 
+	// Rebuilding wasm store is not required when just starting out
+	err = gethexec.WriteToKeyValueStore(wasmDb, gethexec.RebuildingPositionKey, gethexec.RebuildingDone)
+	log.Info("setting codehash position in rebuilding of wasm store to done")
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to set codehash position in rebuilding of wasm store to done: %w", err)
+	}
+
 	if config.Init.ImportFile != "" {
 		initDataReader, err = statetransfer.NewJsonInitDataReader(config.Init.ImportFile)
 		if err != nil {

execution/gethexec/wasmstorerebuilder.go

@@ -0,0 +1,114 @@
+// Copyright 2021-2024, Offchain Labs, Inc.
+// For license information, see https://github.com/nitro/blob/master/LICENSE
+
+package gethexec
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core"
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/offchainlabs/nitro/arbos/arbosState"
+)
+
+var RebuildingPositionKey []byte = []byte("_rebuildingPosition")             // contains the codehash upto which rebuilding of wasm store was last completed. Initialized to common.Hash{} at the start
+var RebuildingStartBlockHashKey []byte = []byte("_rebuildingStartBlockHash") // contains the block hash of starting block when rebuilding of wasm store first began
+var RebuildingDone common.Hash = common.BytesToHash([]byte("_done"))         // indicates that the rebuilding is done, if RebuildingPositionKey holds this value it implies rebuilding was completed
+
+func ReadFromKeyValueStore[T any](store ethdb.KeyValueStore, key []byte) (T, error) {
+	var empty T
+	posBytes, err := store.Get(key)
+	if err != nil {
+		return empty, err
+	}
+	var val T
+	err = rlp.DecodeBytes(posBytes, &val)
+	if err != nil {
+		return empty, fmt.Errorf("error decoding value stored for key in the KeyValueStore: %w", err)
+	}
+	return val, nil
+}
+
+func WriteToKeyValueStore[T any](store ethdb.KeyValueStore, key []byte, val T) error {
+	valBytes, err := rlp.EncodeToBytes(val)
+	if err != nil {
+		return err
+	}
+	err = store.Put(key, valBytes)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// RebuildWasmStore function runs a loop looking at every codehash in diskDb, checking if its an activated stylus contract and
+// saving it to wasm store if it doesnt already exists. When errored it logs them and silently returns
+//
+// It stores the status of rebuilding to wasm store by updating the codehash (of the latest sucessfully checked contract) in
+// RebuildingPositionKey every 50 checks.
+//
+// It also stores a special value that is only set once when rebuilding commenced in RebuildingStartBlockHashKey as the block
+// time of the latest block when rebuilding was first called, this is used to avoid recomputing of assembly and module of
+// contracts that were created after rebuilding commenced since they would anyway already be added during sync.
+func RebuildWasmStore(ctx context.Context, wasmStore ethdb.KeyValueStore, l2Blockchain *core.BlockChain, position, rebuildingStartBlockHash common.Hash) error {
+	var err error
+	var stateDb *state.StateDB
+	latestHeader := l2Blockchain.CurrentBlock()
+	// Attempt to get state at the start block when rebuilding commenced, if not available (in case of non-archival nodes) use latest state
+	rebuildingStartHeader := l2Blockchain.GetHeaderByHash(rebuildingStartBlockHash)
+	stateDb, err = l2Blockchain.StateAt(rebuildingStartHeader.Root)
+	if err != nil {
+		log.Info("error getting state at start block of rebuilding wasm store, attempting rebuilding with latest state", "err", err)
+		stateDb, err = l2Blockchain.StateAt(latestHeader.Root)
+		if err != nil {
+			return fmt.Errorf("error getting state at latest block, aborting rebuilding: %w", err)
+		}
+	}
+	diskDb := stateDb.Database().DiskDB()
+	arbState, err := arbosState.OpenSystemArbosState(stateDb, nil, true)
+	if err != nil {
+		return fmt.Errorf("error getting arbos state, aborting rebuilding: %w", err)
+	}
+	programs := arbState.Programs()
+	iter := diskDb.NewIterator(rawdb.CodePrefix, position[:])
+	defer iter.Release()
+	lastStatusUpdate := time.Now()
+	for iter.Next() {
+		codeHashBytes := bytes.TrimPrefix(iter.Key(), rawdb.CodePrefix)
+		codeHash := common.BytesToHash(codeHashBytes)
+		code := iter.Value()
+		if state.IsStylusProgram(code) {
+			if err := programs.SaveActiveProgramToWasmStore(stateDb, codeHash, code, latestHeader.Time, l2Blockchain.Config().DebugMode(), rebuildingStartHeader.Time); err != nil {
+				return fmt.Errorf("error while rebuilding of wasm store, aborting rebuilding: %w", err)
+			}
+		}
+		// After every one second of work, update the rebuilding position
+		// This also notifies user that we are working on rebuilding
+		if time.Since(lastStatusUpdate) >= time.Second || ctx.Err() != nil {
+			log.Info("Storing rebuilding status to disk", "codeHash", codeHash)
+			if err := WriteToKeyValueStore(wasmStore, RebuildingPositionKey, codeHash); err != nil {
+				return fmt.Errorf("error updating codehash position in rebuilding of wasm store: %w", err)
+			}
+			// If outer context is cancelled we should terminate rebuilding
+			// We attempted to write the latest checked codeHash to wasm store
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			lastStatusUpdate = time.Now()
+		}
+	}
+	// Set rebuilding position to done indicating completion
+	if err := WriteToKeyValueStore(wasmStore, RebuildingPositionKey, RebuildingDone); err != nil {
+		return fmt.Errorf("error updating codehash position in rebuilding of wasm store to done: %w", err)
+	}
+	log.Info("Rebuilding of wasm store was successful")
+	return nil
+}