mirror of
https://github.com/pezkuwichain/pezkuwi-telemetry.git
synced 2026-07-16 08:25:44 +00:00
Update node locations when they come in, and get the real IP addr of nodes
This commit is contained in:
@@ -3,7 +3,7 @@ use std::sync::atomic::AtomicU64;
|
||||
use futures::channel::mpsc;
|
||||
use futures::{ future, Sink, SinkExt };
|
||||
use super::inner_loop;
|
||||
use super::find_location::{ self, find_location };
|
||||
use crate::find_location::find_location;
|
||||
use crate::state::NodeId;
|
||||
use std::net::Ipv4Addr;
|
||||
|
||||
|
||||
@@ -1,206 +0,0 @@
|
||||
use std::net::Ipv4Addr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use parking_lot::RwLock;
|
||||
use rustc_hash::FxHashMap;
|
||||
use serde::Deserialize;
|
||||
use futures::{Sink, SinkExt, StreamExt};
|
||||
use futures::channel::mpsc;
|
||||
|
||||
use common::types::NodeLocation;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
/// The returned location is optional; it may be None if not found.
|
||||
pub type Location = Option<Arc<NodeLocation>>;
|
||||
|
||||
/// This is responsible for taking an IP address and attempting
|
||||
/// to find a geographical location from this
|
||||
pub fn find_location<Id, R>(response_chan: R) -> mpsc::UnboundedSender<(Id, Ipv4Addr)>
|
||||
where
|
||||
R: Sink<(Id, Option<Arc<NodeLocation>>)> + Unpin + Send + Clone + 'static,
|
||||
Id: Clone + Send + 'static
|
||||
{
|
||||
let (tx, mut rx) = mpsc::unbounded();
|
||||
|
||||
// cache entries
|
||||
let mut cache: FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>> = FxHashMap::default();
|
||||
|
||||
// Default entry for localhost
|
||||
cache.insert(
|
||||
Ipv4Addr::new(127, 0, 0, 1),
|
||||
Some(Arc::new(NodeLocation {
|
||||
latitude: 52.516_6667,
|
||||
longitude: 13.4,
|
||||
city: "Berlin".into(),
|
||||
})),
|
||||
);
|
||||
|
||||
// Create a locator with our cache. This is used to obtain locations.
|
||||
let locator = Locator::new(cache);
|
||||
|
||||
// Spawn a loop to handle location requests
|
||||
tokio::spawn(async move {
|
||||
|
||||
// Allow 4 requests at a time. acquiring a token will block while the
|
||||
// number of concurrent location requests is more than this.
|
||||
let semaphore = Arc::new(Semaphore::new(4));
|
||||
|
||||
loop {
|
||||
while let Some((id, ip_address)) = rx.next().await {
|
||||
|
||||
let permit = semaphore.clone().acquire_owned().await.unwrap();
|
||||
let mut response_chan = response_chan.clone();
|
||||
let locator = locator.clone();
|
||||
|
||||
// Once we have acquired our permit, spawn a task to avoid
|
||||
// blocking this loop so that we can handle concurrent requests.
|
||||
tokio::spawn(async move {
|
||||
match locator.locate(ip_address).await {
|
||||
Ok(loc) => {
|
||||
let _ = response_chan.send((id,loc)).await;
|
||||
},
|
||||
Err(e) => {
|
||||
log::debug!("GET error for ip location: {:?}", e);
|
||||
}
|
||||
};
|
||||
|
||||
// ensure permit is moved into task by dropping it explicitly:
|
||||
drop(permit);
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tx
|
||||
}
|
||||
|
||||
/// This struct can be used to make location requests, given
|
||||
/// an IPV4 address.
|
||||
#[derive(Clone)]
|
||||
struct Locator {
|
||||
client: reqwest::Client,
|
||||
cache: Arc<RwLock<FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>>>>,
|
||||
}
|
||||
|
||||
impl Locator {
|
||||
pub fn new(cache: FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>>) -> Self {
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
Locator {
|
||||
client,
|
||||
cache: Arc::new(RwLock::new(cache))
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn locate(&self, ip: Ipv4Addr) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
|
||||
// Return location quickly if it's cached:
|
||||
let cached_loc = {
|
||||
let cache_reader = self.cache.read();
|
||||
cache_reader.get(&ip).map(|o| o.clone())
|
||||
};
|
||||
if let Some(loc) = cached_loc {
|
||||
return Ok(loc);
|
||||
}
|
||||
|
||||
// Look it up via the location services if not cached:
|
||||
let location = self.iplocate_ipapi_co(ip).await?;
|
||||
let location = match location {
|
||||
Some(location) => Ok(Some(location)),
|
||||
None => self.iplocate_ipinfo_io(ip).await,
|
||||
}?;
|
||||
|
||||
self.cache.write().insert(ip, location.clone());
|
||||
Ok(location)
|
||||
}
|
||||
|
||||
async fn iplocate_ipapi_co(&self, ip: Ipv4Addr) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
|
||||
let location = self
|
||||
.query(&format!("https://ipapi.co/{}/json", ip))
|
||||
.await?
|
||||
.map(Arc::new);
|
||||
|
||||
Ok(location)
|
||||
}
|
||||
|
||||
async fn iplocate_ipinfo_io(&self, ip: Ipv4Addr) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
|
||||
let location = self
|
||||
.query(&format!("https://ipinfo.io/{}/json", ip))
|
||||
.await?
|
||||
.and_then(|loc: IPApiLocate| loc.into_node_location().map(Arc::new));
|
||||
|
||||
Ok(location)
|
||||
}
|
||||
|
||||
async fn query<T>(&self, url: &str) -> Result<Option<T>, reqwest::Error>
|
||||
where for<'de> T: Deserialize<'de>
|
||||
{
|
||||
match self.client.get(url).send().await?.json::<T>().await {
|
||||
Ok(result) => Ok(Some(result)),
|
||||
Err(err) => {
|
||||
log::debug!("JSON error for ip location: {:?}", err);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// This is the format returned from ipinfo.co, so we do
|
||||
/// a little conversion to get it into the shape we want.
|
||||
#[derive(Deserialize)]
|
||||
struct IPApiLocate {
|
||||
city: Box<str>,
|
||||
loc: Box<str>,
|
||||
}
|
||||
|
||||
impl IPApiLocate {
|
||||
fn into_node_location(self) -> Option<NodeLocation> {
|
||||
let IPApiLocate { city, loc } = self;
|
||||
|
||||
let mut loc = loc.split(',').map(|n| n.parse());
|
||||
|
||||
let latitude = loc.next()?.ok()?;
|
||||
let longitude = loc.next()?.ok()?;
|
||||
|
||||
// Guarantee that the iterator has been exhausted
|
||||
if loc.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(NodeLocation {
|
||||
latitude,
|
||||
longitude,
|
||||
city,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn ipapi_locate_to_node_location() {
|
||||
let ipapi = IPApiLocate {
|
||||
loc: "12.5,56.25".into(),
|
||||
city: "Foobar".into(),
|
||||
};
|
||||
|
||||
let location = ipapi.into_node_location().unwrap();
|
||||
|
||||
assert_eq!(location.latitude, 12.5);
|
||||
assert_eq!(location.longitude, 56.25);
|
||||
assert_eq!(&*location.city, "Foobar");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ipapi_locate_to_node_location_too_many() {
|
||||
let ipapi = IPApiLocate {
|
||||
loc: "12.5,56.25,1.0".into(),
|
||||
city: "Foobar".into(),
|
||||
};
|
||||
|
||||
let location = ipapi.into_node_location();
|
||||
|
||||
assert!(location.is_none());
|
||||
}
|
||||
}
|
||||
@@ -8,13 +8,13 @@ use common::{
|
||||
util::now
|
||||
};
|
||||
use bimap::BiMap;
|
||||
use std::{iter::FromIterator, net::Ipv4Addr, str::FromStr};
|
||||
use std::{net::{IpAddr, Ipv4Addr}, str::FromStr};
|
||||
use futures::channel::{ mpsc };
|
||||
use futures::{ SinkExt, StreamExt };
|
||||
use std::collections::{ HashMap, HashSet };
|
||||
use crate::state::{ self, State, NodeId };
|
||||
use crate::feed_message::{ self, FeedMessageSerializer };
|
||||
use super::find_location;
|
||||
use crate::find_location;
|
||||
|
||||
/// A unique Id is assigned per websocket connection (or more accurately,
|
||||
/// per feed socket and per shard socket). This can be combined with the
|
||||
@@ -143,7 +143,7 @@ pub struct InnerLoop {
|
||||
/// These feeds want finality info, too.
|
||||
feed_conn_id_finality: HashSet<ConnId>,
|
||||
|
||||
/// Send messages here to make location requests, which are sent back into the loop.
|
||||
/// Send messages here to make geographical location requests.
|
||||
tx_to_locator: mpsc::UnboundedSender<(NodeId, Ipv4Addr)>
|
||||
}
|
||||
|
||||
@@ -184,8 +184,32 @@ impl InnerLoop {
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle messages that come from the node geographical locator.
|
||||
async fn handle_from_find_location(&mut self, node_id: NodeId, location: find_location::Location) {
|
||||
// TODO: Update node location here
|
||||
self.node_state.update_node_location(node_id, location.clone());
|
||||
|
||||
if let Some(loc) = location {
|
||||
let mut feed_message_serializer = FeedMessageSerializer::new();
|
||||
feed_message_serializer.push(feed_message::LocatedNode(
|
||||
node_id,
|
||||
loc.latitude,
|
||||
loc.longitude,
|
||||
&loc.city
|
||||
));
|
||||
|
||||
if let Some(bytes) = feed_message_serializer.into_finalized() {
|
||||
let chain_label = self.node_state
|
||||
.get_node_chain(node_id)
|
||||
.map(|chain| chain.label());
|
||||
|
||||
if let Some(chain_label) = chain_label {
|
||||
// Don't hold onto lifetime from self because we call a mut fn next:
|
||||
let label = chain_label.to_owned();
|
||||
// Update location for any feeds subscribed to the node's chain.
|
||||
self.broadcast_to_chain_feeds(&label, ToFeedWebsocket::Bytes(bytes)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle messages coming from shards.
|
||||
@@ -228,7 +252,10 @@ impl InnerLoop {
|
||||
).await
|
||||
}
|
||||
|
||||
// TODO: The node has been added. use it's IP to find a location.
|
||||
// Currently we only geographically locate IPV4 addresses so ignore IPV6;
|
||||
if let Some(IpAddr::V4(ip_v4)) = ip {
|
||||
let _ = self.tx_to_locator.send((node_id, ip_v4)).await;
|
||||
}
|
||||
},
|
||||
}
|
||||
},
|
||||
@@ -409,6 +436,7 @@ impl InnerLoop {
|
||||
if let Some(feeds) = self.chain_to_feed_conn_ids.get(chain) {
|
||||
for &feed_id in feeds {
|
||||
// How much faster would it be if we processed these in parallel?
|
||||
// Is it practical to do so given lifetimes and such?
|
||||
if let Some(chan) = self.feed_channels.get_mut(&feed_id) {
|
||||
chan.send(message.clone()).await;
|
||||
}
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
mod aggregator;
|
||||
mod inner_loop;
|
||||
mod find_location;
|
||||
|
||||
// Expose the various message types that can be worked with externally:
|
||||
pub use inner_loop::{ FromFeedWebsocket, FromShardWebsocket, ToFeedWebsocket, ToShardWebsocket };
|
||||
|
||||
Reference in New Issue
Block a user