Add --light-client-server flag and state cache utils (#3714)

## Issue Addressed

Part of https://github.com/sigp/lighthouse/issues/3651.

## Proposed Changes

Add a flag for enabling the light client server, which should be checked before gossip/RPC traffic is processed (e.g. https://github.com/sigp/lighthouse/pull/3693, https://github.com/sigp/lighthouse/pull/3711). The flag is available at runtime from `beacon_chain.config.enable_light_client_server`.

Additionally, a new method `BeaconChain::with_mutable_state_for_block` is added which I envisage being used for computing light client updates. Unfortunately its performance will be quite poor on average because it will only run quickly with access to the tree hash cache. Each slot the tree hash cache is only available for a brief window of time between the head block being processed and the state advance at 9s in the slot. When the state advance happens the cache is moved and mutated to get ready for the next slot, which makes it no longer useful for merkle proofs related to the head block. Rather than spend more time trying to optimise this I think we should continue prototyping with this code, and I'll make sure `tree-states` is ready to ship before we enable the light client server in prod (cf. https://github.com/sigp/lighthouse/pull/3206).

## Additional Info

I also fixed a bug in the implementation of `BeaconState::compute_merkle_proof` whereby the tree hash cache was moved with `.take()` but never put back with `.restore()`.

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -997,6 +997,46 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         Ok(self.store.get_state(state_root, slot)?)
     }
 
+    /// Run a function with mutable access to a state for `block_root`.
+    ///
+    /// The primary purpose of this function is to borrow a state with its tree hash cache
+    /// from the snapshot cache *without moving it*. This means that calls to this function should
+    /// be kept to an absolute minimum, because holding the snapshot cache lock has the ability
+    /// to delay block import.
+    ///
+    /// If there is no appropriate state in the snapshot cache then one will be loaded from disk.
+    /// If no state is found on disk then `Ok(None)` will be returned.
+    ///
+    /// The 2nd parameter to the closure is a bool indicating whether the snapshot cache was used,
+    /// which can inform logging/metrics.
+    ///
+    /// NOTE: the medium-term plan is to delete this function and the snapshot cache in favour
+    /// of `tree-states`, where all caches are CoW and everything is good in the world.
+    pub fn with_mutable_state_for_block<F, V, Payload: ExecPayload<T::EthSpec>>(
+        &self,
+        block: &SignedBeaconBlock<T::EthSpec, Payload>,
+        block_root: Hash256,
+        f: F,
+    ) -> Result<Option<V>, Error>
+    where
+        F: FnOnce(&mut BeaconState<T::EthSpec>, bool) -> Result<V, Error>,
+    {
+        if let Some(state) = self
+            .snapshot_cache
+            .try_write_for(BLOCK_PROCESSING_CACHE_LOCK_TIMEOUT)
+            .ok_or(Error::SnapshotCacheLockTimeout)?
+            .borrow_unadvanced_state_mut(block_root)
+        {
+            let cache_hit = true;
+            f(state, cache_hit).map(Some)
+        } else if let Some(mut state) = self.get_state(&block.state_root(), Some(block.slot()))? {
+            let cache_hit = false;
+            f(&mut state, cache_hit).map(Some)
+        } else {
+            Ok(None)
+        }
+    }
+
     /// Return the sync committee at `slot + 1` from the canonical chain.
     ///
     /// This is useful when dealing with sync committee messages, because messages are signed

beacon_node/beacon_chain/src/chain_config.rs

@@ -47,6 +47,8 @@ pub struct ChainConfig {
     pub count_unrealized_full: CountUnrealizedFull,
     /// Optionally set timeout for calls to checkpoint sync endpoint.
     pub checkpoint_sync_url_timeout: u64,
+    /// Whether to enable the light client server protocol.
+    pub enable_light_client_server: bool,
 }
 
 impl Default for ChainConfig {
@@ -68,6 +70,7 @@ impl Default for ChainConfig {
             paranoid_block_proposal: false,
             count_unrealized_full: CountUnrealizedFull::default(),
             checkpoint_sync_url_timeout: 60,
+            enable_light_client_server: false,
         }
     }
 }

beacon_node/beacon_chain/src/snapshot_cache.rs

@@ -298,6 +298,27 @@ impl<T: EthSpec> SnapshotCache<T> {
             })
     }
 
+    /// Borrow the state corresponding to `block_root` if it exists in the cache *unadvanced*.
+    ///
+    /// Care must be taken not to mutate the state in an invalid way. This function should only
+    /// be used to mutate the *caches* of the state, for example the tree hash cache when
+    /// calculating a light client merkle proof.
+    pub fn borrow_unadvanced_state_mut(
+        &mut self,
+        block_root: Hash256,
+    ) -> Option<&mut BeaconState<T>> {
+        self.snapshots
+            .iter_mut()
+            .find(|snapshot| {
+                // If the pre-state exists then state advance has already taken the state for
+                // `block_root` and mutated its tree hash cache. Rather than re-building it while
+                // holding the snapshot cache lock (>1 second), prefer to return `None` from this
+                // function and force the caller to load it from disk.
+                snapshot.beacon_block_root == block_root && snapshot.pre_state.is_none()
+            })
+            .map(|snapshot| &mut snapshot.beacon_state)
+    }
+
     /// If there is a snapshot with `block_root`, clone it and return the clone.
     pub fn get_cloned(
         &self,

beacon_node/src/cli.rs

@@ -868,4 +868,11 @@ pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
                       Useful if you intend to run a non-validating beacon node.")
                 .takes_value(false)
         )
+        .arg(
+            Arg::with_name("light-client-server")
+                .long("light-client-server")
+                .help("Act as a full node supporting light clients on the p2p network \
+                       [experimental]")
+                .takes_value(false)
+        )
 }

beacon_node/src/config.rs

@@ -705,6 +705,9 @@ pub fn get_config<E: EthSpec>(
     client_config.chain.builder_fallback_disable_checks =
         cli_args.is_present("builder-fallback-disable-checks");
 
+    // Light client server config.
+    client_config.chain.enable_light_client_server = cli_args.is_present("light-client-server");
+
     Ok(client_config)
 }
 

consensus/types/src/beacon_state.rs

@@ -1702,12 +1702,12 @@ impl<T: EthSpec> BeaconState<T> {
         };
 
         // 2. Get all `BeaconState` leaves.
-        let cache = self.tree_hash_cache_mut().take();
-        let leaves = if let Some(mut cache) = cache {
-            cache.recalculate_tree_hash_leaves(self)?
-        } else {
-            return Err(Error::TreeHashCacheNotInitialized);
-        };
+        let mut cache = self
+            .tree_hash_cache_mut()
+            .take()
+            .ok_or(Error::TreeHashCacheNotInitialized)?;
+        let leaves = cache.recalculate_tree_hash_leaves(self)?;
+        self.tree_hash_cache_mut().restore(cache);
 
         // 3. Make deposit tree.
         // Use the depth of the `BeaconState` fields (i.e. `log2(32) = 5`).

lighthouse/tests/beacon_node.rs

@@ -1580,3 +1580,18 @@ fn sync_eth1_chain_disable_deposit_contract_sync_flag() {
         .run_with_zero_port()
         .with_config(|config| assert_eq!(config.sync_eth1_chain, false));
 }
+
+#[test]
+fn light_client_server_default() {
+    CommandLineTest::new()
+        .run_with_zero_port()
+        .with_config(|config| assert_eq!(config.chain.enable_light_client_server, false));
+}
+
+#[test]
+fn light_client_server_enabled() {
+    CommandLineTest::new()
+        .flag("light-client-server", None)
+        .run_with_zero_port()
+        .with_config(|config| assert_eq!(config.chain.enable_light_client_server, true));
+}

testing/ef_tests/check_all_files_accessed.py

@@ -39,8 +39,6 @@ excluded_paths = [
     "tests/.*/.*/ssz_static/LightClientOptimistic",
     # LightClientFinalityUpdate
     "tests/.*/.*/ssz_static/LightClientFinalityUpdate",
-    # Merkle-proof tests for light clients
-    "tests/.*/.*/merkle/single_proof",
     # Capella tests are disabled for now.
     "tests/.*/capella",
     # One of the EF researchers likes to pack the tarballs on a Mac

testing/ef_tests/src/cases/merkle_proof_validity.rs

@@ -78,6 +78,10 @@ impl<E: EthSpec> Case for MerkleProofValidity<E> {
                 )));
             }
         }
+
+        // Tree hash cache should still be initialized (not dropped).
+        assert!(state.tree_hash_cache().is_initialized());
+
         Ok(())
     }
 }