Merge pull request #7060 from ethereum-optimism/aj/atomic-put

op-program: Use atomic write pattern for KV store

op-program/host/kvstore/disk.go

@@ -17,8 +17,8 @@ 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 safe for concurrent use with a single DiskKV instance.
-// DiskKV is not safe for concurrent use between different DiskKV instances of the same disk directory:
-// a Put needs to be completed before another DiskKV Get retrieves the values.
+// 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
@@ -37,19 +37,22 @@ func (d *DiskKV) pathKey(k common.Hash) string {
 func (d *DiskKV) Put(k common.Hash, v []byte) error {
 	d.Lock()
 	defer d.Unlock()
-	f, err := os.OpenFile(d.pathKey(k), os.O_WRONLY|os.O_CREATE|os.O_EXCL|os.O_TRUNC, diskPermission)
+	f, err := os.CreateTemp(d.path, k.String()+".txt.*")
 	if err != nil {
-		if errors.Is(err, os.ErrExist) {
-			return ErrAlreadyExists
-		}
-		return fmt.Errorf("failed to open new pre-image file %s: %w", k, err)
+		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 pre-image %s file: %w", k, err)
+		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
 }

op-program/host/kvstore/kv.go

@@ -9,13 +9,9 @@ import (
 // ErrNotFound is returned when a pre-image cannot be found in the KV store.
 var ErrNotFound = errors.New("not found")
 
-// ErrAlreadyExists is returned when a pre-image already exists in the KV store.
-var ErrAlreadyExists = errors.New("already exists")
-
 // KV is a Key-Value store interface for pre-image data.
 type KV interface {
 	// Put puts the pre-image value v in the key-value store with key k.
-	// It returns ErrAlreadyExists when the key already exists.
 	// KV store implementations may return additional errors specific to the KV storage.
 	Put(k common.Hash, v []byte) error
 

op-program/host/kvstore/kv_test.go

@@ -45,9 +45,9 @@ func kvTest(t *testing.T, kv KV) {
 		require.Equal(t, []byte{4, 2}, dat, "pre-image must match")
 	})
 
-	t.Run("not overwriting pre-image", func(t *testing.T) {
+	t.Run("allowing multiple writes for same pre-image", func(t *testing.T) {
 		t.Parallel()
 		require.NoError(t, kv.Put(common.Hash{0xdd}, []byte{4, 2}))
-		require.ErrorIs(t, kv.Put(common.Hash{0xdd}, []byte{4, 2}), ErrAlreadyExists)
+		require.NoError(t, kv.Put(common.Hash{0xdd}, []byte{4, 2}))
 	})
 }

op-program/host/kvstore/mem.go

@@ -23,9 +23,6 @@ func NewMemKV() *MemKV {
 func (m *MemKV) Put(k common.Hash, v []byte) error {
 	m.Lock()
 	defer m.Unlock()
-	if _, ok := m.m[k]; ok {
-		return ErrAlreadyExists
-	}
 	m.m[k] = v
 	return nil
 }

op-program/host/prefetcher/prefetcher.go

@@ -155,10 +155,7 @@ func (p *Prefetcher) storeTrieNodes(values []hexutil.Bytes) error {
 	_, nodes := mpt.WriteTrie(values)
 	for _, node := range nodes {
 		key := preimage.Keccak256Key(crypto.Keccak256Hash(node)).PreimageKey()
-		if err := p.kvStore.Put(key, node); errors.Is(err, kvstore.ErrAlreadyExists) {
-			// It's not uncommon for different tries to contain common nodes (esp for receipts)
-			continue
-		} else if err != nil {
+		if err := p.kvStore.Put(key, node); err != nil {
 			return fmt.Errorf("failed to store node: %w", err)
 		}
 	}