WIP: fix decentralized sync for storage lends

docker-compose.testnet.yml

@@ -2,14 +2,15 @@
 
 services:
   snapchain:
-    image: farcasterxyz/snapchain:latest
-    pull_policy: always
-#    build: # For testing
-#      context: .
-#      dockerfile: Dockerfile
+    # image: farcasterxyz/snapchain:latest
+    # pull_policy: always
+    build:
+     context: .
+     dockerfile: Dockerfile
     init: true # Auto-reap zombie processes and forward process signals
     environment:
       RUST_BACKTRACE: "full"
+      RUST_LOG: "info"
     entrypoint:
       - "/bin/bash"
       - "-c"
@@ -40,6 +41,7 @@ services:
         [snapshot]
         endpoint_url = "https://e1f9f185c6e63471dd39f96abd3413c4.r2.cloudflarestorage.com"
         load_db_from_snapshot=true
+        bootstrap_method = "Replicate"
         EOF
         exec $0 $@ # Now run the original command
     command: [ "./snapchain", "--config-path", "config.toml" ]

proto/definitions/replication.proto

@@ -38,6 +38,7 @@ message GetShardTransactionsRequest {
 
   // If NONE, then start from the left-most leaf node under the prefix
   optional string page_token = 4;
+  optional uint64 fid = 5;
 }
 
 message GetShardTransactionsResponse {

src/bootstrap/replication/client_test.rs

@@ -723,6 +723,7 @@ mod tests {
                     trie_virtual_shard: vts,
                     height,
                     page_token: next_page_token.clone(),
+                    fid: None,
                 };
 
                 // Call the server method and handle the Result

src/bootstrap/replication/rpc_client.rs

@@ -297,6 +297,7 @@ impl RpcClientsManager {
                                 height,
                                 trie_virtual_shard: vts as u32,
                                 page_token: Some(next_page_token),
+                                fid: None,
                             },
                             &config,
                         )
@@ -359,6 +360,7 @@ impl RpcClientsManager {
             height: self.height,
             trie_virtual_shard: vts as u32,
             page_token,
+            fid: None,
         };
 
         let response =

src/bootstrap/replication/service.rs

@@ -5,11 +5,12 @@ use crate::cfg::Config;
 use crate::core::validations;
 use crate::core::validations::message::validate_message_hash;
 use crate::network::gossip;
+use crate::proto::replication_service_client::ReplicationServiceClient;
 use crate::proto::shard_trie_entry_with_message::TrieMessage;
 use crate::proto::{self, MessageType, ReplicationTriePartStatus, ShardSnapshotMetadata};
 use crate::storage::store::block_engine::BlockEngine;
 use crate::storage::store::node_local_state::LocalStateStore;
-use crate::storage::trie::merkle_trie::MerkleTrie;
+use crate::storage::trie::merkle_trie::{DecodedTrieKey, MerkleTrie};
 use crate::storage::{
     constants::RootPrefix,
     db::{PageOptions, RocksDB, RocksDbTransactionBatch},
@@ -949,10 +950,12 @@ impl ReplicatorBootstrap {
             Self::check_fid_roots(
                 &work_item.thread_engine,
                 status.shard_id,
+                status.height,
                 status.virtual_trie_shard as u8,
                 &mut txn_batch,
                 fids_to_check,
-            )?;
+            )
+            .await?;
 
             // 9. Now that the account roots match, commit to DB
             // First, add the work status to the txn_batch so it gets commited atomically with the work done
@@ -982,10 +985,12 @@ impl ReplicatorBootstrap {
                     Self::check_fid_roots(
                         &work_item.thread_engine,
                         status.shard_id,
+                        status.height,
                         status.virtual_trie_shard as u8,
                         &mut txn_batch,
                         vec![last_fid],
-                    )?;
+                    )
+                    .await?;
                 }
                 // Write to the DB that we're all done
                 status.last_response = WorkUnitResponse::Finished as u32;
@@ -1021,10 +1026,83 @@ impl ReplicatorBootstrap {
         return response;
     }
 
+    async fn debug_account_root_mismatch(
+        thread_engine: &Arc<ShardEngine>,
+        shard_id: u32,
+        height: u64,
+        virtual_trie_shard: u8,
+        fid: u64,
+    ) -> Result<(), BootstrapError> {
+        let fid_key = TrieKey::for_fid(fid);
+        let our_trie_keys_set: HashSet<Vec<u8>> = thread_engine
+            .get_stores()
+            .trie
+            .get_all_values(&merkle_trie::Context::new(), &thread_engine.db, &fid_key)?
+            .into_iter()
+            .collect();
+
+        let mut server_trie_keys = vec![];
+
+        let mut page_token = None;
+        let mut client =
+            ReplicationServiceClient::connect("https://tau.farcaster.xyz:3383".to_string()).await?;
+
+        loop {
+            let request = proto::GetShardTransactionsRequest {
+                shard_id,
+                height,
+                trie_virtual_shard: virtual_trie_shard as u32,
+                page_token,
+                fid: Some(fid),
+            };
+
+            let response = client
+                .get_shard_transactions(request)
+                .await
+                .unwrap()
+                .into_inner();
+
+            server_trie_keys.extend(
+                response
+                    .trie_messages
+                    .into_iter()
+                    .map(|trie_message| trie_message.trie_key),
+            );
+
+            page_token = response.next_page_token;
+            if page_token.is_none() {
+                break;
+            };
+        }
+
+        let server_trie_keys_set: HashSet<Vec<u8>> = server_trie_keys.into_iter().collect();
+
+        let unique_to_us: Vec<DecodedTrieKey> = our_trie_keys_set
+            .difference(&server_trie_keys_set)
+            .map(|key| TrieKey::decode(key).unwrap())
+            .collect();
+
+        for trie_key in unique_to_us {
+            info!(fid, "Trie key missing on server {:#?}", trie_key)
+        }
+
+        let unique_to_server: Vec<DecodedTrieKey> = server_trie_keys_set
+            .difference(&our_trie_keys_set)
+            .map(|key| TrieKey::decode(key).unwrap())
+            .collect();
+
+        for trie_key in unique_to_server {
+            info!(fid, "Trie key missing locally {:#?}", trie_key);
+        }
+
+        Ok(())
+    }
+
     // Go over all the FIDs that were just processed, and check that the roots match
-    fn check_fid_roots(
+    async fn check_fid_roots(
         thread_engine: &Arc<ShardEngine>,
         shard_id: u32,
+        height: u64,
         virtual_trie_shard: u8,
         txn_batch: &mut RocksDbTransactionBatch,
         fids_to_check: Vec<u64>,
@@ -1081,6 +1159,15 @@ impl ReplicatorBootstrap {
 
             let expected_root = expected_account.account_root_hash;
             if expected_root != actual_root {
+                Self::debug_account_root_mismatch(
+                    &thread_engine,
+                    shard_id,
+                    height,
+                    virtual_trie_shard,
+                    *fid,
+                )
+                .await
+                .unwrap();
                 return Err(BootstrapError::AccountRootMismatch(format!(
                     "Account root mismatch for fid  {}/{}/{}. expected {}, got {}",
                     shard_id,

src/network/replication/replication_server.rs

@@ -141,6 +141,7 @@ impl proto::replication_service_server::ReplicationService for ReplicationServer
             request.shard_id,
             request.height,
             request.trie_virtual_shard as u8,
+            request.fid,
             request.page_token.clone(),
         );
 

src/network/replication/replicator.rs

@@ -2,8 +2,8 @@ use crate::{
     core::util,
     network::replication::{error::ReplicationError, replication_stores::ReplicationStores},
     proto::{
-        self, shard_trie_entry_with_message::TrieMessage, GetShardTransactionsResponse,
-        MessageType, OnChainEventType,
+        self, shard_trie_entry_with_message::TrieMessage, FarcasterNetwork,
+        GetShardTransactionsResponse, MessageType, OnChainEventType,
     },
     storage::{
         db::{PageOptions, RocksDbTransactionBatch},
@@ -475,6 +475,7 @@ impl Replicator {
         shard_id: u32,
         height: u64,
         trie_virtual_shard: u8,
+        fid: Option<u64>,
         page_token: Option<String>,
     ) -> Result<GetShardTransactionsResponse, ReplicationError> {
         // Get the stores for this shard_id and height
@@ -494,10 +495,16 @@ impl Replicator {
         // First, collect MAX_SIZE trie elements starting at the given page_token and prefix
         let mut trie_keys = vec![];
 
+        let prefix = if let Some(fid) = fid {
+            TrieKey::for_fid(fid)
+        } else {
+            vec![trie_virtual_shard]
+        };
+
         let next_page_token = trie.get_paged_values_of_subtree(
             &merkle_trie::Context::new(),
             &stores.db,
-            &[trie_virtual_shard],
+            &prefix,
             &mut trie_keys,
             Self::MESSAGE_LIMIT,
             page_token,
@@ -516,6 +523,7 @@ impl Replicator {
                 )));
             }
 
+            // TODO(aditi): We put messages into the trie for both fids on storage lends, but it's only stored under 1 fid in the so
             let fid = decoded_key.fid;
             let onchain_message_type = decoded_key.onchain_message_type;
             let message_type = decoded_key.message_type;
@@ -743,8 +751,15 @@ impl Replicator {
         self.stores
             .close_aged_snapshots(msg.shard_id, oldest_valid_timestamp);
 
+        // Take a snapshot for testnet nodes if none exist because there aren't many read nodes running and we may have to wait a long time for the scheduled snapshot after restart.
+        let take_first_snapshot = self.stores.network() == FarcasterNetwork::Testnet
+            && self.stores.max_height_for_shard(msg.shard_id).is_none();
+
         // Check if we can take a snapshot of this block
-        if block_number > 0 && block_number % self.snapshot_options.interval != 0 {
+        if block_number > 0
+            && block_number % self.snapshot_options.interval != 0
+            && !take_first_snapshot
+        {
             return Ok(());
         }
 

src/storage/trie/merkle_trie.rs

@@ -17,6 +17,7 @@ pub const FNAME_MESSAGE_TYPE: u8 = 7;
 
 pub const TRIE_SHARD_SIZE: u32 = 256; // So it fits into 1 byte
 
+#[derive(Debug)]
 pub struct DecodedTrieKey {
     pub virtual_shard: u8,
     pub fid: u64,