From b4dd98b3c6bd80f285585038f300c5914c7757e8 Mon Sep 17 00:00:00 2001
From: Pawan Dhananjay <pawandhananjay@gmail.com>
Date: Mon, 30 Aug 2021 13:46:13 +0000
Subject: [PATCH] Shutdown after sync (#2519)

## Issue Addressed

Resolves #2033

## Proposed Changes

Adds a flag to enable shutting down beacon node right after sync is completed.

## Additional Info

Will need modification after weak subjectivity sync is enabled to change definition of a fully synced node.
---
 beacon_node/eth2_libp2p/src/config.rs |  4 ++++
 beacon_node/network/src/service.rs    | 15 +++++++++++++++
 beacon_node/src/cli.rs                |  6 ++++++
 beacon_node/src/config.rs             |  4 ++++
 lighthouse/tests/beacon_node.rs       | 13 +++++++++++++
 5 files changed, 42 insertions(+)

diff --git a/beacon_node/eth2_libp2p/src/config.rs b/beacon_node/eth2_libp2p/src/config.rs
index e36c34b53..4c6524084 100644
--- a/beacon_node/eth2_libp2p/src/config.rs
+++ b/beacon_node/eth2_libp2p/src/config.rs
@@ -109,6 +109,9 @@ pub struct Config {
     /// prevents sending client identifying information over identify.
     pub private: bool,
 
+    /// Shutdown beacon node after sync is completed.
+    pub shutdown_after_sync: bool,
+
     /// List of extra topics to initially subscribe to as strings.
     pub topics: Vec<GossipKind>,
 }
@@ -183,6 +186,7 @@ impl Default for Config {
             private: false,
             subscribe_all_subnets: false,
             import_all_attestations: false,
+            shutdown_after_sync: false,
             topics: Vec::new(),
         }
     }
diff --git a/beacon_node/network/src/service.rs b/beacon_node/network/src/service.rs
index 26c1e272f..6dd0b58d1 100644
--- a/beacon_node/network/src/service.rs
+++ b/beacon_node/network/src/service.rs
@@ -133,6 +133,8 @@ pub struct NetworkService<T: BeaconChainTypes> {
     next_unsubscribe: Pin<Box<OptionFuture<Sleep>>>,
     /// Subscribe to all the subnets once synced.
     subscribe_all_subnets: bool,
+    /// Shutdown beacon node after sync is complete.
+    shutdown_after_sync: bool,
     /// A timer for updating various network metrics.
     metrics_update: tokio::time::Interval,
     /// gossipsub_parameter_update timer
@@ -254,6 +256,7 @@ impl<T: BeaconChainTypes> NetworkService<T> {
             next_fork_update,
             next_unsubscribe,
             subscribe_all_subnets: config.subscribe_all_subnets,
+            shutdown_after_sync: config.shutdown_after_sync,
             metrics_update,
             gossipsub_parameter_update,
             fork_context,
@@ -436,6 +439,18 @@ fn spawn_service<T: BeaconChainTypes>(
                                 }
                         }
                         NetworkMessage::SubscribeCoreTopics => {
+                            if service.shutdown_after_sync {
+                                let _ = shutdown_sender
+                                .send(ShutdownReason::Success(
+                                    "Beacon node completed sync. Shutting down as --shutdown-after-sync flag is enabled"))
+                                .await
+                                .map_err(|e| warn!(
+                                    service.log,
+                                    "failed to send a shutdown signal";
+                                    "error" => %e
+                                ));
+                                return;
+                            }
                             let mut subscribed_topics: Vec<GossipTopic> = vec![];
                             for topic_kind in eth2_libp2p::types::CORE_TOPICS.iter() {
                                 for fork_digest in service.required_gossip_fork_digests() {
diff --git a/beacon_node/src/cli.rs b/beacon_node/src/cli.rs
index 712f0eae7..b2c3b5bf2 100644
--- a/beacon_node/src/cli.rs
+++ b/beacon_node/src/cli.rs
@@ -51,6 +51,12 @@ pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
                 .help("Disables the discovery packet filter. Useful for testing in smaller networks")
                 .takes_value(false),
         )
+        .arg(
+            Arg::with_name("shutdown-after-sync")
+                .long("shutdown-after-sync")
+                .help("Shutdown beacon node as soon as sync is completed")
+                .takes_value(false),
+        )
         .arg(
             Arg::with_name("zero-ports")
                 .long("zero-ports")
diff --git a/beacon_node/src/config.rs b/beacon_node/src/config.rs
index 628ecf316..9eb5db395 100644
--- a/beacon_node/src/config.rs
+++ b/beacon_node/src/config.rs
@@ -467,6 +467,10 @@ pub fn set_network_config(
         config.import_all_attestations = true;
     }
 
+    if cli_args.is_present("shutdown-after-sync") {
+        config.shutdown_after_sync = true;
+    }
+
     if let Some(listen_address_str) = cli_args.value_of("listen-address") {
         let listen_address = listen_address_str
             .parse()
diff --git a/lighthouse/tests/beacon_node.rs b/lighthouse/tests/beacon_node.rs
index b9ab6ce99..af7f847c8 100644
--- a/lighthouse/tests/beacon_node.rs
+++ b/lighthouse/tests/beacon_node.rs
@@ -308,6 +308,19 @@ fn network_import_all_attestations_flag() {
         .with_config(|config| assert!(config.network.import_all_attestations));
 }
 #[test]
+fn network_shutdown_after_sync_flag() {
+    CommandLineTest::new()
+        .flag("shutdown-after-sync", None)
+        .run()
+        .with_config(|config| assert!(config.network.shutdown_after_sync));
+}
+#[test]
+fn network_shutdown_after_sync_disabled_flag() {
+    CommandLineTest::new()
+        .run()
+        .with_config(|config| assert!(!config.network.shutdown_after_sync));
+}
+#[test]
 fn network_listen_address_flag() {
     let addr = "127.0.0.2".parse::<Ipv4Addr>().unwrap();
     CommandLineTest::new()