feat: implement proximity-based update forwarding (#2228)

Co-authored-by: Claude <noreply@anthropic.com>

Author: sanity

crates/core/src/message.rs

@@ -16,7 +16,9 @@ use crate::{
     },
     ring::{Location, PeerKeyLocation},
 };
-use freenet_stdlib::prelude::{ContractContainer, ContractKey, DelegateKey, WrappedState};
+use freenet_stdlib::prelude::{
+    ContractContainer, ContractInstanceId, ContractKey, DelegateKey, WrappedState,
+};
 pub(crate) use sealed_msg_type::{TransactionType, TransactionTypeId};
 use serde::{Deserialize, Serialize};
 use ulid::Ulid;
@@ -333,6 +335,30 @@ pub(crate) enum NetMessageV1 {
     },
     Update(UpdateMsg),
     Aborted(Transaction),
+    /// Proximity cache protocol message for tracking which neighbors cache which contracts.
+    ProximityCache {
+        message: ProximityCacheMessage,
+    },
+}
+
+/// Messages for the proximity cache protocol.
+///
+/// This protocol allows neighbors to inform each other which contracts they have cached,
+/// enabling UPDATE forwarding to seeders who may not be explicitly subscribed.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[allow(clippy::enum_variant_names)]
+pub enum ProximityCacheMessage {
+    /// Announce changes to our cached contracts.
+    CacheAnnounce {
+        /// Contracts we've started caching.
+        added: Vec<ContractInstanceId>,
+        /// Contracts we've stopped caching.
+        removed: Vec<ContractInstanceId>,
+    },
+    /// Request a neighbor's full cache state (used on new connections).
+    CacheStateRequest,
+    /// Response with the neighbor's full cache state.
+    CacheStateResponse { contracts: Vec<ContractInstanceId> },
 }
 
 trait Versioned {
@@ -357,6 +383,7 @@ impl Versioned for NetMessageV1 {
             NetMessageV1::Unsubscribed { .. } => semver::Version::new(1, 0, 0),
             NetMessageV1::Update(_) => semver::Version::new(1, 0, 0),
             NetMessageV1::Aborted(_) => semver::Version::new(1, 0, 0),
+            NetMessageV1::ProximityCache { .. } => semver::Version::new(1, 0, 0),
         }
     }
 }
@@ -427,6 +454,10 @@ pub(crate) enum NodeEvent {
     ExpectPeerConnection {
         addr: SocketAddr,
     },
+    /// Broadcast a proximity cache message to all connected peers.
+    BroadcastProximityCache {
+        message: ProximityCacheMessage,
+    },
 }
 
 #[derive(Debug, Clone)]
@@ -506,6 +537,9 @@ impl Display for NodeEvent {
             NodeEvent::ExpectPeerConnection { addr } => {
                 write!(f, "ExpectPeerConnection (from {addr})")
             }
+            NodeEvent::BroadcastProximityCache { message } => {
+                write!(f, "BroadcastProximityCache ({message:?})")
+            }
         }
     }
 }
@@ -540,6 +574,7 @@ impl MessageStats for NetMessageV1 {
             NetMessageV1::Update(op) => op.id(),
             NetMessageV1::Aborted(tx) => tx,
             NetMessageV1::Unsubscribed { transaction, .. } => transaction,
+            NetMessageV1::ProximityCache { .. } => Transaction::NULL,
         }
     }
 
@@ -552,6 +587,7 @@ impl MessageStats for NetMessageV1 {
             NetMessageV1::Update(op) => op.target().as_ref().map(|b| b.borrow().clone()),
             NetMessageV1::Aborted(_) => None,
             NetMessageV1::Unsubscribed { .. } => None,
+            NetMessageV1::ProximityCache { .. } => None,
         }
     }
 
@@ -564,6 +600,7 @@ impl MessageStats for NetMessageV1 {
             NetMessageV1::Update(op) => op.requested_location(),
             NetMessageV1::Aborted(_) => None,
             NetMessageV1::Unsubscribed { .. } => None,
+            NetMessageV1::ProximityCache { .. } => None,
         }
     }
 }
@@ -583,6 +620,9 @@ impl Display for NetMessage {
                 Unsubscribed { key, from, .. } => {
                     write!(f, "Unsubscribed {{  key: {key}, from: {from} }}")?;
                 }
+                ProximityCache { message } => {
+                    write!(f, "ProximityCache {{ {message:?} }}")?;
+                }
             },
         };
         write!(f, "}}")

crates/core/src/node/mod.rs

@@ -68,6 +68,7 @@ mod message_processor;
 mod network_bridge;
 mod op_state_manager;
 mod p2p_impl;
+pub(crate) mod proximity_cache;
 mod request_router;
 pub(crate) mod testing_impl;
 
@@ -864,6 +865,13 @@ async fn process_message_v1<CB>(
                 op_manager.ring.remove_subscriber(key, &peer_id);
                 break;
             }
+            NetMessageV1::ProximityCache { .. } => {
+                // Legacy path doesn't have source_addr - ProximityCache requires connection-based routing
+                tracing::warn!(
+                    "ProximityCache message received via legacy path (no source address)"
+                );
+                break;
+            }
             _ => break, // Exit the loop if no applicable message type is found
         }
     }
@@ -1105,6 +1113,30 @@ where
                 op_manager.ring.remove_subscriber(key, &peer_id);
                 break;
             }
+            NetMessageV1::ProximityCache { ref message } => {
+                let Some(source) = source_addr else {
+                    tracing::warn!(
+                        "Received ProximityCache message without source address (pure network)"
+                    );
+                    break;
+                };
+                tracing::debug!(
+                    from = %source,
+                    "Processing ProximityCache message (pure network)"
+                );
+                if let Some(response) = op_manager
+                    .proximity_cache
+                    .handle_message(source, message.clone())
+                {
+                    // Send response back to sender
+                    let response_msg =
+                        NetMessage::V1(NetMessageV1::ProximityCache { message: response });
+                    if let Err(err) = conn_manager.send(source, response_msg).await {
+                        tracing::error!(%err, %source, "Failed to send ProximityCache response");
+                    }
+                }
+                break;
+            }
             _ => break, // Exit the loop if no applicable message type is found
         }
     }

crates/core/src/node/network_bridge/p2p_protoc.rs

@@ -799,6 +799,12 @@ impl P2pConnManager {
                                             peer.pub_key().clone(),
                                         ))
                                         .await;
+
+                                    // Clean up proximity cache for disconnected peer
+                                    ctx.bridge
+                                        .op_manager
+                                        .proximity_cache
+                                        .on_peer_disconnected(&peer_addr);
                                     if let Some(conn) = ctx.connections.remove(&peer_addr) {
                                         // Also remove from reverse lookup
                                         if let Some(pub_key) = pub_key_to_remove {
@@ -1271,6 +1277,27 @@ impl P2pConnManager {
                                     }
                                 }
                             }
+                            NodeEvent::BroadcastProximityCache { message } => {
+                                // Broadcast ProximityCache message to all connected peers
+                                tracing::debug!(
+                                    ?message,
+                                    peer_count = ctx.connections.len(),
+                                    "Broadcasting ProximityCache message to connected peers"
+                                );
+
+                                let msg = crate::message::NetMessage::V1(
+                                    crate::message::NetMessageV1::ProximityCache {
+                                        message: message.clone(),
+                                    },
+                                );
+
+                                for peer_addr in ctx.connections.keys() {
+                                    tracing::debug!(%peer_addr, "Sending ProximityCache to peer");
+                                    if let Err(e) = ctx.bridge.send(*peer_addr, msg.clone()).await {
+                                        tracing::warn!(%peer_addr, "Failed to send ProximityCache: {e}");
+                                    }
+                                }
+                            }
                             NodeEvent::Disconnect { cause } => {
                                 tracing::info!(
                                     "Disconnecting from network{}",

crates/core/src/node/op_state_manager.rs

@@ -32,7 +32,10 @@ use crate::{
     ring::{ConnectionManager, LiveTransactionTracker, PeerKeyLocation, Ring},
 };
 
-use super::{network_bridge::EventLoopNotificationsSender, NetEventRegister, NodeConfig};
+use super::{
+    network_bridge::EventLoopNotificationsSender, proximity_cache::ProximityCacheManager,
+    NetEventRegister, NodeConfig,
+};
 
 #[cfg(debug_assertions)]
 macro_rules! check_id_op {
@@ -213,6 +216,8 @@ pub(crate) struct OpManager {
     pub is_gateway: bool,
     /// Sub-operation tracking for atomic operation execution
     sub_op_tracker: SubOperationTracker,
+    /// Proximity cache manager for tracking neighbor contract caches
+    pub proximity_cache: Arc<ProximityCacheManager>,
 }
 
 impl Clone for OpManager {
@@ -228,6 +233,7 @@ impl Clone for OpManager {
             peer_ready: self.peer_ready.clone(),
             is_gateway: self.is_gateway,
             sub_op_tracker: self.sub_op_tracker.clone(),
+            proximity_cache: self.proximity_cache.clone(),
         }
     }
 }
@@ -284,6 +290,8 @@ impl OpManager {
             tracing::debug!("Regular peer node: peer_ready will be set after first handshake");
         }
 
+        let proximity_cache = Arc::new(ProximityCacheManager::new());
+
         Ok(Self {
             ring,
             ops,
@@ -295,6 +303,7 @@ impl OpManager {
             peer_ready,
             is_gateway,
             sub_op_tracker,
+            proximity_cache,
         })
     }