use std::{
borrow::Cow,
collections::{HashMap, HashSet},
str,
sync::Arc,
};
use flume::{Receiver, Sender};
use tokio::sync::RwLock;
use zenoh::{
internal::Value,
key_expr::{KeyExpr, OwnedKeyExpr},
query::Selector,
sample::{Sample, SampleBuilder},
time::Timestamp,
Session,
};
use super::{Digest, EraType, LogEntry, Snapshotter, CONTENTS, ERA, INTERVALS, SUBINTERVALS};
pub struct Aligner {
session: Arc<Session>,
digest_key: OwnedKeyExpr,
snapshotter: Arc<Snapshotter>,
rx_digest: Receiver<(String, Digest)>,
tx_sample: Sender<Sample>,
digests_processed: RwLock<HashSet<u64>>,
}
impl Aligner {
pub async fn start_aligner(
session: Arc<Session>,
digest_key: OwnedKeyExpr,
rx_digest: Receiver<(String, Digest)>,
tx_sample: Sender<Sample>,
snapshotter: Arc<Snapshotter>,
) {
let aligner = Aligner {
session,
digest_key,
snapshotter,
rx_digest,
tx_sample,
digests_processed: RwLock::new(HashSet::new()),
};
aligner.start().await;
}
pub async fn start(&self) {
while let Ok((from, incoming_digest)) = self.rx_digest.recv_async().await {
if self.in_processed(incoming_digest.checksum).await {
tracing::trace!(
"[ALIGNER]Skipping already processed digest: {}",
incoming_digest.checksum
);
continue;
} else if self.snapshotter.get_digest().await.checksum == incoming_digest.checksum {
tracing::trace!(
"[ALIGNER]Skipping matching digest: {}",
incoming_digest.checksum
);
continue;
} else {
tracing::debug!(
"[ALIGNER]Processing digest: {:?} from {}",
incoming_digest,
from
);
self.process_incoming_digest(incoming_digest, &from).await;
}
}
}
async fn in_processed(&self, checksum: u64) -> bool {
let processed_set = self.digests_processed.read().await;
processed_set.contains(&checksum)
}
async fn process_incoming_digest(&self, other: Digest, from: &str) {
let checksum = other.checksum;
let timestamp = other.timestamp;
let (missing_content, no_content_err) = self.get_missing_content(&other, from).await;
tracing::debug!(
"[ALIGNER] Missing {} entries; query corresponding samples",
missing_content.len()
);
if !missing_content.is_empty() {
let (missing_data, no_data_err) = self
.get_missing_data(&missing_content, timestamp, from)
.await;
tracing::debug!("[ALIGNER] Received {} queried samples", missing_data.len());
tracing::trace!("[ALIGNER] Received queried samples: {missing_data:?}");
for (key, (ts, value)) in missing_data {
let sample = SampleBuilder::put(key, value.payload().clone())
.encoding(value.encoding().clone())
.timestamp(ts)
.into();
tracing::debug!("[ALIGNER] Adding {:?} to storage", sample);
self.tx_sample.send_async(sample).await.unwrap_or_else(|e| {
tracing::error!("[ALIGNER] Error adding sample to storage: {}", e)
});
}
if no_content_err && no_data_err {
let mut processed = self.digests_processed.write().await;
(*processed).insert(checksum);
}
}
}
async fn get_missing_data(
&self,
missing_content: &[LogEntry],
timestamp: Timestamp,
from: &str,
) -> (HashMap<OwnedKeyExpr, (Timestamp, Value)>, bool) {
let mut result = HashMap::new();
let parameters = format!(
"timestamp={}&{}={}",
timestamp,
CONTENTS,
serde_json::to_string(missing_content).unwrap()
);
let (replies, no_err) = self.perform_query(from, parameters.clone()).await;
for sample in replies {
result.insert(
sample.key_expr().clone().into(),
(*sample.timestamp().unwrap(), Value::from(sample)),
);
}
(result, no_err)
}
async fn get_missing_content(&self, other: &Digest, from: &str) -> (Vec<LogEntry>, bool) {
tracing::debug!("[ALIGNER] Get missing content from {from} ...");
let this = &self.snapshotter.get_digest().await;
let cold_alignment =
self.perform_era_alignment(&EraType::Cold, this, from.to_string(), other);
let warm_alignment =
self.perform_era_alignment(&EraType::Warm, this, from.to_string(), other);
let hot_alignment =
self.perform_era_alignment(&EraType::Hot, this, from.to_string(), other);
let ((cold_data, no_cold_err), (warm_data, no_warm_err), (hot_data, no_hot_err)) =
futures::join!(cold_alignment, warm_alignment, hot_alignment);
tracing::debug!("[ALIGNER] Missing content from {from} in Cold era: {cold_data:?}");
tracing::debug!("[ALIGNER] Missing content from {from} in Warm era: {warm_data:?}");
tracing::debug!("[ALIGNER] Missing content from {from} in Hot era: {hot_data:?}");
(
[cold_data, warm_data, hot_data].concat(),
no_cold_err && no_warm_err && no_hot_err,
)
}
async fn perform_era_alignment(
&self,
era: &EraType,
this: &Digest,
other_rep: String,
other: &Digest,
) -> (Vec<LogEntry>, bool) {
if !this.era_has_diff(era, &other.eras) {
return (Vec::new(), true);
}
let (diff_intervals, no_era_err) =
self.get_interval_diff(era, this, other, &other_rep).await;
let (diff_subintervals, no_int_err) = self
.get_subinterval_diff(era, diff_intervals, this, other, &other_rep)
.await;
let (diff_content, no_sub_err) = self
.get_content_diff(diff_subintervals, this, other, &other_rep)
.await;
(diff_content, no_era_err && no_int_err && no_sub_err)
}
async fn get_interval_diff(
&self,
era: &EraType,
this: &Digest,
other: &Digest,
other_rep: &str,
) -> (HashSet<u64>, bool) {
let (other_intervals, no_err) = if era.eq(&EraType::Cold) {
let parameters = format!("timestamp={}&{}=cold", other.timestamp, ERA);
let (reply_content, mut no_err) = self.perform_query(other_rep, parameters).await;
let mut other_intervals: HashMap<u64, u64> = HashMap::new();
for each in reply_content {
match serde_json::from_reader(each.payload().reader()) {
Ok((i, c)) => {
other_intervals.insert(i, c);
}
Err(e) => {
tracing::error!("[ALIGNER] Error decoding reply: {}", e);
no_err = false;
}
}
}
(other_intervals, no_err)
} else {
(other.get_era_content(era), true)
};
(this.get_interval_diff(other_intervals), no_err)
}
async fn get_subinterval_diff(
&self,
era: &EraType,
diff_intervals: HashSet<u64>,
this: &Digest,
other: &Digest,
other_rep: &str,
) -> (HashSet<u64>, bool) {
if !diff_intervals.is_empty() {
let (other_subintervals, no_err) = if era.eq(&EraType::Hot) {
(other.get_interval_content(diff_intervals), true)
} else {
let mut diff_string = Vec::new();
for each_int in diff_intervals {
diff_string.push(each_int.to_string());
}
let parameters = format!(
"timestamp={}&{}=[{}]",
other.timestamp,
INTERVALS,
diff_string.join(",")
);
let (reply_content, mut no_err) = self.perform_query(other_rep, parameters).await;
let mut other_subintervals: HashMap<u64, u64> = HashMap::new();
for each in reply_content {
match serde_json::from_reader(each.payload().reader()) {
Ok((i, c)) => {
other_subintervals.insert(i, c);
}
Err(e) => {
tracing::error!("[ALIGNER] Error decoding reply: {}", e);
no_err = false;
}
}
}
(other_subintervals, no_err)
};
(this.get_subinterval_diff(other_subintervals), no_err)
} else {
(HashSet::new(), true)
}
}
async fn get_content_diff(
&self,
diff_subintervals: HashSet<u64>,
this: &Digest,
other: &Digest,
other_rep: &str,
) -> (Vec<LogEntry>, bool) {
if !diff_subintervals.is_empty() {
let mut diff_string = Vec::new();
for each_sub in diff_subintervals {
diff_string.push(each_sub.to_string());
}
let parameters = format!(
"timestamp={}&{}=[{}]",
other.timestamp,
SUBINTERVALS,
diff_string.join(",")
);
let (reply_content, mut no_err) = self.perform_query(other_rep, parameters).await;
let mut other_content: HashMap<u64, Vec<LogEntry>> = HashMap::new();
for each in reply_content {
match serde_json::from_reader(each.payload().reader()) {
Ok((i, c)) => {
other_content.insert(i, c);
}
Err(e) => {
tracing::error!("[ALIGNER] Error decoding reply: {}", e);
no_err = false;
}
}
}
let result = this.get_full_content_diff(other_content);
(result, no_err)
} else {
(Vec::new(), true)
}
}
async fn perform_query(&self, from: &str, parameters: String) -> (Vec<Sample>, bool) {
let mut no_err = true;
let selector = Selector::owned(
KeyExpr::from(&self.digest_key).join(&from).unwrap(),
parameters,
);
tracing::trace!("[ALIGNER] Sending Query '{}'...", selector);
let mut return_val = Vec::new();
match self
.session
.get(&selector)
.consolidation(zenoh::query::ConsolidationMode::None)
.accept_replies(zenoh::query::ReplyKeyExpr::Any)
.await
{
Ok(replies) => {
while let Ok(reply) = replies.recv_async().await {
match reply.into_result() {
Ok(sample) => {
tracing::trace!(
"[ALIGNER] Received ('{}': '{}')",
sample.key_expr().as_str(),
sample
.payload()
.deserialize::<Cow<str>>()
.unwrap_or(Cow::Borrowed("<malformed>"))
);
return_val.push(sample);
}
Err(err) => {
tracing::error!(
"[ALIGNER] Received error for query on selector {} :{:?}",
selector,
err
);
no_err = false;
}
}
}
}
Err(err) => {
tracing::error!(
"[ALIGNER] Query failed on selector `{}`: {:?}",
selector,
err
);
no_err = false;
}
};
tracing::trace!(
"[ALIGNER] On Query '{selector}' received: {return_val:?} (no_err:{no_err})"
);
(return_val, no_err)
}
}