1  // Copyright (c) 2025-2026 ACDC Network
  2  // This file is part of the alphaos library.
  3  //
  4  // Alpha Chain | Delta Chain Protocol
  5  // International Monetary Graphite.
  6  //
  7  // Derived from Aleo (https://aleo.org) and ProvableHQ (https://provable.com).
  8  // They built world-class ZK infrastructure. We installed the EASY button.
  9  // Their cryptography: elegant. Our modifications: bureaucracy-compatible.
 10  // Original brilliance: theirs. Robert's Rules: ours. Bugs: definitely ours.
 11  //
 12  // Original Aleo/ProvableHQ code subject to Apache 2.0 https://www.apache.org/licenses/LICENSE-2.0
 13  // All modifications and new work: CC0 1.0 Universal Public Domain Dedication.
 14  // No rights reserved. No permission required. No warranty. No refunds.
 15  //
 16  // https://creativecommons.org/publicdomain/zero/1.0/
 17  // SPDX-License-Identifier: CC0-1.0
 18  
 19  mod common;
 20  use common::*;
 21  
 22  use alphaos_node_network::PeerPoolHandling;
 23  use alphaos_node_router::{Outbound, Routing};
 24  use alphaos_node_tcp::protocols::{Disconnect, Handshake, OnConnect};
 25  use alphavm::{prelude::Rng, utilities::TestRng};
 26  use deadline::deadline;
 27  use peak_alloc::PeakAlloc;
 28  
 29  use core::time::Duration;
 30  
 31  #[global_allocator]
 32  static PEAK_ALLOC: PeakAlloc = PeakAlloc;
 33  
 34  #[tokio::test]
 35  async fn test_connection_cleanups() {
 36      // The number of connections to start and close.
 37      const NUM_CONNECTIONS: usize = 10;
 38  
 39      // Initialize an Rng.
 40      let mut rng = TestRng::default();
 41  
 42      // Create 2 routers of random types.
 43      let mut nodes = Vec::with_capacity(2);
 44      for _ in 0..2 {
 45          let node = match rng.gen_range(0..=1) {
 46              0 => client(0, 1, &mut rng).await,
 47              1 => prover(0, 1, &mut rng).await,
 48              // TODO => validator(0, 1, &[], false, &mut rng).await,
 49              _ => unreachable!(),
 50          };
 51  
 52          nodes.push(node);
 53      }
 54  
 55      // Enable handshake handling.
 56      nodes[0].enable_handshake().await;
 57      nodes[1].enable_handshake().await;
 58  
 59      nodes[0].enable_disconnect().await;
 60      nodes[1].enable_disconnect().await;
 61  
 62      nodes[0].enable_on_connect().await;
 63      nodes[1].enable_on_connect().await;
 64  
 65      nodes[0].enable_listener().await;
 66      nodes[1].enable_listener().await;
 67  
 68      // We'll want to register heap use after a single connection, after the related collections are initialized.
 69      let mut heap_after_one_conn = None;
 70  
 71      // Connect and disconnect in a loop.
 72      for i in 0..NUM_CONNECTIONS {
 73          // Connect one of the nodes to the other one.
 74          nodes[1].connect(nodes[0].local_ip());
 75  
 76          // Wait until the connection is complete.
 77          let node0 = nodes[0].clone();
 78          let node1 = nodes[1].clone();
 79          deadline!(Duration::from_secs(3), move || node0.router().number_of_connected_peers() == 1
 80              && node1.router().number_of_connected_peers() == 1);
 81  
 82          // Since the connectee doesn't read from the connector, it can't tell that the connector disconnected
 83          // from it, so it needs to disconnect from it manually.
 84          nodes[0].disconnect(nodes[1].local_ip());
 85          nodes[1].disconnect(nodes[0].local_ip());
 86  
 87          // Wait until the disconnect is complete.
 88          let node0 = nodes[0].clone();
 89          let node1 = nodes[1].clone();
 90          deadline!(Duration::from_secs(3), move || node0.router().number_of_connected_peers() == 0
 91              && node1.router().number_of_connected_peers() == 0);
 92  
 93          // Register heap use after a single connection.
 94          if i == 0 {
 95              heap_after_one_conn = Some(PEAK_ALLOC.current_usage());
 96          }
 97      }
 98  
 99      // Register final heap use.
100      let heap_after_loop = PEAK_ALLOC.current_usage();
101  
102      // Final heap use should equal that after the first connection.
103      assert_eq!(heap_after_one_conn.unwrap(), heap_after_loop);
104  }