Stop connecting outbound on shutdown in addition to stopping listen

[ldk-sample] / src / main.rs

diff --git a/src/main.rs b/src/main.rs
index 07afa66436164c15b0b5af0e018fc8135549d917..49fdcb81e8080183a1bcec4b1095c605a39da15c 100644 (file)
--- a/src/main.rs
+++ b/src/main.rs
@@ -48,6 +48,7 @@ use std::io;
 use std::io::Write;
 use std::ops::Deref;
 use std::path::Path;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, SystemTime};
 
@@ -542,6 +543,8 @@ async fn start_ldk() {
 
        let peer_manager_connection_handler = peer_manager.clone();
        let listening_port = args.ldk_peer_listening_port;
+       let stop_listen_connect = Arc::new(AtomicBool::new(false));
+       let stop_listen = Arc::clone(&stop_listen_connect);
        tokio::spawn(async move {
                let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", listening_port))
                        .await
@@ -549,6 +552,9 @@ async fn start_ldk() {
                loop {
                        let peer_mgr = peer_manager_connection_handler.clone();
                        let tcp_stream = listener.accept().await.unwrap().0;
+                       if stop_listen.load(Ordering::Acquire) {
+                               return;
+                       }
                        tokio::spawn(async move {
                                lightning_net_tokio::setup_inbound(
                                        peer_mgr.clone(),
@@ -657,6 +663,7 @@ async fn start_ldk() {
        let connect_cm = Arc::clone(&channel_manager);
        let connect_pm = Arc::clone(&peer_manager);
        let peer_data_path = format!("{}/channel_peer_data", ldk_data_dir.clone());
+       let stop_connect = Arc::clone(&stop_listen_connect);
        tokio::spawn(async move {
                let mut interval = tokio::time::interval(Duration::from_secs(1));
                loop {
@@ -670,6 +677,9 @@ async fn start_ldk() {
                                                .map(|chan| chan.counterparty.node_id)
                                                .filter(|id| !peers.contains(id))
                                        {
+                                               if stop_connect.load(Ordering::Acquire) {
+                                                       return;
+                                               }
                                                for (pubkey, peer_addr) in info.iter() {
                                                        if *pubkey == node_id {
                                                                let _ = cli::do_connect_peer(
@@ -709,10 +719,11 @@ async fn start_ldk() {
 
        // Start the CLI.
        cli::poll_for_user_input(
-               invoice_payer.clone(),
-               peer_manager.clone(),
-               channel_manager.clone(),
-               keys_manager.clone(),
+               Arc::clone(&invoice_payer),
+               Arc::clone(&peer_manager),
+               Arc::clone(&channel_manager),
+               Arc::clone(&keys_manager),
+               Arc::clone(&network_graph),
                inbound_payments,
                outbound_payments,
                ldk_data_dir.clone(),
@@ -720,6 +731,11 @@ async fn start_ldk() {
        )
        .await;
 
+       // Disconnect our peers and stop accepting new connections. This ensures we don't continue
+       // updating our channel data after we've stopped the background processor.
+       stop_listen_connect.store(true, Ordering::Release);
+       peer_manager.disconnect_all_peers();
+
        // Stop the background processor.
        background_processor.stop().unwrap();
 }