Update best snapshot picking algorithm

Signed-off-by: Stanislav Frolov <stanislav@thirdhash.com>

src/snapshot/p2p_processing.cpp

@@ -16,19 +16,35 @@
 
 namespace snapshot {
 
+const finalization::FinalizationState *FindFinalizationState(const CBlockIndex *best_index) {
+  auto repo = GetComponent<finalization::StateRepository>();
+  AssertLockHeld(repo->GetLock());
+  for (const CBlockIndex *walk = best_index; walk != nullptr; walk = walk->pprev) {
+    if (auto state = repo->Find(*walk)) {
+      return state;
+    }
+  }
+  return nullptr;
+}
+
 inline const CBlockIndex *LookupFinalizedBlockIndex(const uint256 &hash) {
   const CBlockIndex *bi = LookupBlockIndex(hash);
   if (!bi) {
-    LogPrint(BCLog::SNAPSHOT, "%s: block=%s not found\n", __func__, hash.GetHex());
+    return nullptr;
+  }
+
+  if (pindexBestHeader == nullptr) {
     return nullptr;
   }
 
   {
-    auto fin_repo = GetComponent<finalization::StateRepository>();
-    LOCK(fin_repo->GetLock());
-    const finalization::FinalizationState *fin_state = fin_repo->GetTipState();
-    if (bi->nHeight <= fin_state->GetEpochCheckpointHeight(fin_state->GetLastFinalizedEpoch())) {
-      LogPrint(BCLog::SNAPSHOT, "%s: block=%s height=%d is not finalized\n", hash.GetHex(), bi->nHeight);
+    LOCK(GetComponent<finalization::StateRepository>()->GetLock());
+    const finalization::FinalizationState *fin_state = FindFinalizationState(pindexBestHeader);
+    if (fin_state == nullptr) {
+      return nullptr;
+    }
+    if (bi->nHeight > fin_state->GetEpochCheckpointHeight(fin_state->GetLastFinalizedEpoch())) {
+      LogPrint(BCLog::SNAPSHOT, "block=%s height=%d is not finalized\n", hash.GetHex(), bi->nHeight);
       return nullptr;
     }
   }
@@ -527,11 +543,6 @@ void P2PState::SetIfBestSnapshot(const SnapshotHeader &best_snapshot) {
     return;
   }
 
-  if (m_best_snapshot.IsNull()) {
-    m_best_snapshot = best_snapshot;
-    return;
-  }
-
   // if a peer has the snapshot which matches with one node downloads, mark it the best
   if (!m_downloading_snapshot.IsNull() && best_snapshot == m_downloading_snapshot) {
     m_best_snapshot = best_snapshot;
@@ -540,17 +551,27 @@ void P2PState::SetIfBestSnapshot(const SnapshotHeader &best_snapshot) {
 
   // don't switch the snapshot once it's decided to download it
   // and there are peers that can support it
-  if (m_downloading_snapshot == m_best_snapshot) {
+  if (!m_best_snapshot.IsNull() && m_downloading_snapshot == m_best_snapshot) {
     return;
   }
 
   // compare heights to find the best snapshot
   LOCK(cs_main);
+
+  const CBlockIndex *const new_bi = LookupFinalizedBlockIndex(best_snapshot.block_hash);
+
+  if (new_bi == nullptr) {
+    return;
+  }
+
+  if (m_best_snapshot.IsNull()) {
+    m_best_snapshot = best_snapshot;
+  }
+
   const CBlockIndex *const cur_bi = LookupBlockIndex(m_best_snapshot.block_hash);
   assert(cur_bi);
 
-  const CBlockIndex *const new_bi = LookupFinalizedBlockIndex(best_snapshot.block_hash);
-  if (new_bi && new_bi->nHeight > cur_bi->nHeight) {
+  if (new_bi->nHeight > cur_bi->nHeight) {
     m_best_snapshot = best_snapshot;
     return;
   }

src/test/snapshot/p2p_processing_tests.cpp

@@ -7,6 +7,7 @@
 #include <memory>
 
 #include <chainparams.h>
+#include <injector.h>
 #include <net.h>
 #include <netaddress.h>
 #include <serialize.h>
@@ -15,6 +16,7 @@
 #include <snapshot/messages.h>
 #include <snapshot/snapshot_index.h>
 #include <snapshot/state.h>
+#include <test/esperanza/finalizationstate_utils.h>
 #include <test/test_unite.h>
 #include <validation.h>
 #include <version.h>
@@ -157,6 +159,36 @@ BOOST_AUTO_TEST_CASE(process_snapshot) {
   BOOST_CHECK_EQUAL(best_snapshot.total_utxo_subsets, total);
 }
 
+class FinalizationStateCopyable : public finalization::FinalizationState {
+ public:
+  FinalizationStateCopyable &operator=(const FinalizationStateCopyable &other) {
+    m_checkpoints = other.m_checkpoints;
+    m_epoch_to_dynasty = other.m_epoch_to_dynasty;
+    m_dynasty_start_epoch = other.m_dynasty_start_epoch;
+    m_validators = other.m_validators;
+    m_dynasty_deltas = other.m_dynasty_deltas;
+    m_deposit_scale_factor = other.m_deposit_scale_factor;
+    m_total_slashed = other.m_total_slashed;
+    m_current_epoch = other.m_current_epoch;
+    m_current_dynasty = other.m_current_dynasty;
+    m_cur_dyn_deposits = other.m_cur_dyn_deposits;
+    m_prev_dyn_deposits = other.m_prev_dyn_deposits;
+    m_expected_source_epoch = other.m_expected_source_epoch;
+    m_last_finalized_epoch = other.m_last_finalized_epoch;
+    m_last_justified_epoch = other.m_last_justified_epoch;
+    m_recommended_target_hash = other.m_recommended_target_hash;
+    m_recommended_target_epoch = other.m_recommended_target_epoch;
+    m_last_voter_rescale = other.m_last_voter_rescale;
+    m_last_non_voter_rescale = other.m_last_non_voter_rescale;
+    m_reward_factor = other.m_reward_factor;
+    m_admin_state = other.m_admin_state;
+    return *this;
+  }
+  FinalizationStateCopyable &operator=(const FinalizationState &other) {
+    return (*this = static_cast<const FinalizationStateCopyable &>(other));
+  }
+};
+
 BOOST_AUTO_TEST_CASE(start_initial_snapshot_download) {
   snapshot::InitP2P(Params().GetSnapshotParams());
   snapshot::EnableISDMode();
@@ -171,6 +203,21 @@ BOOST_AUTO_TEST_CASE(start_initial_snapshot_download) {
   b1->nHeight = 1;
   b2->nHeight = 2;
 
+  pindexBestHeader = b2;
+
+  FinalizationStateSpy spy;
+  spy.SetLastFinalizedEpoch(1);
+  auto repo = GetComponent<finalization::StateRepository>();
+  {
+    // UNIT-E TODO: These dirty things mean we must factor out snapshot to the component
+    // and mock repository dependency.
+    LOCK(repo->GetLock());
+    repo->ResetToTip(*b2);
+    finalization::FinalizationState *fin_state = repo->Find(*b2);
+    assert(fin_state != nullptr);
+    *(static_cast<FinalizationStateCopyable *>(fin_state)) = spy;
+  }
+
   snapshot::SnapshotHeader best;
   best.snapshot_hash = uint256S("a2");
   best.block_hash = b2->GetBlockHash();

src/validation.cpp

@@ -3044,7 +3044,7 @@ CBlockIndex* CChainState::AddToBlockIndex(const CBlockHeader& block)
     pindexNew->nTimeMax = (pindexNew->pprev ? std::max(pindexNew->pprev->nTimeMax, pindexNew->nTime) : pindexNew->nTime);
     pindexNew->nChainWork = (pindexNew->pprev ? pindexNew->pprev->nChainWork : 0) + GetBlockProof(*pindexNew);
     pindexNew->RaiseValidity(BLOCK_VALID_TREE);
-    if (pindexBestHeader == nullptr || pindexBestHeader->nChainWork < pindexNew->nChainWork)
+    if (pindexBestHeader == nullptr || CBlockIndexWorkComparator()(pindexBestHeader, pindexNew))
         pindexBestHeader = pindexNew;
 
     setDirtyBlockIndex.insert(pindexNew);

test/functional/p2p_snapshot.py

@@ -121,9 +121,9 @@ def on_getdata(self, message):
             if i.hash in self.parent_blocks:
                 self.send_message(msg_witness_block(self.parent_blocks[i.hash]))
 
-    def update_snapshot_from(self, node):
+    def update_snapshot_from(self, node, finalized=True):
         # take the latest finalized
-        res = next(s for s in reversed(node.listsnapshots()) if s['snapshot_finalized'])
+        res = next(s for s in reversed(node.listsnapshots()) if s['snapshot_finalized'] == finalized)
         self.snapshot_header = SnapshotHeader(
             snapshot_hash=uint256_from_rev_hex(res['snapshot_hash']),
             block_hash=uint256_from_rev_hex(res['block_hash']),
@@ -462,6 +462,7 @@ def test_invalid_snapshot(self):
         3. node - the node which syncs the snapshot
         4. broken_p2p - mini node that claims has the best snapshot but it's broken
         5. valid_p2p - mini node that sends a valid snapshot
+        6. not_finalized_p2p - mini node that claims has the best snapshot but it's not finalied
         """
 
         snap_node = self.nodes[4]
@@ -492,11 +493,16 @@ def test_invalid_snapshot(self):
         broken_p2p.update_headers_and_blocks_from(snap_node)
         valid_p2p.update_headers_and_blocks_from(snap_node)
 
+        not_finalized_p2p = node.add_p2p_connection(WaitNode(), services=SERVICE_FLAGS_WITH_SNAPSHOT)
+        not_finalized_p2p.update_snapshot_from(snap_node, finalized=False)
+        not_finalized_p2p.update_headers_and_blocks_from(snap_node)
+
         network_thread_start()
 
-        # make sure that node knows about both peers
+        # make sure that node knows about all the peers
         valid_p2p.wait_for_verack()
         broken_p2p.wait_for_verack()
+        not_finalized_p2p.wait_for_verack()
 
         # node must pick the best snapshot
         wait_until(lambda: broken_p2p.snapshot_chunk1_requested, timeout=10)
@@ -520,6 +526,10 @@ def test_invalid_snapshot(self):
         valid_p2p.return_parent_block = True
         valid_p2p.on_getsnapshot(valid_p2p.last_getsnapshot_message)
 
+        # node doesn't request not finalied snapshot
+        assert_equal(not_finalized_p2p.snapshot_header_requested, True)
+        assert_equal(not_finalized_p2p.snapshot_chunk1_requested, False)
+
         # node requests parent block and finishes ISD
         wait_until(lambda: node.getblockcount() == 16, timeout=20)
         node.disconnect_p2ps()