Merge pull request #6204 from oasisprotocol/martin/feature/parallelize-checkpoint-creation-v2

go/storage/mkvs: Parallelize checkpoint creation

Author: martintomazic

.changelog/6204.feature.md

@@ -0,0 +1 @@
+go/storage/mkvs: Parallelize checkpoint creation

go/consensus/cometbft/abci/mux.go

@@ -63,6 +63,7 @@ type ApplicationConfig struct { // nolint: maligned
 
 	DisableCheckpointer       bool
 	CheckpointerCheckInterval time.Duration
+	ChunkerThreads            uint16
 
 	// Identity is the local node identity.
 	Identity *identity.Identity

go/consensus/cometbft/abci/snapshots.go

@@ -212,7 +212,11 @@ func (mux *abciMux) ApplySnapshotChunk(req types.RequestApplySnapshotChunk) type
 			"err", err,
 		)
 
-		return types.ResponseApplySnapshotChunk{Result: types.ResponseApplySnapshotChunk_REJECT_SNAPSHOT}
+		// Given that snapshot was invalid, the sender must have been malicious.
+		return types.ResponseApplySnapshotChunk{
+			Result:        types.ResponseApplySnapshotChunk_REJECT_SNAPSHOT,
+			RejectSenders: []string{req.Sender},
+		}
 	default:
 		// Unspecified error during restoration.
 		mux.logger.Error("error during chunk restoration, aborting state sync",

go/consensus/cometbft/abci/state.go

@@ -775,6 +775,7 @@ func newApplicationState(ctx context.Context, upgrader upgrade.Backend, cfg *App
 					NumKept:        params.StateCheckpointNumKept,
 					ChunkSize:      params.StateCheckpointChunkSize,
 					InitialVersion: cfg.InitialHeight,
+					ChunkerThreads: cfg.ChunkerThreads,
 				}, nil
 			},
 		}

go/consensus/cometbft/config/config.go

@@ -115,6 +115,8 @@ type CheckpointerConfig struct {
 	Disabled bool `yaml:"disabled"`
 	// ABCI state checkpointer check interval.
 	CheckInterval time.Duration `yaml:"check_interval"`
+	// ParallelChunker specifies if the new parallel chunking algorithm is used.
+	ParallelChunker bool `yaml:"parallel_chunker"`
 }
 
 // StateSyncConfig is the consensus state sync configuration structure.
@@ -247,8 +249,9 @@ func DefaultConfig() Config {
 			NumLightBlocksKept: 10000,
 		},
 		Checkpointer: CheckpointerConfig{
-			Disabled:      false,
-			CheckInterval: 1 * time.Minute,
+			Disabled:        false,
+			CheckInterval:   1 * time.Minute,
+			ParallelChunker: false,
 		},
 		StateSync: StateSyncConfig{
 			Enabled: false,

go/consensus/cometbft/full/full.go

@@ -70,6 +70,12 @@ const (
 	tmSubscriberID = "oasis-core"
 
 	exportsSubDir = "exports"
+
+	// chunkerThreads is target number of subtrees during parallel checkpoint creation.
+	// It is intentionally non-configurable since we want operators to produce
+	// same checkpoint hashes. The current value was chosen based on the benchmarks
+	// done on the modern developer machine.
+	chunkerThreads = 12
 )
 
 var (
@@ -544,6 +550,11 @@ func (t *fullService) lazyInit() error { // nolint: gocyclo
 	pruneCfg.NumKept = config.GlobalConfig.Consensus.Prune.NumKept
 	pruneCfg.PruneInterval = max(config.GlobalConfig.Consensus.Prune.Interval, time.Second)
 
+	var threads uint16
+	if config.GlobalConfig.Storage.Checkpointer.ParallelChunker {
+		threads = chunkerThreads
+	}
+
 	appConfig := &abci.ApplicationConfig{
 		DataDir:                   filepath.Join(t.dataDir, tmcommon.StateDir),
 		StorageBackend:            config.GlobalConfig.Storage.Backend,
@@ -554,6 +565,7 @@ func (t *fullService) lazyInit() error { // nolint: gocyclo
 		Identity:                  t.identity,
 		DisableCheckpointer:       config.GlobalConfig.Consensus.Checkpointer.Disabled,
 		CheckpointerCheckInterval: config.GlobalConfig.Consensus.Checkpointer.CheckInterval,
+		ChunkerThreads:            threads,
 		InitialHeight:             uint64(t.genesisHeight),
 		ChainContext:              t.chainContext,
 	}

go/go.mod

@@ -58,6 +58,7 @@ require (
 	golang.org/x/crypto v0.39.0
 	golang.org/x/exp v0.0.0-20250128182459-e0ece0dbea4c
 	golang.org/x/net v0.41.0
+	golang.org/x/sync v0.15.0
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1
 	google.golang.org/grpc v1.68.0
 	google.golang.org/grpc/security/advancedtls v0.0.0-20221004221323-12db695f1648
@@ -210,7 +211,6 @@ require (
 	go.uber.org/mock v0.5.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/mod v0.25.0 // indirect
-	golang.org/x/sync v0.15.0 // indirect
 	golang.org/x/sys v0.33.0 // indirect
 	golang.org/x/text v0.26.0 // indirect
 	golang.org/x/tools v0.33.0 // indirect

go/oasis-test-runner/oasis/compute.go

@@ -37,10 +37,11 @@ type Compute struct { // nolint: maligned
 
 	sentryIndices []int
 
-	storageBackend          string
-	disablePublicRPC        bool
-	checkpointSyncDisabled  bool
-	checkpointCheckInterval time.Duration
+	storageBackend            string
+	disablePublicRPC          bool
+	checkpointSyncDisabled    bool
+	checkpointCheckInterval   time.Duration
+	checkpointParallelChunker bool
 
 	sentryPubKey  signature.PublicKey
 	consensusPort uint16
@@ -62,10 +63,11 @@ type ComputeCfg struct {
 
 	SentryIndices []int
 
-	StorageBackend          string
-	DisablePublicRPC        bool
-	CheckpointSyncDisabled  bool
-	CheckpointCheckInterval time.Duration
+	StorageBackend            string
+	DisablePublicRPC          bool
+	CheckpointSyncDisabled    bool
+	CheckpointCheckInterval   time.Duration
+	CheckpointParallelChunker bool
 }
 
 // UpdateRuntimes updates the worker node runtimes.
@@ -172,6 +174,7 @@ func (worker *Compute) ModifyConfig() error {
 	worker.Config.Storage.CheckpointSyncDisabled = worker.checkpointSyncDisabled
 	worker.Config.Storage.Checkpointer.Enabled = true
 	worker.Config.Storage.Checkpointer.CheckInterval = worker.checkpointCheckInterval
+	worker.Config.Storage.Checkpointer.ParallelChunker = worker.checkpointParallelChunker
 
 	// Sentry configuration.
 	sentries, err := resolveSentries(worker.net, worker.sentryIndices)
@@ -230,18 +233,19 @@ func (net *Network) NewCompute(cfg *ComputeCfg) (*Compute, error) {
 	}
 
 	worker := &Compute{
-		Node:                    host,
-		storageBackend:          cfg.StorageBackend,
-		sentryIndices:           cfg.SentryIndices,
-		disablePublicRPC:        cfg.DisablePublicRPC,
-		checkpointSyncDisabled:  cfg.CheckpointSyncDisabled,
-		checkpointCheckInterval: cfg.CheckpointCheckInterval,
-		sentryPubKey:            sentryPubKey,
-		runtimeProvisioner:      cfg.RuntimeProvisioner,
-		consensusPort:           host.getProvisionedPort(nodePortConsensus),
-		p2pPort:                 host.getProvisionedPort(nodePortP2P),
-		runtimes:                cfg.Runtimes,
-		runtimeConfig:           cfg.RuntimeConfig,
+		Node:                      host,
+		storageBackend:            cfg.StorageBackend,
+		sentryIndices:             cfg.SentryIndices,
+		disablePublicRPC:          cfg.DisablePublicRPC,
+		checkpointSyncDisabled:    cfg.CheckpointSyncDisabled,
+		checkpointCheckInterval:   cfg.CheckpointCheckInterval,
+		checkpointParallelChunker: cfg.CheckpointParallelChunker,
+		sentryPubKey:              sentryPubKey,
+		runtimeProvisioner:        cfg.RuntimeProvisioner,
+		consensusPort:             host.getProvisionedPort(nodePortConsensus),
+		p2pPort:                   host.getProvisionedPort(nodePortP2P),
+		runtimes:                  cfg.Runtimes,
+		runtimeConfig:             cfg.RuntimeConfig,
 	}
 
 	// Remove any exploded bundles on cleanup.

go/oasis-test-runner/oasis/fixture.go

@@ -401,8 +401,9 @@ type ComputeWorkerFixture struct {
 
 	Sentries []int `json:"sentries,omitempty"`
 
-	CheckpointCheckInterval time.Duration `json:"checkpoint_check_interval,omitempty"`
-	CheckpointSyncEnabled   bool          `json:"checkpoint_sync_enabled,omitempty"`
+	CheckpointCheckInterval   time.Duration `json:"checkpoint_check_interval,omitempty"`
+	CheckpointSyncEnabled     bool          `json:"checkpoint_sync_enabled,omitempty"`
+	CheckpointParallelChunker bool          `json:"checkpoint_parallel_chunker,omitempty"`
 
 	// Runtimes contains the indexes of the runtimes to enable.
 	Runtimes []int `json:"runtimes,omitempty"`
@@ -434,10 +435,11 @@ func (f *ComputeWorkerFixture) Create(net *Network) (*Compute, error) {
 			Entity:                      entity,
 			ExtraArgs:                   f.ExtraArgs,
 		},
-		RuntimeProvisioner:      f.RuntimeProvisioner,
-		StorageBackend:          f.StorageBackend,
-		SentryIndices:           f.Sentries,
-		CheckpointCheckInterval: f.CheckpointCheckInterval,
+		RuntimeProvisioner:        f.RuntimeProvisioner,
+		StorageBackend:            f.StorageBackend,
+		SentryIndices:             f.Sentries,
+		CheckpointCheckInterval:   f.CheckpointCheckInterval,
+		CheckpointParallelChunker: f.CheckpointParallelChunker,
 		// The checkpoint syncing flag is intentionally flipped here.
 		// Syncing should normally be enabled, but normally disabled in tests.
 		CheckpointSyncDisabled: !f.CheckpointSyncEnabled,

go/oasis-test-runner/scenario/e2e/runtime/scenario.go

@@ -264,7 +264,13 @@ func (sc *Scenario) Fixture() (*oasis.NetworkFixture, error) {
 					},
 				},
 			},
-			{RuntimeProvisioner: runtimeProvisioner, Entity: 1, Runtimes: []int{1}},
+			{
+				RuntimeProvisioner: runtimeProvisioner,
+				Entity:             1,
+				Runtimes:           []int{1},
+				// Use a parallel chunking algorithm for the second compute node to show interoperability.
+				CheckpointParallelChunker: true,
+			},
 		},
 		Sentries: []oasis.SentryFixture{},
 		Seeds:    []oasis.SeedFixture{{}},