feature: Implement data column support for different storage layouts (#15014)

* Implement data column support for different storage layouts

* Fix errors

* Fix linting

* `slotFromFile`: First try to decode as a data column.

---------

Co-authored-by: Manu NALEPA <[email protected]>

Author: 0x00101010

beacon-chain/blockchain/kzg/BUILD.bazel

@@ -30,7 +30,8 @@ go_test(
     deps = [
         "//consensus-types/blocks:go_default_library",
         "//testing/require:go_default_library",
-        "//testing/util:go_default_library",
+        "@com_github_consensys_gnark_crypto//ecc/bls12-381/fr:go_default_library",
         "@com_github_crate_crypto_go_kzg_4844//:go_default_library",
+        "@com_github_sirupsen_logrus//:go_default_library",
     ],
 )

beacon-chain/blockchain/kzg/validation_test.go

@@ -1,12 +1,16 @@
 package kzg
 
 import (
+	"bytes"
+	"crypto/sha256"
+	"encoding/binary"
 	"testing"
 
+	"github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
 	GoKZG "github.com/crate-crypto/go-kzg-4844"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/blocks"
 	"github.com/prysmaticlabs/prysm/v5/testing/require"
-	"github.com/prysmaticlabs/prysm/v5/testing/util"
+	"github.com/sirupsen/logrus"
 )
 
 func GenerateCommitmentAndProof(blob GoKZG.Blob) (GoKZG.KZGCommitment, GoKZG.KZGProof, error) {
@@ -37,11 +41,44 @@ func TestBytesToAny(t *testing.T) {
 }
 
 func TestGenerateCommitmentAndProof(t *testing.T) {
-	blob := util.GetRandBlob(123)
+	blob := getRandBlob(123)
 	commitment, proof, err := GenerateCommitmentAndProof(blob)
 	require.NoError(t, err)
 	expectedCommitment := GoKZG.KZGCommitment{180, 218, 156, 194, 59, 20, 10, 189, 186, 254, 132, 93, 7, 127, 104, 172, 238, 240, 237, 70, 83, 89, 1, 152, 99, 0, 165, 65, 143, 62, 20, 215, 230, 14, 205, 95, 28, 245, 54, 25, 160, 16, 178, 31, 232, 207, 38, 85}
 	expectedProof := GoKZG.KZGProof{128, 110, 116, 170, 56, 111, 126, 87, 229, 234, 211, 42, 110, 150, 129, 206, 73, 142, 167, 243, 90, 149, 240, 240, 236, 204, 143, 182, 229, 249, 81, 27, 153, 171, 83, 70, 144, 250, 42, 1, 188, 215, 71, 235, 30, 7, 175, 86}
 	require.Equal(t, expectedCommitment, commitment)
 	require.Equal(t, expectedProof, proof)
 }
+
+func deterministicRandomness(seed int64) [32]byte {
+	// Converts an int64 to a byte slice
+	buf := new(bytes.Buffer)
+	err := binary.Write(buf, binary.BigEndian, seed)
+	if err != nil {
+		logrus.WithError(err).Error("Failed to write int64 to bytes buffer")
+		return [32]byte{}
+	}
+	bytes := buf.Bytes()
+
+	return sha256.Sum256(bytes)
+}
+
+// Returns a serialized random field element in big-endian
+func getRandFieldElement(seed int64) [32]byte {
+	bytes := deterministicRandomness(seed)
+	var r fr.Element
+	r.SetBytes(bytes[:])
+
+	return GoKZG.SerializeScalar(r)
+}
+
+// Returns a random blob using the passed seed as entropy
+func getRandBlob(seed int64) GoKZG.Blob {
+	var blob GoKZG.Blob
+	bytesPerBlob := GoKZG.ScalarsPerBlob * GoKZG.SerializedScalarSize
+	for i := 0; i < bytesPerBlob; i += GoKZG.SerializedScalarSize {
+		fieldElementBytes := getRandFieldElement(seed + int64(i))
+		copy(blob[i:i+GoKZG.SerializedScalarSize], fieldElementBytes[:])
+	}
+	return blob
+}

beacon-chain/db/filesystem/BUILD.bazel

@@ -26,6 +26,7 @@ go_library(
         "//consensus-types/primitives:go_default_library",
         "//encoding/bytesutil:go_default_library",
         "//io/file:go_default_library",
+        "//proto/prysm/v1alpha1:go_default_library",
         "//runtime/logging:go_default_library",
         "//time/slots:go_default_library",
         "@com_github_ethereum_go_ethereum//common/hexutil:go_default_library",
@@ -49,10 +50,12 @@ go_test(
     ],
     embed = [":go_default_library"],
     deps = [
+        "//beacon-chain/blockchain/kzg:go_default_library",
         "//beacon-chain/db:go_default_library",
         "//beacon-chain/verification:go_default_library",
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",
+        "//consensus-types/blocks:go_default_library",
         "//consensus-types/primitives:go_default_library",
         "//encoding/bytesutil:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",

beacon-chain/db/filesystem/blob.go

@@ -401,7 +401,7 @@ func (bs *BlobStorage) GetColumn(root [32]byte, idx uint64) (blocks.VerifiedRODa
 	return verification.VerifiedRODataColumnFromDisk(bs.fs, root, bs.layout.sszPath(ident))
 }
 
-// Remove removes all blobs for a given root.
+// Remove removes all blobs or data columns for a given root.
 func (bs *BlobStorage) Remove(root [32]byte) error {
 	dirIdent, err := bs.layout.dirIdent(root)
 	if err != nil {

beacon-chain/db/filesystem/blob_test.go

@@ -2,16 +2,19 @@ package filesystem
 
 import (
 	"bytes"
+	"fmt"
 	"math"
 	"os"
 	"path"
 	"sync"
 	"testing"
 
 	ssz "github.com/prysmaticlabs/fastssz"
+	"github.com/prysmaticlabs/prysm/v5/beacon-chain/blockchain/kzg"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/db"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/verification"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
+	"github.com/prysmaticlabs/prysm/v5/consensus-types/blocks"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
 	ethpb "github.com/prysmaticlabs/prysm/v5/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/v5/testing/require"
@@ -56,7 +59,7 @@ func TestBlobStorage_SaveBlobData(t *testing.T) {
 		require.NoError(t, bs.Save(sc))
 		actualSc, err := bs.Get(sc.BlockRoot(), sc.Index)
 		require.NoError(t, err)
-		expectedIdx := blobIndexMask{false, false, true, false, false, false}
+		expectedIdx := dataIndexMask{false, false, true, false, false, false}
 		actualIdx := bs.Summary(actualSc.BlockRoot()).mask
 		require.NoError(t, err)
 		require.DeepEqual(t, expectedIdx, actualIdx)
@@ -221,3 +224,158 @@ func TestLayoutNames(t *testing.T) {
 	_, err := newLayout(badLayoutName, nil, nil, nil)
 	require.ErrorIs(t, err, errInvalidLayoutName)
 }
+
+func TestBlobStorage_DataColumn_WithAllLayouts(t *testing.T) {
+	for _, layout := range LayoutNames {
+		t.Run(fmt.Sprintf("layout=%s", layout), func(t *testing.T) {
+			sidecars := setupDataColumnTest(t)
+
+			t.Run("no error for duplicate", func(t *testing.T) {
+				fs, bs := NewEphemeralBlobStorageAndFs(t, WithLayout(layout))
+				sidecar := sidecars[0]
+
+				columnPath := bs.layout.sszPath(identForDataColumnSidecar(sidecar))
+				data, err := ssz.MarshalSSZ(sidecar)
+				require.NoError(t, err)
+
+				require.NoError(t, bs.SaveDataColumn(sidecar))
+				// No error when attempting to write twice.
+				require.NoError(t, bs.SaveDataColumn(sidecar))
+
+				content, err := afero.ReadFile(fs, columnPath)
+				require.NoError(t, err)
+				require.Equal(t, true, bytes.Equal(data, content))
+
+				// Deserialize the DataColumnSidecar from the saved file data.
+				saved := &ethpb.DataColumnSidecar{}
+				err = saved.UnmarshalSSZ(content)
+				require.NoError(t, err)
+
+				// Compare the original Sidecar and the saved Sidecar.
+				require.DeepSSZEqual(t, sidecar.DataColumnSidecar, saved)
+			})
+
+			t.Run("indices", func(t *testing.T) {
+				bs := NewEphemeralBlobStorage(t, WithLayout(layout))
+				sidecar := sidecars[2]
+				require.NoError(t, bs.SaveDataColumn(sidecar))
+				actual, err := bs.GetColumn(sidecar.BlockRoot(), sidecar.ColumnIndex)
+				require.NoError(t, err)
+				require.DeepEqual(t, sidecar, actual)
+
+				expectedIdx := make(dataIndexMask, params.BeaconConfig().NumberOfColumns)
+				expectedIdx[2] = true
+				actualIdx := bs.Summary(actual.BlockRoot()).mask
+				require.NoError(t, err)
+				require.DeepEqual(t, expectedIdx, actualIdx)
+
+				sidecar = sidecars[10]
+				expectedIdx[10] = true
+				require.NoError(t, bs.SaveDataColumn(sidecar))
+				actual, err = bs.GetColumn(sidecar.BlockRoot(), sidecar.ColumnIndex)
+				require.NoError(t, err)
+				require.DeepEqual(t, sidecar, actual)
+				actualIdx = bs.Summary(actual.BlockRoot()).mask
+				require.NoError(t, err)
+				require.DeepEqual(t, expectedIdx, actualIdx)
+			})
+
+			t.Run("write -> read -> delete", func(t *testing.T) {
+				bs := NewEphemeralBlobStorage(t, WithLayout(layout))
+				err := bs.SaveDataColumn(sidecars[0])
+				require.NoError(t, err)
+
+				expected := sidecars[0]
+				actual, err := bs.GetColumn(expected.BlockRoot(), expected.ColumnIndex)
+				require.NoError(t, err)
+				require.DeepEqual(t, expected, actual)
+
+				require.NoError(t, bs.Remove(expected.BlockRoot()))
+				for i := range params.BeaconConfig().NumberOfColumns {
+					_, err = bs.GetColumn(expected.BlockRoot(), uint64(i))
+					require.Equal(t, true, db.IsNotFound(err))
+				}
+			})
+
+			t.Run("clear", func(t *testing.T) {
+				bs := NewEphemeralBlobStorage(t, WithLayout(layout))
+				err := bs.SaveDataColumn(sidecars[0])
+				require.NoError(t, err)
+				res, err := bs.GetColumn(sidecars[0].BlockRoot(), sidecars[0].ColumnIndex)
+				require.NoError(t, err)
+				require.NotNil(t, res)
+				require.NoError(t, bs.Clear())
+				// After clearing, the blob should not exist in the db.
+				_, err = bs.GetColumn(sidecars[0].BlockRoot(), sidecars[0].ColumnIndex)
+				require.ErrorIs(t, err, os.ErrNotExist)
+			})
+		})
+	}
+}
+
+func TestBlobStorage_DataColumn_WithMigrationFromFlatToByEpoch(t *testing.T) {
+	sidecars := setupDataColumnTest(t)
+
+	// Setup flat layout
+	fs, bs := NewEphemeralBlobStorageAndFs(t, WithLayout(LayoutNameFlat))
+	sidecar := sidecars[0]
+	columnPath := bs.layout.sszPath(identForDataColumnSidecar(sidecar))
+	data, err := ssz.MarshalSSZ(sidecar)
+	require.NoError(t, err)
+	require.NoError(t, bs.SaveDataColumn(sidecar))
+	content, err := afero.ReadFile(fs, columnPath)
+	require.NoError(t, err)
+	require.Equal(t, true, bytes.Equal(data, content))
+
+	// Setup by-epoch layout
+	bs = NewWarmedEphemeralBlobStorageUsingFs(t, fs, WithLayout(LayoutNameByEpoch))
+
+	// Verify data is the same
+	columnPath = bs.layout.sszPath(identForDataColumnSidecar(sidecar))
+	content, err = afero.ReadFile(fs, columnPath)
+	require.NoError(t, err)
+	require.Equal(t, true, bytes.Equal(data, content))
+}
+
+func TestBlobStorage_DataColumn_WithMigrationFromByEpochToFlat(t *testing.T) {
+	sidecars := setupDataColumnTest(t)
+
+	// Setup by-epoch layout
+	fs, bs := NewEphemeralBlobStorageAndFs(t, WithLayout(LayoutNameFlat))
+	for _, sidecar := range sidecars {
+		require.NoError(t, bs.SaveDataColumn(sidecar))
+	}
+	columnPath := bs.layout.sszPath(identForDataColumnSidecar(sidecars[0]))
+	content, err := afero.ReadFile(fs, columnPath)
+	require.NoError(t, err)
+	data, err := ssz.MarshalSSZ(sidecars[0])
+	require.NoError(t, err)
+	require.Equal(t, true, bytes.Equal(data, content))
+
+	// Setup flat layout
+	bs = NewWarmedEphemeralBlobStorageUsingFs(t, fs, WithLayout(LayoutNameByEpoch))
+
+	// Verify data is the same
+	columnPath = bs.layout.sszPath(identForDataColumnSidecar(sidecars[0]))
+	content, err = afero.ReadFile(fs, columnPath)
+	require.NoError(t, err)
+	require.Equal(t, true, bytes.Equal(data, content))
+}
+
+func setupDataColumnTest(t *testing.T) []blocks.VerifiedRODataColumn {
+	// load trusted setup
+	err := kzg.Start()
+	require.NoError(t, err)
+
+	// Setup right fork epoch
+	params.SetupTestConfigCleanup(t)
+	cfg := params.BeaconConfig().Copy()
+	cfg.CapellaForkEpoch = 0
+	cfg.DenebForkEpoch = 0
+	cfg.ElectraForkEpoch = 0
+	cfg.FuluForkEpoch = 0
+	params.OverrideBeaconConfig(cfg)
+
+	_, scs := util.GenerateTestFuluBlockWithSidecar(t, [32]byte{}, 0, 1)
+	return verification.FakeVerifyDataColumnSliceForTest(t, scs)
+}

beacon-chain/db/filesystem/cache.go

@@ -10,27 +10,24 @@ import (
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
 )
 
-// blobIndexMask is a bitmask representing the set of blob indices that are currently set.
-// TODO: Separate blobs from data columns
-type blobIndexMask []bool
+// dataIndexMask is a bitmask representing the set of blob or data column indices that are currently set.
+type dataIndexMask []bool
 
-// type blobIndexMask [fieldparams.NumberOfColumns]bool
-
-// BlobStorageSummary represents cached information about the BlobSidecars on disk for each root the cache knows about.
+// BlobStorageSummary represents cached information about the BlobSidecars or DataColumnSidecars on disk for each root the cache knows about.
 type BlobStorageSummary struct {
 	epoch primitives.Epoch
-	mask  blobIndexMask
+	mask  dataIndexMask
 }
 
-// HasIndex returns true if the BlobSidecar at the given index is available in the filesystem.
+// HasIndex returns true if the BlobSidecar or DataColumnSidecar at the given index is available in the filesystem.
 func (s BlobStorageSummary) HasIndex(idx uint64) bool {
 	if idx >= uint64(len(s.mask)) {
 		return false
 	}
 	return s.mask[idx]
 }
 
-// HasDataColumnIndex true if the DataColumnSidecar at the given index is available in the filesystem.
+// HasDataColumnIndex returns true if the DataColumnSidecar at the given index is available in the filesystem.
 func (s BlobStorageSummary) HasDataColumnIndex(idx uint64) bool {
 	// Protect from panic, but assume callers are sophisticated enough to not need an error telling them they have an invalid idx.
 	numberOfColumns := params.BeaconConfig().NumberOfColumns
@@ -134,7 +131,7 @@ func (s *blobStorageSummaryCache) ensure(ident blobIdent) error {
 	v := s.cache[ident.root]
 	v.epoch = ident.epoch
 	if v.mask == nil {
-		v.mask = make(blobIndexMask, maskSize)
+		v.mask = make(dataIndexMask, maskSize)
 	}
 	if !v.mask[ident.index] {
 		s.updateMetrics(1)

beacon-chain/db/filesystem/cache_test.go

@@ -24,7 +24,7 @@ func TestSlotByRoot_Summary(t *testing.T) {
 	cases := []struct {
 		name     string
 		root     [32]byte
-		expected blobIndexMask
+		expected dataIndexMask
 	}{
 		{
 			name: "not found",
@@ -193,7 +193,7 @@ func TestHasDataColumnIndex(t *testing.T) {
 				}
 			}
 
-			mask := make(blobIndexMask, maxIndex+1)
+			mask := make(dataIndexMask, maxIndex+1)
 
 			for idx := range storedIndices {
 				mask[idx] = true
@@ -263,7 +263,7 @@ func TestAllDataColumnAvailable(t *testing.T) {
 				}
 			}
 
-			mask := make(blobIndexMask, maxIndex+1)
+			mask := make(dataIndexMask, maxIndex+1)
 
 			for idx := range c.storedIndices {
 				mask[idx] = true

beacon-chain/db/filesystem/iteration_test.go

@@ -1,7 +1,6 @@
 package filesystem
 
 import (
-	"bytes"
 	"fmt"
 	"io"
 	"math"
@@ -220,28 +219,6 @@ func TestListDir(t *testing.T) {
 	}
 }
 
-func TestSlotFromBlob(t *testing.T) {
-	cases := []struct {
-		slot primitives.Slot
-	}{
-		{slot: 0},
-		{slot: 2},
-		{slot: 1123581321},
-		{slot: math.MaxUint64},
-	}
-	for _, c := range cases {
-		t.Run(fmt.Sprintf("slot %d", c.slot), func(t *testing.T) {
-			_, sidecars := util.GenerateTestDenebBlockWithSidecar(t, [32]byte{}, c.slot, 1)
-			sc := sidecars[0]
-			enc, err := sc.MarshalSSZ()
-			require.NoError(t, err)
-			slot, err := slotFromBlob(bytes.NewReader(enc))
-			require.NoError(t, err)
-			require.Equal(t, c.slot, slot)
-		})
-	}
-}
-
 func TestIterationComplete(t *testing.T) {
 	targets := []migrationTestTarget{
 		{