feat(op-program): Use `PebbleDB` for `DiskKV` (#11705)

* feat(op-program): Use `PebbleDB` for `DiskKV`

* close db

* fix `testFaultProofProgramScenario` tests

* switch to snappy compression

https://github.com/cockroachdb/pebble/issues/3434

* fix tempdir

* update compat release

* defer k/v until preimage server and hinter have both exited

op-e2e/system_fpp_test.go

@@ -86,6 +86,7 @@ func applySpanBatchActivation(active bool, dp *genesis.DeployConfig) {
 // - update the state root via a tx
 // - run program
 func testVerifyL2OutputRootEmptyBlock(t *testing.T, detached bool, spanBatchActivated bool) {
+	t.Helper()
 	InitParallel(t)
 	ctx := context.Background()
 
@@ -186,6 +187,7 @@ func testVerifyL2OutputRootEmptyBlock(t *testing.T, detached bool, spanBatchActi
 }
 
 func testVerifyL2OutputRoot(t *testing.T, detached bool, spanBatchActivated bool) {
+	t.Helper()
 	InitParallel(t)
 	ctx := context.Background()
 
@@ -278,6 +280,7 @@ type FaultProofProgramTestScenario struct {
 // testFaultProofProgramScenario runs the fault proof program in several contexts, given a test scenario.
 func testFaultProofProgramScenario(t *testing.T, ctx context.Context, sys *System, s *FaultProofProgramTestScenario) {
 	preimageDir := t.TempDir()
+
 	fppConfig := oppconf.NewConfig(sys.RollupConfig, sys.L2GenesisCfg.Config, s.L1Head, s.L2Head, s.L2OutputRoot, common.Hash(s.L2Claim), s.L2ClaimBlockNumber)
 	fppConfig.L1URL = sys.NodeEndpoint("l1").RPC()
 	fppConfig.L2URL = sys.NodeEndpoint("sequencer").RPC()

op-program/host/host.go

@@ -122,6 +122,10 @@ func FaultProofProgram(ctx context.Context, logger log.Logger, cfg *config.Confi
 func PreimageServer(ctx context.Context, logger log.Logger, cfg *config.Config, preimageChannel preimage.FileChannel, hintChannel preimage.FileChannel) error {
 	var serverDone chan error
 	var hinterDone chan error
+	logger.Info("Starting preimage server")
+	var kv kvstore.KV
+
+	// Close the preimage/hint channels, and then kv store once the server and hinter have exited.
 	defer func() {
 		preimageChannel.Close()
 		hintChannel.Close()
@@ -133,9 +137,12 @@ func PreimageServer(ctx context.Context, logger log.Logger, cfg *config.Config,
 			// Wait for hinter to complete
 			<-hinterDone
 		}
+
+		if kv != nil {
+			kv.Close()
+		}
 	}()
-	logger.Info("Starting preimage server")
-	var kv kvstore.KV
+
 	if cfg.DataDir == "" {
 		logger.Info("Using in-memory storage")
 		kv = kvstore.NewMemKV()

op-program/host/kvstore/disk.go

 package kvstore
 
 import (
-	"encoding/hex"
 	"errors"
 	"fmt"
-	"io"
-	"os"
-	"path"
+	"runtime"
 	"sync"
 
+	"github.com/cockroachdb/pebble"
 	"github.com/ethereum/go-ethereum/common"
 )
 
-// read/write mode for user/group/other, not executable.
-const diskPermission = 0666
-
-// DiskKV is a disk-backed key-value store, every key-value pair is a hex-encoded .txt file, with the value as content.
+// DiskKV is a disk-backed key-value store, with PebbleDB as the underlying DBMS.
 // DiskKV is safe for concurrent use with a single DiskKV instance.
-// DiskKV is safe for concurrent use between different DiskKV instances of the same disk directory as long as the
-// file system supports atomic renames.
 type DiskKV struct {
 	sync.RWMutex
-	path string
+	db *pebble.DB
 }
 
 // NewDiskKV creates a DiskKV that puts/gets pre-images as files in the given directory path.
 // The path must exist, or subsequent Put/Get calls will error when it does not.
 func NewDiskKV(path string) *DiskKV {
-	return &DiskKV{path: path}
-}
+	opts := &pebble.Options{
+		Cache:                    pebble.NewCache(int64(32 * 1024 * 1024)),
+		MaxConcurrentCompactions: runtime.NumCPU,
+		Levels: []pebble.LevelOptions{
+			{Compression: pebble.SnappyCompression},
+		},
+	}
+	db, err := pebble.Open(path, opts)
+	if err != nil {
+		panic(fmt.Errorf("failed to open pebbledb at %s: %w", path, err))
+	}
 
-func (d *DiskKV) pathKey(k common.Hash) string {
-	return path.Join(d.path, k.String()+".txt")
+	return &DiskKV{db: db}
 }
 
 func (d *DiskKV) Put(k common.Hash, v []byte) error {
 	d.Lock()
 	defer d.Unlock()
-	f, err := openTempFile(d.path, k.String()+".txt.*")
-	if err != nil {
-		return fmt.Errorf("failed to open temp file for pre-image %s: %w", k, err)
-	}
-	defer os.Remove(f.Name()) // Clean up the temp file if it doesn't actually get moved into place
-	if _, err := f.Write([]byte(hex.EncodeToString(v))); err != nil {
-		_ = f.Close()
-		return fmt.Errorf("failed to write pre-image %s to disk: %w", k, err)
-	}
-	if err := f.Close(); err != nil {
-		return fmt.Errorf("failed to close temp pre-image %s file: %w", k, err)
-	}
-
-	targetFile := d.pathKey(k)
-	if err := os.Rename(f.Name(), targetFile); err != nil {
-		return fmt.Errorf("failed to move temp dir %v to final destination %v: %w", f.Name(), targetFile, err)
-	}
-	return nil
-}
-
-func openTempFile(dir string, nameTemplate string) (*os.File, error) {
-	f, err := os.CreateTemp(dir, nameTemplate)
-	// Directory has been deleted out from underneath us. Recreate it.
-	if errors.Is(err, os.ErrNotExist) {
-		if mkdirErr := os.MkdirAll(dir, 0777); mkdirErr != nil {
-			return nil, errors.Join(fmt.Errorf("failed to create directory %v: %w", dir, mkdirErr), err)
-		}
-		f, err = os.CreateTemp(dir, nameTemplate)
-	}
-	if err != nil {
-		return nil, err
-	}
-	return f, nil
+	return d.db.Set(k.Bytes(), v, pebble.NoSync)
 }
 
 func (d *DiskKV) Get(k common.Hash) ([]byte, error) {
 	d.RLock()
 	defer d.RUnlock()
-	f, err := os.OpenFile(d.pathKey(k), os.O_RDONLY, diskPermission)
+
+	dat, closer, err := d.db.Get(k.Bytes())
 	if err != nil {
-		if errors.Is(err, os.ErrNotExist) {
+		if errors.Is(err, pebble.ErrNotFound) {
 			return nil, ErrNotFound
 		}
-		return nil, fmt.Errorf("failed to open pre-image file %s: %w", k, err)
-	}
-	defer f.Close() // fine to ignore closing error here
-	dat, err := io.ReadAll(f)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read pre-image from file %s: %w", k, err)
+		return nil, err
 	}
-	return hex.DecodeString(string(dat))
+	ret := make([]byte, len(dat))
+	copy(ret, dat)
+	closer.Close()
+	return ret, nil
+}
+
+func (d *DiskKV) Close() error {
+	d.Lock()
+	defer d.Unlock()
+
+	return d.db.Close()
 }
 
 var _ KV = (*DiskKV)(nil)

op-program/host/kvstore/kv.go

@@ -19,4 +19,7 @@ type KV interface {
 	// It returns ErrNotFound when the pre-image cannot be found.
 	// KV store implementations may return additional errors specific to the KV storage.
 	Get(k common.Hash) ([]byte, error)
+
+	// Closes the KV store.
+	Close() error
 }

op-program/host/kvstore/mem.go

@@ -37,3 +37,7 @@ func (m *MemKV) Get(k common.Hash) ([]byte, error) {
 	}
 	return slices.Clone(v), nil
 }
+
+func (m *MemKV) Close() error {
+	return nil
+}

op-program/scripts/run-compat.sh

@@ -5,7 +5,7 @@ SCRIPTS_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 COMPAT_DIR="${SCRIPTS_DIR}/../temp/compat"
 
 TESTNAME="${1?Must specify compat file to run}"
-BASEURL="${2:-https://github.com/ethereum-optimism/chain-test-data/releases/download/2024-08-02}"
+BASEURL="${2:-https://github.com/ethereum-optimism/chain-test-data/releases/download/2024-09-01}"
 
 URL="${BASEURL}/${TESTNAME}.tar.bz"