139 lines
3.7 KiB
Rust
139 lines
3.7 KiB
Rust
use daqbuf_err as err;
|
|
use err::thiserror;
|
|
use err::ThisError;
|
|
use futures_util::StreamExt;
|
|
use netpod::log::*;
|
|
use netpod::ttl::RetentionTime;
|
|
use netpod::TsMs;
|
|
use netpod::EMIT_ACCOUNTING_SNAP;
|
|
use scylla::prepared_statement::PreparedStatement;
|
|
use scylla::Session as ScySession;
|
|
|
|
#[derive(Debug, ThisError)]
|
|
#[cstm(name = "AccountingToplist")]
|
|
pub enum Error {
|
|
ScyllaQuery(#[from] scylla::transport::errors::QueryError),
|
|
ScyllaNextRow(#[from] scylla::transport::iterator::NextRowError),
|
|
UsageDataMalformed,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct UsageData {
|
|
ts: TsMs,
|
|
series: Vec<u64>,
|
|
counts: Vec<u64>,
|
|
bytes: Vec<u64>,
|
|
}
|
|
|
|
impl UsageData {
|
|
pub fn new(ts: TsMs) -> Self {
|
|
Self {
|
|
ts,
|
|
series: Vec::new(),
|
|
counts: Vec::new(),
|
|
bytes: Vec::new(),
|
|
}
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.series.len()
|
|
}
|
|
|
|
pub fn ts(&self) -> TsMs {
|
|
self.ts
|
|
}
|
|
|
|
pub fn series(&self) -> &[u64] {
|
|
&self.series
|
|
}
|
|
|
|
pub fn counts(&self) -> &[u64] {
|
|
&self.counts
|
|
}
|
|
|
|
pub fn bytes(&self) -> &[u64] {
|
|
&self.bytes
|
|
}
|
|
|
|
pub fn sort_by_counts(&mut self) {
|
|
let mut tmp: Vec<_> = self
|
|
.counts
|
|
.iter()
|
|
.map(|&x| x)
|
|
.enumerate()
|
|
.map(|(i, x)| (x, i))
|
|
.collect();
|
|
tmp.sort_unstable();
|
|
let tmp: Vec<_> = tmp.into_iter().rev().map(|x| x.1).collect();
|
|
self.reorder_by_index_list(&tmp);
|
|
}
|
|
|
|
pub fn sort_by_bytes(&mut self) {
|
|
let mut tmp: Vec<_> = self.bytes.iter().map(|&x| x).enumerate().map(|(i, x)| (x, i)).collect();
|
|
tmp.sort_unstable();
|
|
let tmp: Vec<_> = tmp.into_iter().rev().map(|x| x.1).collect();
|
|
self.reorder_by_index_list(&tmp);
|
|
}
|
|
|
|
fn reorder_by_index_list(&mut self, tmp: &[usize]) {
|
|
self.series = tmp.iter().map(|&x| self.series[x]).collect();
|
|
self.counts = tmp.iter().map(|&x| self.counts[x]).collect();
|
|
self.bytes = tmp.iter().map(|&x| self.bytes[x]).collect();
|
|
}
|
|
|
|
fn verify(&self) -> Result<(), Error> {
|
|
if self.counts.len() != self.series.len() {
|
|
Err(Error::UsageDataMalformed)
|
|
} else if self.bytes.len() != self.series.len() {
|
|
Err(Error::UsageDataMalformed)
|
|
} else {
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn read_ts(ks: &str, rt: RetentionTime, ts: TsMs, scy: &ScySession) -> Result<UsageData, Error> {
|
|
// TODO toplist::read_ts refactor
|
|
let snap = EMIT_ACCOUNTING_SNAP.ms();
|
|
let ts = TsMs::from_ms_u64(ts.ms() / snap * snap);
|
|
let ret = read_ts_inner(ks, rt, ts, scy).await?;
|
|
Ok(ret)
|
|
}
|
|
|
|
async fn read_ts_inner(ks: &str, rt: RetentionTime, ts: TsMs, scy: &ScySession) -> Result<UsageData, Error> {
|
|
type RowType = (i64, i64, i64);
|
|
let cql = format!(
|
|
concat!(
|
|
"select series, count, bytes",
|
|
" from {}.{}account_00",
|
|
" where part = ? and ts = ?"
|
|
),
|
|
ks,
|
|
rt.table_prefix()
|
|
);
|
|
let qu = prep(&cql, scy).await?;
|
|
let ts_sec = ts.ms() as i64 / 1000;
|
|
let mut ret = UsageData::new(ts);
|
|
for part in 0..255_u32 {
|
|
let mut res = scy
|
|
.execute_iter(qu.clone(), (part as i32, ts_sec))
|
|
.await?
|
|
.into_typed::<RowType>();
|
|
while let Some(row) = res.next().await {
|
|
let row = row?;
|
|
let series = row.0 as u64;
|
|
let count = row.1 as u64;
|
|
let bytes = row.2 as u64;
|
|
ret.series.push(series);
|
|
ret.counts.push(count);
|
|
ret.bytes.push(bytes);
|
|
}
|
|
}
|
|
ret.verify()?;
|
|
Ok(ret)
|
|
}
|
|
|
|
async fn prep(cql: &str, scy: &ScySession) -> Result<PreparedStatement, Error> {
|
|
Ok(scy.prepare(cql).await?)
|
|
}
|