track useful shards

Author: VorobyevIlya

crates/starfish/core/src/authority_service.rs

@@ -3,6 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::{
+    cmp::max,
     collections::{BTreeMap, BTreeSet, VecDeque},
     pin::Pin,
     sync::Arc,
@@ -122,6 +123,13 @@ pub(crate) struct AuthorityService<C: CoreThreadDispatcher> {
     /// useful to them (communicated inside their block bundles).
     /// Keyed by the peer’s AuthorityIndex.
     useful_authorities_to_peer: Arc<RwLock<BTreeMap<AuthorityIndex, BTreeSet<AuthorityIndex>>>>,
+    /// For each peer `i`, stores a vector where each entry at index `j`
+    /// represents the last round in which a received bundle from peer `i`
+    /// contained a useful shard from a block created by authority `j`.
+    last_round_with_useful_shards_from_peer: Arc<RwLock<Vec<Vec<Round>>>>,
+    /// For each peer `i`, stores a set of authority indices representing the
+    /// authors of blocks whose shards were reported as useful by peer `i`.
+    peers_reporting_useful_shards: Arc<RwLock<Vec<BTreeSet<AuthorityIndex>>>>,
 }
 
 impl<C: CoreThreadDispatcher> AuthorityService<C> {
@@ -141,6 +149,7 @@ impl<C: CoreThreadDispatcher> AuthorityService<C> {
             context.clone(),
             core_dispatcher.clone(),
         ));
+        let committee_size = context.committee.size();
 
         Self {
             context,
@@ -157,6 +166,17 @@ impl<C: CoreThreadDispatcher> AuthorityService<C> {
             useful_authorities_from_peer: Arc::new(RwLock::new(BTreeMap::new())),
             useful_authorities_to_peer: Arc::new(RwLock::new(BTreeMap::new())),
             transaction_message_sender,
+            last_round_with_useful_shards_from_peer: Arc::new(RwLock::new(vec![
+                vec![
+                    GENESIS_ROUND;
+                    committee_size
+                ];
+                committee_size
+            ])),
+            peers_reporting_useful_shards: Arc::new(RwLock::new(vec![
+                BTreeSet::new();
+                committee_size
+            ])),
         }
     }
 }
@@ -181,6 +201,11 @@ impl<C: CoreThreadDispatcher> NetworkService for AuthorityService<C> {
             let mut guard = self.useful_authorities_to_peer.write();
             guard.insert(peer, useful_authorities_to_peer);
         }
+        let useful_shards_authors_to_peer = serialized_block_bundle_parts.useful_shards_authors();
+        {
+            let mut guard = self.peers_reporting_useful_shards.write();
+            guard[peer] = useful_shards_authors_to_peer;
+        }
 
         // 1. Create a verified block and make some preliminary checks
         let SerializedHeaderAndTransactions {
@@ -556,6 +581,7 @@ impl<C: CoreThreadDispatcher> NetworkService for AuthorityService<C> {
             .map_err(|_| ConsensusError::Shutdown)?;
 
         // 11. Add the block to dag, add its missing ancestors to the set
+        let block_round = verified_block.round();
         let (missing_block_ancestors, missing_block_committed_transactions) = self
             .core_dispatcher
             .add_blocks(vec![verified_block])
@@ -597,6 +623,20 @@ impl<C: CoreThreadDispatcher> NetworkService for AuthorityService<C> {
                 self.useful_authorities_from_peer.write();
             useful_authorities_from_peer_write_guard.insert(peer, useful_authorities);
         }
+        {
+            let useful_shard_authors = additional_block_headers
+                .iter()
+                .map(|block_header| block_header.author())
+                .collect::<Vec<_>>();
+            let mut last_round_with_useful_shards_from_peer_write_guard =
+                self.last_round_with_useful_shards_from_peer.write();
+            for author in useful_shard_authors {
+                last_round_with_useful_shards_from_peer_write_guard[peer][author] = max(
+                    last_round_with_useful_shards_from_peer_write_guard[peer][author],
+                    block_round,
+                );
+            }
+        }
 
         // 14. schedule the fetching of missing ancestors (if any) from this peer
         if !missing_ancestors.is_empty() {
@@ -662,6 +702,9 @@ impl<C: CoreThreadDispatcher> NetworkService for AuthorityService<C> {
             let dag_state = Arc::clone(&self.dag_state);
             let useful_authorities_from_peer = Arc::clone(&self.useful_authorities_from_peer);
             let useful_authorities_to_peer = Arc::clone(&self.useful_authorities_to_peer);
+            let last_round_with_useful_shards_from_peer =
+                Arc::clone(&self.last_round_with_useful_shards_from_peer);
+            let peers_reporting_useful_shards = Arc::clone(&self.peers_reporting_useful_shards);
             let committee = self
                 .context
                 .committee
@@ -675,6 +718,12 @@ impl<C: CoreThreadDispatcher> NetworkService for AuthorityService<C> {
                     .get(&peer)
                     .map(|authorities| authorities.iter().cloned().collect::<BTreeSet<_>>());
                 drop(useful_authorities_to_peer_guard);
+
+                let peers_reporting_useful_shards_guard = peers_reporting_useful_shards.read();
+                let peers_reporting_useful_shards_read =
+                    peers_reporting_useful_shards_guard[peer].clone();
+                drop(peers_reporting_useful_shards_guard);
+
                 let useful_authorities_from_peer_guard = useful_authorities_from_peer.read();
                 let useful_authorities_from_peer_read =
                     match useful_authorities_from_peer_guard.get(&peer) {
@@ -683,6 +732,12 @@ impl<C: CoreThreadDispatcher> NetworkService for AuthorityService<C> {
                     };
                 drop(useful_authorities_from_peer_guard);
 
+                let last_round_with_useful_shards_from_peer_guard =
+                    last_round_with_useful_shards_from_peer.read();
+                let last_round_with_useful_shards_from_peer_read =
+                    last_round_with_useful_shards_from_peer_guard[peer].clone();
+                drop(last_round_with_useful_shards_from_peer_guard);
+
                 let mut dag_state_guard = dag_state.write();
                 let block_headers = match useful_authorities_to_peer_read {
                     None => dag_state_guard.take_unknown_headers_for_authority(peer, block.round()),
@@ -692,15 +747,25 @@ impl<C: CoreThreadDispatcher> NetworkService for AuthorityService<C> {
                         authorities,
                     ),
                 };
-                let serialized_shards =
-                    dag_state_guard.take_unknown_shards_for_authority(peer, block.round());
+
+                let serialized_shards = dag_state_guard.take_useful_shards_for_authority(
+                    peer,
+                    block.round(),
+                    peers_reporting_useful_shards_read,
+                );
                 drop(dag_state_guard);
 
+                let useful_shards_authors = DagState::get_useful_shards_authors(
+                    last_round_with_useful_shards_from_peer_read,
+                    block.round(),
+                );
+
                 let block_bundle = BlockBundle {
                     verified_block: block,
                     verified_headers: block_headers,
                     serialized_shards,
                     useful_authorities: useful_authorities_from_peer_read,
+                    useful_shards_authors,
                 };
                 async move {
                     match SerializedBlockBundle::try_from(block_bundle) {
@@ -1672,6 +1737,7 @@ mod tests {
             verified_headers: headers.clone(),
             serialized_shards: vec![],
             useful_authorities: (0u8..(committee_size as u8)).map(Into::into).collect(),
+            useful_shards_authors: (0u8..(committee_size as u8)).map(Into::into).collect(),
         };
         let serialized_big_block_bundle = SerializedBlockBundle::try_from(
             SerializedBlockBundleParts::try_from(big_block_bundle).unwrap(),
@@ -1701,6 +1767,7 @@ mod tests {
             verified_headers: headers.clone(),
             serialized_shards: vec![],
             useful_authorities: (0u8..(committee_size as u8)).map(Into::into).collect(),
+            useful_shards_authors: (0u8..(committee_size as u8)).map(Into::into).collect(),
         };
         let serialized_block_bundle_with_big_round = SerializedBlockBundle::try_from(
             SerializedBlockBundleParts::try_from(block_bundle_with_big_rounds).unwrap(),
@@ -1738,6 +1805,7 @@ mod tests {
             verified_headers: headers.clone(),
             serialized_shards: vec![],
             useful_authorities: (0u8..(committee_size as u8)).map(Into::into).collect(),
+            useful_shards_authors: (0u8..(committee_size as u8)).map(Into::into).collect(),
         };
         let serialized_block_bundle = SerializedBlockBundle::try_from(
             SerializedBlockBundleParts::try_from(block_bundle).unwrap(),
@@ -2235,6 +2303,9 @@ mod tests {
                     useful_authorities: (0u8..(context.committee.size() as u8))
                         .map(Into::into)
                         .collect(),
+                    useful_shards_authors: (0u8..(context.committee.size() as u8))
+                        .map(Into::into)
+                        .collect(),
                 };
                 let serialized_block_bundle = SerializedBlockBundle::try_from(
                     SerializedBlockBundleParts::try_from(block_bundle).unwrap(),
@@ -2378,6 +2449,9 @@ mod tests {
                     useful_authorities: (0u8..(context.committee.size() as u8))
                         .map(Into::into)
                         .collect(),
+                    useful_shards_authors: (0u8..(context.committee.size() as u8))
+                        .map(Into::into)
+                        .collect(),
                 };
                 let serialized_block_bundle = SerializedBlockBundle::try_from(
                     SerializedBlockBundleParts::try_from(block_bundle).unwrap(),

crates/starfish/core/src/dag_state.rs

@@ -41,6 +41,7 @@ use crate::{
 pub(crate) const MAX_TRANSACTIONS_ACK_DEPTH: Round = 50;
 pub(crate) const MAX_HEADERS_PER_BUNDLE: usize = 150;
 pub(crate) const MAX_SHARDS_PER_BUNDLE: usize = 150;
+const MAX_ROUND_GAP_FOR_USEFUL_SHARDS: usize = 5;
 
 /// DagState provides the API to write and read accepted blocks from the DAG.
 /// Only uncommitted and last committed blocks are cached in memory.
@@ -1326,7 +1327,7 @@ impl DagState {
                 break;
             }
             if useful_authorities.contains(&block_ref.author) {
-                to_take.push(block_ref.clone());
+                to_take.push(*block_ref);
             }
         }
 
@@ -1347,6 +1348,36 @@ impl DagState {
         to_take
     }
 
+    fn take_block_refs_of_useful_shards_for_authority(
+        &mut self,
+        authority_index: AuthorityIndex,
+        round_upper_bound_exclusive: Round,
+        useful_authorities: BTreeSet<AuthorityIndex>,
+    ) -> Vec<BlockRef> {
+        let set = &mut self.shards_not_known_by_authority[authority_index.value()];
+
+        // Collect candidate block_refs we want to take out
+        let mut to_take = Vec::new();
+
+        for block_ref in set.iter() {
+            if to_take.len() >= MAX_SHARDS_PER_BUNDLE
+                || block_ref.round >= round_upper_bound_exclusive
+            {
+                break;
+            }
+            if useful_authorities.contains(&block_ref.author) {
+                to_take.push(*block_ref);
+            }
+        }
+
+        // Remove the references from the set and update cordial knowledge
+        for block_ref in &to_take {
+            set.remove(block_ref);
+        }
+
+        to_take
+    }
+
     /// Retrieves up to `MAX_HEADERS_PER_BUNDLE` previously unknown block
     /// headers for the given authority, restricted to rounds strictly below
     /// `round_upper_bound_exclusive`.
@@ -1395,6 +1426,28 @@ impl DagState {
             .collect()
     }
 
+    pub(crate) fn take_useful_shards_for_authority(
+        &mut self,
+        authority_index: AuthorityIndex,
+        round_upper_bound_exclusive: Round,
+        useful_authorities: BTreeSet<AuthorityIndex>,
+    ) -> Vec<Bytes> {
+        let block_refs = self.take_block_refs_of_useful_shards_for_authority(
+            authority_index,
+            round_upper_bound_exclusive,
+            useful_authorities,
+        );
+        block_refs
+            .iter()
+            .map(|block_ref| {
+                self.recent_shards
+                    .get(block_ref)
+                    .expect("Shard should be in DagState")
+                    .clone()
+            })
+            .collect::<Vec<_>>()
+    }
+
     pub(crate) fn take_unknown_shards_for_authority(
         &mut self,
         authority_index: AuthorityIndex,
@@ -1424,6 +1477,19 @@ impl DagState {
         }
         shards
     }
+
+    pub(crate) fn get_useful_shards_authors(
+        last_useful_round: Vec<Round>,
+        block_round: Round,
+    ) -> BTreeSet<AuthorityIndex> {
+        let mut useful_authorities = BTreeSet::new();
+        for (i, round) in last_useful_round.iter().enumerate() {
+            if block_round - round < MAX_ROUND_GAP_FOR_USEFUL_SHARDS as u32 {
+                useful_authorities.insert(AuthorityIndex::from(i as u8));
+            }
+        }
+        useful_authorities
+    }
     pub(crate) fn take_commit_votes(&mut self, limit: usize) -> Vec<CommitVote> {
         let mut votes = Vec::new();
         while !self.pending_commit_votes.is_empty() && votes.len() < limit {

crates/starfish/core/src/network/mod.rs

@@ -247,6 +247,7 @@ pub(crate) struct BlockBundle {
     pub(crate) verified_headers: Vec<VerifiedBlockHeader>,
     pub(crate) serialized_shards: Vec<Bytes>,
     pub(crate) useful_authorities: BTreeSet<AuthorityIndex>,
+    pub(crate) useful_shards_authors: BTreeSet<AuthorityIndex>,
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
@@ -255,21 +256,39 @@ pub(crate) struct SerializedBlockBundleParts {
     pub(crate) serialized_headers: Vec<Bytes>,
     pub(crate) serialized_shards: Vec<Bytes>,
     pub(crate) useful_authorities_bitmask: [u64; 4],
+    pub(crate) useful_shards_authors_bitmask: [u64; 4],
 }
 
+fn authority_set_to_bitmask(authorities: &BTreeSet<AuthorityIndex>) -> [u64; 4] {
+    let mut bitmask = [0u64; 4];
+    for authority_index in authorities {
+        let index = authority_index.value();
+        let array_index = index / 64;
+        let bit_pos = index % 64;
+        bitmask[array_index] |= 1u64 << bit_pos;
+    }
+    bitmask
+}
+
+fn bitmask_to_authority_set(bitmask: [u64; 4]) -> BTreeSet<AuthorityIndex> {
+    let mut set = BTreeSet::new();
+    for array_index in 0..4 {
+        let mut bits = bitmask[array_index];
+        let base = array_index * 64;
+        while bits != 0 {
+            let bit = bits.trailing_zeros() as usize;
+            set.insert(AuthorityIndex::from((base + bit) as u8));
+            bits &= bits - 1;
+        }
+    }
+    set
+}
 impl SerializedBlockBundleParts {
     pub(crate) fn useful_authorities(&self) -> BTreeSet<AuthorityIndex> {
-        let mut useful_authorities = BTreeSet::new();
-        for array_index in 0..4 {
-            let mut bits = self.useful_authorities_bitmask[array_index];
-            let base = array_index * 64;
-            while bits != 0 {
-                let bit = bits.trailing_zeros() as usize;
-                useful_authorities.insert(AuthorityIndex::from((base + bit) as u8));
-                bits &= bits - 1;
-            }
-        }
-        useful_authorities
+        bitmask_to_authority_set(self.useful_authorities_bitmask)
+    }
+    pub(crate) fn useful_shards_authors(&self) -> BTreeSet<AuthorityIndex> {
+        bitmask_to_authority_set(self.useful_shards_authors_bitmask)
     }
 }
 
@@ -293,6 +312,7 @@ impl TryFrom<VerifiedBlock> for SerializedBlockBundleParts {
             serialized_headers: vec![],
             serialized_shards: vec![],
             useful_authorities_bitmask: [0u64; 4],
+            useful_shards_authors_bitmask: [0u64; 4],
         })
     }
 }
@@ -312,18 +332,14 @@ impl TryFrom<BlockBundle> for SerializedBlockBundleParts {
         for block_header in block_bundle.verified_headers.iter() {
             serialized_block_headers.push(block_header.serialized().clone());
         }
-        let mut useful_authorities_bitmask = [0u64; 4];
-        for authority_index in block_bundle.useful_authorities {
-            let index = authority_index.value();
-            let array_index = index / 64;
-            let bit_pos = index % 64;
-            useful_authorities_bitmask[array_index] |= 1u64 << bit_pos;
-        }
         Ok(Self {
             serialized_block: Bytes::from(bytes),
             serialized_headers: serialized_block_headers,
             serialized_shards: block_bundle.serialized_shards,
-            useful_authorities_bitmask,
+            useful_authorities_bitmask: authority_set_to_bitmask(&block_bundle.useful_authorities),
+            useful_shards_authors_bitmask: authority_set_to_bitmask(
+                &block_bundle.useful_shards_authors,
+            ),
         })
     }
 }
@@ -389,6 +405,7 @@ mod tests {
             verified_headers: vec![],
             serialized_shards: vec![],
             useful_authorities: useful_authorities.clone(),
+            useful_shards_authors: useful_authorities.clone(),
         };
         let serialized_bundle = SerializedBlockBundle::try_from(block_bundle).unwrap();
         let serialized_bundle_parts =