Pārlūkot izejas kodu

script/research/dhtd: renaming and added optional init lookup map

aggstam 4 gadi atpakaļ
vecāks
revīzija
7b0d19e8dd

+ 27 - 27
script/research/dhtd/src/main.rs

@@ -34,7 +34,7 @@ mod error;
 use error::{server_error, RpcError};
 
 mod structures;
-use structures::{KeyRequest, KeyResponse, LookupRequest, State, StatePtr};
+use structures::{Dht, DhtPtr, KeyRequest, KeyResponse, LookupRequest};
 
 mod protocol;
 use protocol::Protocol;
@@ -81,12 +81,12 @@ struct Args {
     verbose: u8,
 }
 
-/// Struct representing DHT daemon state.
+/// Struct representing DHT daemon.
 /// This example/temp-impl stores String data.
 /// In final version everything will be in bytes (Vec<u8).
 pub struct Dhtd {
-    /// Daemon state
-    state: StatePtr,
+    /// Daemon dht state
+    dht: DhtPtr,
     /// P2P network pointer
     p2p: P2pPtr,
     /// Channel to receive responses from P2P
@@ -97,12 +97,12 @@ pub struct Dhtd {
 
 impl Dhtd {
     pub async fn new(
-        state: StatePtr,
+        dht: DhtPtr,
         p2p: P2pPtr,
         p2p_recv_channel: async_channel::Receiver<KeyResponse>,
         stop_signal: async_channel::Receiver<()>,
     ) -> Result<Self> {
-        Ok(Self { state, p2p, p2p_recv_channel, stop_signal })
+        Ok(Self { dht, p2p, p2p_recv_channel, stop_signal })
     }
 
     // RPCAPI:
@@ -117,7 +117,7 @@ impl Dhtd {
 
         // Node verifies the key exist in the lookup map.
         let key = params[0].to_string();
-        let peers = match self.state.read().await.lookup.get(&key) {
+        let peers = match self.dht.read().await.lookup.get(&key) {
             Some(v) => v.clone(),
             None => {
                 info!("Did not find key: {}", key);
@@ -130,7 +130,7 @@ impl Dhtd {
         // Each node holds a local map, acting as its cache.
         // When the node receives a request for a key it doesn't hold,
         // it will query the P2P network and saves the response in its local cache.
-        match self.state.read().await.map.get(&key) {
+        match self.dht.read().await.map.get(&key) {
             Some(v) => {
                 let string = std::str::from_utf8(&v).unwrap();
                 return JsonResponse::new(json!(string), id).into()
@@ -148,7 +148,7 @@ impl Dhtd {
         }
 
         // We create a key request, and broadcast it to the network
-        let daemon = self.state.read().await.id.to_string();
+        let daemon = self.dht.read().await.id.to_string();
         // We choose last known peer as request recipient
         let peer = peers.iter().last().unwrap().to_string();
         let request = KeyRequest::new(daemon.clone(), peer, key.clone());
@@ -216,14 +216,14 @@ impl Dhtd {
         self.insert_pair(id, key, value).await
     }
 
-    /// Auxilary function to handle pair insertion to state
+    /// Auxilary function to handle pair insertion to dht
     async fn insert_pair(&self, id: Value, key: String, value: String) -> JsonResult {
-        if let Err(e) = self.state.write().await.insert(key.clone(), value.as_bytes().to_vec()) {
+        if let Err(e) = self.dht.write().await.insert(key.clone(), value.as_bytes().to_vec()) {
             error!("Failed to insert key: {}", e);
             return server_error(RpcError::KeyInsertFail, id)
         }
 
-        let daemon = self.state.read().await.id.to_string();
+        let daemon = self.dht.read().await.id.to_string();
         let request = LookupRequest::new(daemon, key.clone(), 0);
         if let Err(e) = self.p2p.broadcast(request).await {
             error!("Failed broadcasting request: {}", e);
@@ -244,13 +244,13 @@ impl Dhtd {
 
         let key = params[0].to_string();
         // Check if key value pair existed and act accordingly
-        let result = self.state.write().await.remove(key.clone());
+        let result = self.dht.write().await.remove(key.clone());
         match result {
             Ok(option) => match option {
                 Some(k) => {
                     info!("Key removed: {}", k);
 
-                    let daemon = self.state.read().await.id.to_string();
+                    let daemon = self.dht.read().await.id.to_string();
                     let request = LookupRequest::new(daemon, key.clone(), 1);
                     if let Err(e) = self.p2p.broadcast(request).await {
                         error!("Failed broadcasting request: {}", e);
@@ -276,7 +276,7 @@ impl Dhtd {
     // --> {"jsonrpc": "2.0", "method": "map", "params": [], "id": 1}
     // <-- {"jsonrpc": "2.0", "result": "map", "id": 1}
     pub async fn map(&self, id: Value, _params: &[Value]) -> JsonResult {
-        let map = self.state.read().await.map.clone();
+        let map = self.dht.read().await.map.clone();
         JsonResponse::new(json!(map), id).into()
     }
 
@@ -285,7 +285,7 @@ impl Dhtd {
     // --> {"jsonrpc": "2.0", "method": "lookup", "params": [], "id": 1}
     // <-- {"jsonrpc": "2.0", "result": "lookup", "id": 1}
     pub async fn lookup(&self, id: Value, _params: &[Value]) -> JsonResult {
-        let lookup = self.state.read().await.lookup.clone();
+        let lookup = self.dht.read().await.lookup.clone();
         JsonResponse::new(json!(lookup), id).into()
     }
 }
@@ -312,7 +312,7 @@ impl RequestHandler for Dhtd {
 
 // Auxilary function to periodically prun seen messages, based on when they were received.
 // This helps us to prevent broadcasting loops.
-async fn prune_seen_messages(state: StatePtr) {
+async fn prune_seen_messages(dht: DhtPtr) {
     loop {
         sleep(SEEN_DURATION as u64).await;
         debug!("Pruning seen messages");
@@ -320,7 +320,7 @@ async fn prune_seen_messages(state: StatePtr) {
         let now = Utc::now().timestamp();
 
         let mut prune = vec![];
-        let map = state.read().await.seen.clone();
+        let map = dht.read().await.seen.clone();
         for (k, v) in map.iter() {
             if now - v > SEEN_DURATION {
                 prune.push(k);
@@ -332,7 +332,7 @@ async fn prune_seen_messages(state: StatePtr) {
             map.remove(i);
         }
 
-        state.write().await.seen = map;
+        dht.write().await.seen = map;
     }
 }
 
@@ -347,8 +347,8 @@ async fn realmain(args: Args, ex: Arc<Executor<'_>>) -> Result<()> {
     })
     .unwrap();
 
-    // Initialize daemon state
-    let state = State::new().await?;
+    // Initialize daemon dht
+    let dht = Dht::new(None).await?;
 
     // P2P network
     let network_settings = net::Settings {
@@ -365,21 +365,21 @@ async fn realmain(args: Args, ex: Arc<Executor<'_>>) -> Result<()> {
     let registry = p2p.protocol_registry();
 
     info!("Registering P2P protocols...");
-    let _state = state.clone();
+    let _dht = dht.clone();
     registry
         .register(net::SESSION_ALL, move |channel, p2p| {
             let sender = p2p_send_channel.clone();
-            let state = _state.clone();
-            async move { Protocol::init(channel, sender, state, p2p).await.unwrap() }
+            let dht = _dht.clone();
+            async move { Protocol::init(channel, sender, dht, p2p).await.unwrap() }
         })
         .await;
 
-    // Initialize program state
-    let dhtd = Dhtd::new(state.clone(), p2p.clone(), p2p_recv_channel, shutdown.clone()).await?;
+    // Initialize daemon
+    let dhtd = Dhtd::new(dht.clone(), p2p.clone(), p2p_recv_channel, shutdown.clone()).await?;
     let dhtd = Arc::new(dhtd);
 
     // Task to periodically clean up daemon seen messages
-    ex.spawn(prune_seen_messages(state.clone())).detach();
+    ex.spawn(prune_seen_messages(dht.clone())).detach();
 
     // JSON-RPC server
     info!("Starting JSON-RPC server");

+ 15 - 15
script/research/dhtd/src/protocol.rs

@@ -12,7 +12,7 @@ use darkfi::{
     Result,
 };
 
-use crate::structures::{KeyRequest, KeyResponse, LookupRequest, StatePtr};
+use crate::structures::{DhtPtr, KeyRequest, KeyResponse, LookupRequest};
 
 pub struct Protocol {
     channel: ChannelPtr,
@@ -21,7 +21,7 @@ pub struct Protocol {
     resp_sub: MessageSubscription<KeyResponse>,
     lookup_sub: MessageSubscription<LookupRequest>,
     jobsman: ProtocolJobsManagerPtr,
-    state: StatePtr,
+    dht: DhtPtr,
     p2p: P2pPtr,
 }
 
@@ -29,7 +29,7 @@ impl Protocol {
     pub async fn init(
         channel: ChannelPtr,
         notify_queue_sender: async_channel::Sender<KeyResponse>,
-        state: StatePtr,
+        dht: DhtPtr,
         p2p: P2pPtr,
     ) -> Result<ProtocolBasePtr> {
         debug!("Adding Protocol to the protocol registry");
@@ -49,7 +49,7 @@ impl Protocol {
             resp_sub,
             lookup_sub,
             jobsman: ProtocolJobsManager::new("Protocol", channel),
-            state,
+            dht,
             p2p,
         }))
     }
@@ -69,14 +69,14 @@ impl Protocol {
             let req_copy = (*req).clone();
             debug!("Protocol::handle_receive_request(): req: {:?}", req_copy);
 
-            if self.state.read().await.seen.contains_key(&req_copy.id) {
+            if self.dht.read().await.seen.contains_key(&req_copy.id) {
                 debug!("Protocol::handle_receive_request(): We have already seen this request.");
                 continue
             }
 
-            self.state.write().await.seen.insert(req_copy.id.clone(), Utc::now().timestamp());
+            self.dht.write().await.seen.insert(req_copy.id.clone(), Utc::now().timestamp());
 
-            let daemon = self.state.read().await.id.to_string();
+            let daemon = self.dht.read().await.id.to_string();
             if daemon != req_copy.to {
                 if let Err(e) =
                     self.p2p.broadcast_with_exclude(req_copy.clone(), &exclude_list).await
@@ -86,7 +86,7 @@ impl Protocol {
                 };
             }
 
-            match self.state.read().await.map.get(&req_copy.key) {
+            match self.dht.read().await.map.get(&req_copy.key) {
                 Some(value) => {
                     let response =
                         KeyResponse::new(daemon, req_copy.from, req_copy.key, value.clone());
@@ -119,14 +119,14 @@ impl Protocol {
             let resp_copy = (*resp).clone();
             debug!("Protocol::handle_receive_response(): resp: {:?}", resp_copy);
 
-            if self.state.read().await.seen.contains_key(&resp_copy.id) {
+            if self.dht.read().await.seen.contains_key(&resp_copy.id) {
                 debug!("Protocol::handle_receive_response(): We have already seen this response.");
                 continue
             }
 
-            self.state.write().await.seen.insert(resp_copy.id.clone(), Utc::now().timestamp());
+            self.dht.write().await.seen.insert(resp_copy.id.clone(), Utc::now().timestamp());
 
-            if self.state.read().await.id.to_string() != resp_copy.to {
+            if self.dht.read().await.id.to_string() != resp_copy.to {
                 if let Err(e) =
                     self.p2p.broadcast_with_exclude(resp_copy.clone(), &exclude_list).await
                 {
@@ -159,23 +159,23 @@ impl Protocol {
                 continue
             }
 
-            if self.state.read().await.seen.contains_key(&req_copy.id) {
+            if self.dht.read().await.seen.contains_key(&req_copy.id) {
                 debug!(
                     "Protocol::handle_receive_lookup_request(): We have already seen this request."
                 );
                 continue
             }
 
-            self.state.write().await.seen.insert(req_copy.id.clone(), Utc::now().timestamp());
+            self.dht.write().await.seen.insert(req_copy.id.clone(), Utc::now().timestamp());
 
             let result = match req_copy.req_type {
                 0 => self
-                    .state
+                    .dht
                     .write()
                     .await
                     .lookup_insert(req_copy.key.clone(), req_copy.daemon.clone()),
                 _ => self
-                    .state
+                    .dht
                     .write()
                     .await
                     .lookup_remove(req_copy.key.clone(), req_copy.daemon.clone()),

+ 11 - 8
script/research/dhtd/src/structures.rs

@@ -9,16 +9,16 @@ use darkfi::{
     Result,
 };
 
-/// Atomic pointer to DHT daemon state
-pub type StatePtr = Arc<RwLock<State>>;
+/// Atomic pointer to DHT state
+pub type DhtPtr = Arc<RwLock<Dht>>;
 
 // TODO: lookup table to be based on directly connected peers, not broadcast based
 // TODO: replace Strings with blake3 hashes
 // Using string in structures because we are at an external crate
 // and cant use blake3 serialization. To be replaced once merged with core src.
 
-/// Struct representing DHT daemon state.
-pub struct State {
+/// Struct representing DHT state.
+pub struct Dht {
     /// Daemon id
     pub id: blake3::Hash,
     /// Daemon hasmap
@@ -30,17 +30,20 @@ pub struct State {
     pub seen: FxHashMap<String, i64>,
 }
 
-impl State {
-    pub async fn new() -> Result<StatePtr> {
+impl Dht {
+    pub async fn new(initial: Option<FxHashMap<String, HashSet<String>>>) -> Result<DhtPtr> {
         // Generate a random id
         let mut rng = rand::thread_rng();
         let n: u16 = rng.gen();
         let id = blake3::hash(&serialize(&n));
         let map = FxHashMap::default();
-        let lookup = FxHashMap::default();
+        let lookup = match initial {
+            Some(l) => l,
+            None => FxHashMap::default(),
+        };
         let seen = FxHashMap::default();
 
-        let state = Arc::new(RwLock::new(State { id, map, lookup, seen }));
+        let state = Arc::new(RwLock::new(Dht { id, map, lookup, seen }));
 
         Ok(state)
     }