feat: Add re-registration on expiry for p2p node (#685)

docs/src/storage-provider-cli/server.md

@@ -113,6 +113,10 @@ Rendezvous point address that the registration node connects to or the bootstrap
 
 Peer ID of the rendezvous point that the registration node connects to. Only needed if running a registration P2P node.
 
+### `--registration-ttl`
+
+The TTL of the p2p registration in seconds. After the node registration expires, the server automatically re-registers itself.
+
 ### `--config`
 
 Takes in a path to a configuration file, it supports both JSON and TOML (files _must_ have the right extension).
@@ -133,6 +137,7 @@ The supported configuration parameters are:
 | `p2p_key`                  | NA                             |
 | `rendezvous_point_address` | NA                             |
 | `rendezvous_point`         | `None`                         |
+| `registration_ttl`         | `24 hours`                     |
 
 #### Bare bones configuration
 

storage-provider/server/src/config.rs

@@ -15,6 +15,8 @@ use crate::{
     DEFAULT_NODE_ADDRESS,
 };
 
+pub const DEFAULT_REGISTRATION_TTL: u64 = 86400;
+
 /// Default address to bind the RPC server to.
 const fn default_rpc_listen_address() -> SocketAddr {
     SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8000)
@@ -31,6 +33,11 @@ const fn default_parallel_prove_commits() -> NonZero<usize> {
     unsafe { NonZero::new_unchecked(2) }
 }
 
+/// Default registration TTL, how long the node is registered.
+const fn default_registration_ttl() -> u64 {
+    DEFAULT_REGISTRATION_TTL
+}
+
 fn default_node_address() -> Url {
     Url::parse(DEFAULT_NODE_ADDRESS).expect("DEFAULT_NODE_ADDRESS must be a valid Url")
 }
@@ -123,4 +130,9 @@ pub struct ConfigurationArgs {
     #[serde(default, deserialize_with = "string_to_peer_id_option")]
     #[arg(long)]
     pub(crate) rendezvous_point: Option<PeerId>,
+
+    /// TTL of the p2p registration in seconds
+    #[serde(default = "default_registration_ttl")]
+    #[arg(long, default_value_t = DEFAULT_REGISTRATION_TTL)]
+    pub(crate) registration_ttl: u64,
 }

storage-provider/server/src/main.rs

@@ -273,6 +273,9 @@ pub struct Server {
     /// PeerID of the bootstrap node used by the registration node.
     /// Optional because it is not used by the bootstrap node.
     rendezvous_point: Option<PeerId>,
+
+    /// TTL of the p2p registration in seconds
+    registration_ttl: u64,
 }
 
 impl TryFrom<ServerCli> for Server {
@@ -354,6 +357,7 @@ impl TryFrom<ServerCli> for Server {
             p2p_key: args.p2p_key,
             rendezvous_point_address: args.rendezvous_point_address,
             rendezvous_point: args.rendezvous_point,
+            registration_ttl: args.registration_ttl,
         })
     }
 }
@@ -482,6 +486,7 @@ impl Server {
             p2p_key: self.p2p_key,
             rendezvous_point_address: self.rendezvous_point_address,
             rendezvous_point: self.rendezvous_point,
+            registration_ttl: self.registration_ttl,
         };
 
         Ok(SetupOutput {
@@ -569,6 +574,7 @@ fn spawn_p2p_task(
                 p2p_state.p2p_key,
                 p2p_state.rendezvous_point_address,
                 rendezvous_point,
+                p2p_state.registration_ttl,
             );
             Ok(tokio::spawn(run_register_node(config, cancellation_token)))
         }

storage-provider/server/src/p2p/bootstrap.rs

@@ -10,6 +10,7 @@ use libp2p::{
 };
 
 use super::P2PError;
+use crate::config::DEFAULT_REGISTRATION_TTL;
 
 #[derive(NetworkBehaviour)]
 pub struct BootstrapBehaviour {
@@ -39,7 +40,7 @@ impl BootstrapConfig {
             .with_behaviour(|key| BootstrapBehaviour {
                 // Rendezvous server behaviour for serving new peers to connecting nodes.
                 rendezvous: rendezvous::server::Behaviour::new(
-                    rendezvous::server::Config::default(),
+                    rendezvous::server::Config::default().with_max_ttl(DEFAULT_REGISTRATION_TTL), // Max TTL of 24 hours
                 ),
                 // The identify behaviour is used to share the external address and the public key with connecting clients.
                 identify: identify::Behaviour::new(identify::Config::new(
@@ -65,11 +66,8 @@ pub(crate) async fn bootstrap(
     swarm.listen_on(addr)?;
     while let Some(event) = swarm.next().await {
         match event {
-            SwarmEvent::ConnectionEstablished { peer_id, .. } => {
-                tracing::info!("Connected to {}", peer_id);
-            }
-            SwarmEvent::ConnectionClosed { peer_id, .. } => {
-                tracing::info!("Disconnected from {}", peer_id);
+            SwarmEvent::NewListenAddr { address, .. } => {
+                tracing::info!("Listening on {}", address);
             }
             SwarmEvent::Behaviour(BootstrapBehaviourEvent::Rendezvous(
                 rendezvous::server::Event::PeerRegistered { peer, registration },
@@ -95,7 +93,16 @@ pub(crate) async fn bootstrap(
                     );
                 }
             }
-            _other => {}
+            SwarmEvent::Behaviour(BootstrapBehaviourEvent::Rendezvous(
+                rendezvous::server::Event::RegistrationExpired(registration),
+            )) => {
+                tracing::info!(
+                    "Registration for peer {} expired in namespace {}",
+                    registration.record.peer_id(),
+                    registration.namespace
+                );
+            }
+            other => tracing::debug!("Encountered event: {other:?}"),
         }
     }
     Ok(())

storage-provider/server/src/p2p/mod.rs

@@ -6,7 +6,7 @@ use ed25519_dalek::{pkcs8::DecodePrivateKey, SigningKey};
 use libp2p::{identity::Keypair, rendezvous::Namespace, Multiaddr, PeerId};
 use register::register;
 use serde::{de, Deserialize};
-use tokio_util::{sync::CancellationToken, task::TaskTracker};
+use tokio_util::sync::CancellationToken;
 
 mod bootstrap;
 mod register;
@@ -17,6 +17,7 @@ pub(crate) use register::RegisterConfig;
 const P2P_NAMESPACE: &str = "polka-storage";
 
 #[derive(Default, Debug, Clone, Copy, ValueEnum, Deserialize)]
+#[serde(rename_all = "lowercase")]
 pub enum NodeType {
     #[default]
     Bootstrap,
@@ -67,6 +68,9 @@ pub(crate) struct P2PState {
     /// PeerID of the bootstrap node used by the registration node.
     /// Optional because it is not used by the bootstrap node.
     pub(crate) rendezvous_point: Option<PeerId>,
+
+    /// TTL of the p2p registration in seconds
+    pub(crate) registration_ttl: u64,
 }
 
 /// Deserializes a ED25519 private key into a Keypair.
@@ -112,7 +116,6 @@ pub async fn run_bootstrap_node(
     token: CancellationToken,
 ) -> Result<(), P2PError> {
     tracing::info!("Starting P2P bootstrap node");
-    let tracker = TaskTracker::new();
     let (swarm, addr) = config.create_swarm()?;
 
     tokio::select! {
@@ -127,9 +130,6 @@ pub async fn run_bootstrap_node(
         },
     }
 
-    tracker.close();
-    tracker.wait().await;
-
     Ok(())
 }
 
@@ -140,15 +140,17 @@ pub async fn run_register_node(
     token: CancellationToken,
 ) -> Result<(), P2PError> {
     tracing::info!("Starting P2P register node");
-    let tracker = TaskTracker::new();
-    let (swarm, rendezvous_point_address, rendezvous_point) = config.create_swarm()?;
+    let rendezvous_point = config.rendezvous_point;
+    let rendezvous_point_address = config.rendezvous_point_address.clone();
+    let registration_ttl = config.registration_ttl;
+    let mut swarm = config.create_swarm()?;
 
     tokio::select! {
         res = register(
-            swarm,
+            &mut swarm,
             rendezvous_point,
             rendezvous_point_address,
-            None,
+            registration_ttl,
             Namespace::from_static(P2P_NAMESPACE),
         ) => {
             if let Err(e) = res {
@@ -161,8 +163,5 @@ pub async fn run_register_node(
         },
     }
 
-    tracker.close();
-    tracker.wait().await;
-
     Ok(())
 }