Avoid unnecessary database lookups in data column RPC requests (#7897)

This PR is an optimisation to avoid unnecessary database lookups when peer requests data columns that the node doesn't custody (advertised via `cgc`).

e.g. an extreme but realistic example - a full node only store 4 custody columns by default, but it may receive a range request of 32 slots with all 128 columns, and this would result in 4096 database lookups but the node is only able to get 128 (4 * 32) of them.


  - Filter data column RPC requests (`DataColumnsByRoot`, `DataColumnsByRange`) to only lookup columns the node custodies
- Prevents unnecessary database queries that would always fail for non-custody columns

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -7153,6 +7153,17 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             .custody_context()
             .sampling_columns_for_epoch(epoch, &self.spec)
     }
+
+    /// Returns a list of column indices that the node is expected to custody for a given epoch.
+    /// i.e. the node must have validated and persisted the column samples and should be able to
+    /// serve them to peers.
+    ///
+    /// If epoch is `None`, this function computes the custody columns at head.
+    pub fn custody_columns_for_epoch(&self, epoch_opt: Option<Epoch>) -> &[ColumnIndex] {
+        self.data_availability_checker
+            .custody_context()
+            .custody_columns_for_epoch(epoch_opt, &self.spec)
+    }
 }
 
 impl<T: BeaconChainTypes> Drop for BeaconChain<T> {

beacon_node/beacon_chain/src/validator_custody.rs

@@ -323,6 +323,37 @@ impl<E: EthSpec> CustodyContext<E> {
             .expect("all_custody_columns_ordered should be initialized");
         &all_columns_ordered[..num_of_columns_to_sample]
     }
+
+    /// Returns the ordered list of column indices that the node is assigned to custody
+    /// (and advertised to peers) at the given epoch. If epoch is `None`, this function
+    /// computes the custody columns at head.
+    ///
+    /// This method differs from [`self::sampling_columns_for_epoch`] which returns all sampling columns.
+    /// The columns returned by this method are either identical to or a subset of the sampling columns,
+    /// representing only those columns that this node is responsible for maintaining custody of.
+    ///
+    /// # Parameters
+    /// * `epoch_opt` - Optional epoch to determine custody columns for.
+    ///
+    /// # Returns
+    /// A slice of ordered custody column indices for this epoch based on the node's custody configuration
+    pub fn custody_columns_for_epoch(
+        &self,
+        epoch_opt: Option<Epoch>,
+        spec: &ChainSpec,
+    ) -> &[ColumnIndex] {
+        let custody_group_count = if let Some(epoch) = epoch_opt {
+            self.custody_group_count_at_epoch(epoch, spec) as usize
+        } else {
+            self.custody_group_count_at_head(spec) as usize
+        };
+
+        let all_columns_ordered = self
+            .all_custody_columns_ordered
+            .get()
+            .expect("all_custody_columns_ordered should be initialized");
+        &all_columns_ordered[..custody_group_count]
+    }
 }
 
 /// The custody count changed because of a change in the
@@ -670,4 +701,82 @@ mod tests {
             assert_eq!(updated_custody_count_opt, expected_cgc_change);
         }
     }
+
+    #[test]
+    fn custody_columns_for_epoch_no_validators_fullnode() {
+        let custody_context = CustodyContext::<E>::new(false);
+        let spec = E::default_spec();
+        let all_custody_groups_ordered = (0..spec.number_of_custody_groups).collect::<Vec<_>>();
+
+        custody_context
+            .init_ordered_data_columns_from_custody_groups(all_custody_groups_ordered, &spec)
+            .expect("should initialise ordered data columns");
+
+        assert_eq!(
+            custody_context.custody_columns_for_epoch(None, &spec).len(),
+            spec.custody_requirement as usize
+        );
+    }
+
+    #[test]
+    fn custody_columns_for_epoch_no_validators_supernode() {
+        let custody_context = CustodyContext::<E>::new(true);
+        let spec = E::default_spec();
+        let all_custody_groups_ordered = (0..spec.number_of_custody_groups).collect::<Vec<_>>();
+
+        custody_context
+            .init_ordered_data_columns_from_custody_groups(all_custody_groups_ordered, &spec)
+            .expect("should initialise ordered data columns");
+
+        assert_eq!(
+            custody_context.custody_columns_for_epoch(None, &spec).len(),
+            spec.number_of_custody_groups as usize
+        );
+    }
+
+    #[test]
+    fn custody_columns_for_epoch_with_validators_should_match_cgc() {
+        let custody_context = CustodyContext::<E>::new(false);
+        let spec = E::default_spec();
+        let all_custody_groups_ordered = (0..spec.number_of_custody_groups).collect::<Vec<_>>();
+        let val_custody_units = 10;
+
+        custody_context
+            .init_ordered_data_columns_from_custody_groups(all_custody_groups_ordered, &spec)
+            .expect("should initialise ordered data columns");
+
+        let _ = custody_context.register_validators(
+            vec![(
+                0,
+                val_custody_units * spec.balance_per_additional_custody_group,
+            )],
+            Slot::new(10),
+            &spec,
+        );
+
+        assert_eq!(
+            custody_context.custody_columns_for_epoch(None, &spec).len(),
+            val_custody_units as usize
+        );
+    }
+
+    #[test]
+    fn custody_columns_for_epoch_specific_epoch_uses_epoch_cgc() {
+        let custody_context = CustodyContext::<E>::new(false);
+        let spec = E::default_spec();
+        let all_custody_groups_ordered = (0..spec.number_of_custody_groups).collect::<Vec<_>>();
+        let test_epoch = Epoch::new(5);
+
+        custody_context
+            .init_ordered_data_columns_from_custody_groups(all_custody_groups_ordered, &spec)
+            .expect("should initialise ordered data columns");
+
+        let expected_cgc = custody_context.custody_group_count_at_epoch(test_epoch, &spec);
+        assert_eq!(
+            custody_context
+                .custody_columns_for_epoch(Some(test_epoch), &spec)
+                .len(),
+            expected_cgc as usize
+        );
+    }
 }