Use peerdas.CustodyColumnSubnets and peerdas.CustodyColumns.

prysmaticlabs · Apr 24, 2024 · 0c92512 · 0c92512
1 parent b983b33
commit 0c92512
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 51 deletions.
diff --git a/beacon-chain/p2p/BUILD.bazel b/beacon-chain/p2p/BUILD.bazel
@@ -46,6 +46,7 @@ go_library(
         "//beacon-chain/core/altair:go_default_library",
         "//beacon-chain/core/feed/state:go_default_library",
         "//beacon-chain/core/helpers:go_default_library",
+        "//beacon-chain/core/peerdas:go_default_library",
         "//beacon-chain/core/time:go_default_library",
         "//beacon-chain/db:go_default_library",
         "//beacon-chain/p2p/encoder:go_default_library",
@@ -60,7 +61,6 @@ go_library(
         "//consensus-types/primitives:go_default_library",
         "//consensus-types/wrapper:go_default_library",
         "//container/leaky-bucket:go_default_library",
-        "//container/slice:go_default_library",
         "//crypto/ecdsa:go_default_library",
         "//crypto/hash:go_default_library",
         "//encoding/bytesutil:go_default_library",

diff --git a/beacon-chain/p2p/subnets.go b/beacon-chain/p2p/subnets.go
@@ -13,10 +13,10 @@ import (
 	"github.com/prysmaticlabs/go-bitfield"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/helpers"
+	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/peerdas"
 	"github.com/prysmaticlabs/prysm/v5/cmd/beacon-chain/flags"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/wrapper"
-	"github.com/prysmaticlabs/prysm/v5/container/slice"
 	"github.com/prysmaticlabs/prysm/v5/crypto/hash"
 	"github.com/prysmaticlabs/prysm/v5/encoding/bytesutil"
 	mathutil "github.com/prysmaticlabs/prysm/v5/math"
@@ -208,10 +208,16 @@ func initializePersistentColumnSubnets(id enode.ID) error {
 	if ok && expTime.After(time.Now()) {
 		return nil
 	}
-	subs, err := computeSubscribedColumnSubnets(id)
+	subsMap, err := peerdas.CustodyColumnSubnets(id, params.BeaconConfig().CustodyRequirement)
 	if err != nil {
 		return err
 	}
+
+	subs := make([]uint64, 0, len(subsMap))
+	for sub := range subsMap {
+		subs = append(subs, sub)
+	}
+
 	cache.ColumnSubnetIDs.AddColumnSubnets(subs)
 	return nil
 }
@@ -235,46 +241,6 @@ func computeSubscribedSubnets(nodeID enode.ID, epoch primitives.Epoch) ([]uint64
 	return subs, nil
 }
 
-func ComputeCustodyColumns(nodeID enode.ID) ([]uint64, error) {
-	subs, err := computeSubscribedColumnSubnets(nodeID)
-	if err != nil {
-		return nil, err
-	}
-	colsPerSub := params.BeaconConfig().NumberOfColumns / params.BeaconConfig().DataColumnSidecarSubnetCount
-	colIdxs := []uint64{}
-	for _, sub := range subs {
-		for i := uint64(0); i < colsPerSub; i++ {
-			colId := params.BeaconConfig().DataColumnSidecarSubnetCount*i + sub
-			colIdxs = append(colIdxs, colId)
-		}
-	}
-	return colIdxs, nil
-}
-
-func computeSubscribedColumnSubnets(nodeID enode.ID) ([]uint64, error) {
-	subnetsPerNode := params.BeaconConfig().CustodyRequirement
-	subs := make([]uint64, 0, subnetsPerNode)
-
-	for i := uint64(0); i < subnetsPerNode; i++ {
-		sub, err := computeSubscribedColumnSubnet(nodeID, i)
-		if err != nil {
-			return nil, err
-		}
-		if slice.IsInUint64(sub, subs) {
-			continue
-		}
-		subs = append(subs, sub)
-	}
-	isubnetsPerNode, err := mathutil.Int(subnetsPerNode)
-	if err != nil {
-		return nil, err
-	}
-	if len(subs) != isubnetsPerNode {
-		return nil, errors.Errorf("inconsistent subnet assignment: %d vs %d", len(subs), isubnetsPerNode)
-	}
-	return subs, nil
-}
-
 //	Spec pseudocode definition:
 //
 // def compute_subscribed_subnet(node_id: NodeID, epoch: Epoch, index: int) -> SubnetID:

diff --git a/beacon-chain/sync/BUILD.bazel b/beacon-chain/sync/BUILD.bazel
@@ -74,6 +74,7 @@ go_library(
         "//beacon-chain/core/feed/operation:go_default_library",
         "//beacon-chain/core/feed/state:go_default_library",
         "//beacon-chain/core/helpers:go_default_library",
+        "//beacon-chain/core/peerdas:go_default_library",
         "//beacon-chain/core/signing:go_default_library",
         "//beacon-chain/core/transition:go_default_library",
         "//beacon-chain/core/transition/interop:go_default_library",

diff --git a/beacon-chain/sync/rpc_data_column_sidecars_by_root.go b/beacon-chain/sync/rpc_data_column_sidecars_by_root.go
@@ -9,6 +9,7 @@ import (
 
 	libp2pcore "github.com/libp2p/go-libp2p/core"
 	"github.com/pkg/errors"
+	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/peerdas"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/db"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/p2p/types"
@@ -72,19 +73,14 @@ func (s *Service) dataColumnSidecarByRootRPCHandler(ctx context.Context, msg int
 		}
 		s.rateLimiter.add(stream, 1)
 		root, idx := bytesutil.ToBytes32(columnIdents[i].BlockRoot), columnIdents[i].Index
-		custodiedColumns, err := p2p.ComputeCustodyColumns(s.cfg.p2p.NodeID())
+		custodiedColumns, err := peerdas.CustodyColumns(s.cfg.p2p.NodeID(), params.BeaconConfig().CustodyRequirement)
 		if err != nil {
 			log.WithError(err).Errorf("unexpected error retrieving the node id")
 			s.writeErrorResponseToStream(responseCodeServerError, types.ErrGeneric.Error(), stream)
 			return err
 		}
-		isCustodied := false
-		for _, col := range custodiedColumns {
-			if col == idx {
-				isCustodied = true
-				break
-			}
-		}
+
+		isCustodied := custodiedColumns[idx]
 		if !isCustodied {
 			s.cfg.p2p.Peers().Scorers().BadResponsesScorer().Increment(stream.Conn().RemotePeer())
 			s.writeErrorResponseToStream(responseCodeInvalidRequest, types.ErrInvalidColumnIndex.Error(), stream)