Files
daqingest/stats_proc/src/stats_proc.rs
Dominik Werder 79aa8d0466 WIP
2024-08-30 16:29:47 +02:00

727 lines
20 KiB
Rust

use proc_macro::TokenStream;
use quote::quote;
use syn::parse::ParseStream;
use syn::parse_macro_input;
use syn::Ident;
type PunctExpr = syn::punctuated::Punctuated<syn::Expr, syn::token::Comma>;
struct FuncCallWithArgs {
name: Ident,
args: PunctExpr,
}
#[derive(Clone, Debug)]
struct StatsStructDef {
name: syn::Ident,
prefix: Option<syn::Ident>,
counters: Vec<syn::Ident>,
values: Vec<syn::Ident>,
histolog2s: Vec<syn::Ident>,
}
#[derive(Debug)]
struct AggStructDef {
name: syn::Ident,
parent: syn::Ident,
// TODO this currently describes our input (especially the input's name):
stats: StatsStructDef,
}
#[derive(Debug)]
struct DiffStructDef {
name: syn::Ident,
input: syn::Ident,
}
fn extend_str(mut a: String, x: impl AsRef<str>) -> String {
a.push_str(x.as_ref());
a
}
fn stats_struct_impl(st: &StatsStructDef) -> String {
use std::fmt::Write;
let name = &st.name;
let inits1 = st
.counters
.iter()
.map(|x| format!("{:12}{}: stats_types::Counter::new()", "", x.to_string()));
let inits2 = st
.values
.iter()
.map(|x| format!("{:12}{}: stats_types::Value::new()", "", x.to_string()));
let init_histolog2s = st
.histolog2s
.iter()
.map(|x| format!("{:12}{}: stats_types::HistoLog2::new()", "", x.to_string()));
let inits: Vec<_> = inits1.into_iter().chain(inits2).chain(init_histolog2s).collect();
let inits = inits.join(",\n");
let incers: String = st
.counters
.iter()
.map(|nn| {
format!(
"
pub fn {nn}(&self) -> &stats_types::Counter {{
&self.{nn}
}}
//pub fn {nn}_dur(&self, v: Duration) {{
// self.{nn}.fetch_add((v * 1000000).as_secs(), Ordering::AcqRel);
//}}
"
)
})
.fold(String::new(), |mut a, x| {
a.push_str(&x);
a
});
let values = {
let mut buf = String::new();
for nn in &st.values {
write!(
buf,
"
pub fn {nn}(&self) -> &stats_types::Value {{
&self.{nn}
}}
"
)
.unwrap();
}
buf
};
let histolog2s = {
let mut buf = String::new();
for nn in &st.histolog2s {
write!(
buf,
"
pub fn {nn}(&self) -> &stats_types::HistoLog2 {{
&self.{nn}
}}
"
)
.unwrap();
}
buf
};
let fn_prometheus = {
let mut buf = String::new();
for x in &st.counters {
let n = x.to_string();
let pre = if let Some(x) = &st.prefix {
format!("_{}", x)
} else {
String::new()
};
buf.push_str(&format!(
"ret.push_str(&format!(\"daqingest{}_{} {{}}\\n\", self.{}.load()));\n",
pre, n, n
));
}
for x in &st.values {
let n = x.to_string();
let pre = if let Some(x) = &st.prefix {
format!("_{}", x)
} else {
String::new()
};
buf.push_str(&format!(
"ret.push_str(&format!(\"daqingest{}_{} {{}}\\n\", self.{}.load()));\n",
pre, n, n
));
}
for x in &st.histolog2s {
let n = x.to_string();
let fullname = if let Some(x) = &st.prefix {
format!("daqingest_{x}_{n}")
} else {
format!("daqingest_{n}")
};
buf.push_str(&format!("let fullname = \"{fullname}\";\n"));
buf.push_str(&format!("let s = self.{n}.to_prometheus(fullname);\n"));
buf.push_str(&format!("ret.push_str(&s);\n"));
}
format!(
"
pub fn prometheus(&self) -> String {{
let mut ret = String::new();
{buf}
ret
}}
"
)
};
let fn_json = {
let mut buf = String::new();
for x in &st.counters {
buf.push_str(&format!(
"ret.insert(\"{x}\".to_string(), Value::Number(Number::from(self.{x}.load())));\n"
));
}
for x in &st.values {
buf.push_str(&format!(
"ret.insert(\"{x}\".to_string(), Value::Number(Number::from(self.{x}.load())));\n"
));
}
for x in &st.histolog2s {
buf.push_str(&format!("let v = self.{x}.to_json(\"__dummyname__\");\n"));
buf.push_str(&format!("ret.insert(\"{x}\".to_string(), v);\n"));
}
format!(
"
pub fn json(&self) -> stats_types::serde_json::Value {{
use serde_json::Map;
use serde_json::Number;
use serde_json::Value;
use stats_types::serde_json;
let mut ret = Map::new();
// {buf}
Value::Object(ret)
}}
"
)
};
let fn_snapshot = {
let mut init_counters = String::new();
for x in &st.counters {
let n = x.to_string();
init_counters.push_str(&format!("ret.{}.__set(self.{}.load());\n", n, n));
}
let mut init_values = String::new();
for x in &st.values {
let n = x.to_string();
init_values.push_str(&format!("ret.{}.set(self.{}.load());\n", n, n));
}
format!(
"
pub fn snapshot(&self) -> Self {{
let ret = Self::new();
{init_counters}
{init_values}
ret
}}
"
)
};
format!(
"
impl {name} {{
pub fn new() -> Self {{
Self {{
ts_create: Instant::now(),
dropped: stats_types::Value::new(),
{inits}
}}
}}
{incers}
{values}
{histolog2s}
{fn_prometheus}
{fn_json}
{fn_snapshot}
}}
impl stats_types::DropMark for {name} {{
fn field(&self) -> &stats_types::Value {{
&self.dropped
}}
}}
"
)
}
fn stats_struct_decl_impl(st: &StatsStructDef) -> String {
let name = &st.name;
let counters_decl = st
.counters
.iter()
.map(|x| format!("{:4}pub {}: stats_types::Counter,\n", "", x.to_string()))
.fold(String::new(), extend_str);
let values_decl = st
.values
.iter()
.map(|x| format!("{:4}pub {}: stats_types::Value,\n", "", x.to_string()))
.fold(String::new(), extend_str);
let histolog2s_decl = st
.histolog2s
.iter()
.map(|x| format!("{:4}pub {}: stats_types::HistoLog2,\n", "", x.to_string()))
.fold(String::new(), extend_str);
let structt = format!(
"
pub struct {name} {{
pub ts_create: Instant,
dropped: stats_types::Value,
{counters_decl}
{values_decl}
{histolog2s_decl}
}}
"
);
let code = format!("{}\n\n{}", structt, stats_struct_impl(st),);
code
}
fn agg_decl_impl(st: &StatsStructDef, ag: &AggStructDef) -> String {
let name = &ag.name;
let name_inp = &st.name;
let counters_decl = st
.counters
.iter()
.map(|x| format!("{:4}pub {}: stats_types::Counter,\n", "", x.to_string()))
.fold(String::new(), extend_str);
let mut code = String::new();
let s = format!(
"
// Agg decl
pub struct {name} {{
pub ts_create: Instant,
pub aggcount: stats_types::Counter,
{counters_decl}
}}
"
);
code.push_str(&s);
let clone_counters = st
.counters
.iter()
.map(|x| {
let n = x.to_string();
format!("{:12}{}: stats_types::Counter::init(self.{}.load()),\n", "", n, n)
})
.fold(String::new(), extend_str);
let s = format!(
"
impl Clone for {name} {{
fn clone(&self) -> Self {{
Self {{
ts_create: self.ts_create.clone(),
aggcount: stats_types::Counter::init(self.aggcount.load()),
{clone_counters}
}}
}}
}}
"
);
code.push_str(&s);
let inits = st
.counters
.iter()
.map(|x| format!("{:12}{}: stats_types::Counter::new(),\n", "", x.to_string()))
.fold(String::new(), extend_str);
let s = format!(
"
// Agg impl
impl {name} {{
pub fn new() -> Self {{
Self {{
ts_create: Instant::now(),
aggcount: stats_types::Counter::new(),
{inits}
}}
}}
"
);
code.push_str(&s);
let counters_add = st
.counters
.iter()
.map(|x| format!("self.{}.add(inp.{}.load());\n", x.to_string(), x.to_string()))
.fold(String::new(), extend_str);
let s = format!(
"
pub fn push(&self, inp: &{name_inp}) {{
self.aggcount.inc();
{counters_add}
}}
"
);
code.push_str(&s);
{
let mut buf = String::new();
for x in &st.counters {
let n = x.to_string();
buf.push_str(&format!(
"ret.push_str(&format!(\"daqingest_{} {{}}\\n\", self.{}.load()));\n",
n, n
));
}
let s = format!(
"
pub fn prometheus(&self) -> String {{
let mut ret = String::new();
ret.push_str(&format!(\"daqingest_aggcount {{}}\\n\", self.aggcount.load()));
{buf}
ret
}}
"
);
code.push_str(&s);
}
code.push_str(
"
}
",
);
code
}
// TODO maybe basic and agg structs need a different treatment?
// Should probably implement the methods behind a common trait.
fn diff_decl_impl(st: &DiffStructDef, inp: &StatsStructDef) -> String {
let name = &st.name;
let inp_ty = &inp.name;
let decl = inp
.counters
.iter()
.map(|x| format!("{:4}pub {}: stats_types::Counter,\n", "", x.to_string()))
.fold(String::new(), extend_str);
let mut code = String::new();
let s = format!(
"
pub struct {name} {{
pub dt: AtomicU64,
{decl}
}}
"
);
code.push_str(&s);
code.push_str(&format!(
"impl {name} {{
"
));
let diffs = inp
.counters
.iter()
.map(|x| {
let n = x.to_string();
format!(
"{:12}let {} = stats_types::Counter::init(b.{}.load() - a.{}.load());\n",
"", n, n, n
)
})
.fold(String::new(), extend_str);
let inits = inp
.counters
.iter()
.map(|x| {
let n = x.to_string();
format!("{:16}{},\n", "", n)
})
.fold(String::new(), extend_str);
let s = format!(
"
pub fn diff_from(a: &{inp_ty}, b: &{inp_ty}) -> Self {{
let dur = b.ts_create.duration_since(a.ts_create);
{diffs}
Self {{
dt: AtomicU64::new(dur.as_secs() * SEC + dur.subsec_nanos() as u64),
{inits}
}}
}}
"
);
code.push_str(&s);
let mut a = String::new();
let mut b = String::new();
for h in &inp.counters {
a.push_str(&format!("{} {{}} ", h.to_string()));
b.push_str(&format!("self.{}.load(), ", h.to_string()));
}
let s = format!(
"
pub fn display(&self) -> String {{
format!(\"dt {{}} {a}\", self.dt.load(Ordering::Acquire) / 1000000, {b})
}}
"
);
code.push_str(&s);
code.push_str(
"
}
",
);
code
}
fn ident_from_expr(inp: syn::Expr) -> syn::Result<syn::Ident> {
use syn::spanned::Spanned;
match inp {
syn::Expr::Path(k) => {
if k.path.segments.len() == 1 {
Ok(k.path.segments[0].ident.clone())
} else {
Err(syn::Error::new(k.span(), "Expect identifier"))
}
}
_ => Err(syn::Error::new(inp.span(), "Expect identifier")),
}
}
fn idents_from_exprs(inp: PunctExpr) -> syn::Result<Vec<syn::Ident>> {
let mut ret = Vec::new();
for k in inp {
let g = ident_from_expr(k)?;
ret.push(g);
}
Ok(ret)
}
fn func_name_from_expr(inp: syn::Expr) -> syn::Result<syn::Ident> {
use syn::spanned::Spanned;
use syn::Error;
use syn::Expr;
match inp {
Expr::Path(k) => {
if k.path.segments.len() != 1 {
return Err(Error::new(k.span(), "Expect function name"));
}
let res = k.path.segments[0].ident.clone();
Ok(res)
}
_ => {
return Err(Error::new(inp.span(), "Expect function name"));
}
}
}
impl FuncCallWithArgs {
fn from_expr(inp: syn::Expr) -> Result<Self, syn::Error> {
use syn::spanned::Spanned;
use syn::Error;
use syn::Expr;
let span_all = inp.span();
match inp {
Expr::Call(k) => {
let name = func_name_from_expr(*k.func)?;
let args = k.args;
let ret = FuncCallWithArgs { name, args };
Ok(ret)
}
_ => {
return Err(Error::new(span_all, format!("BAD {:?}", inp)));
}
}
}
}
impl StatsStructDef {
fn empty() -> Self {
Self {
name: syn::parse_str("__empty").unwrap(),
prefix: syn::parse_str("__empty").unwrap(),
counters: Vec::new(),
values: Vec::new(),
histolog2s: Vec::new(),
}
}
fn from_args(inp: PunctExpr) -> syn::Result<Self> {
let mut name = None;
let mut prefix = None;
let mut counters = None;
let mut values = None;
let mut histolog2s = None;
for k in inp {
let fa = FuncCallWithArgs::from_expr(k)?;
if fa.name == "name" {
let ident = ident_from_expr(fa.args[0].clone())?;
name = Some(ident);
} else if fa.name == "prefix" {
let ident = ident_from_expr(fa.args[0].clone())?;
prefix = Some(ident);
} else if fa.name == "counters" {
let idents = idents_from_exprs(fa.args)?;
counters = Some(idents);
} else if fa.name == "values" {
let idents = idents_from_exprs(fa.args)?;
values = Some(idents);
} else if fa.name == "histolog2s" {
let idents = idents_from_exprs(fa.args)?;
histolog2s = Some(idents);
} else {
panic!("fa.name: {:?}", fa.name);
}
}
let ret = StatsStructDef {
name: name.expect("Expect name for StatsStructDef"),
prefix,
counters: counters.unwrap_or(Vec::new()),
values: values.unwrap_or(Vec::new()),
histolog2s: histolog2s.unwrap_or(Vec::new()),
};
Ok(ret)
}
}
impl AggStructDef {
fn from_args(inp: PunctExpr) -> syn::Result<Self> {
let mut name = None;
let mut parent = None;
for k in inp {
let fa = FuncCallWithArgs::from_expr(k)?;
if fa.name == "name" {
let ident = ident_from_expr(fa.args[0].clone())?;
name = Some(ident);
}
if fa.name == "parent" {
let ident = ident_from_expr(fa.args[0].clone())?;
parent = Some(ident);
}
}
let ret = AggStructDef {
name: name.expect("Expect name for AggStructDef"),
// Will get resolved later:
stats: StatsStructDef::empty(),
parent: parent.expect("Expect parent"),
};
Ok(ret)
}
}
impl DiffStructDef {
fn from_args(inp: PunctExpr) -> syn::Result<Self> {
let mut name = None;
let mut input = None;
for k in inp {
let fa = FuncCallWithArgs::from_expr(k)?;
if fa.name == "name" {
let ident = ident_from_expr(fa.args[0].clone())?;
name = Some(ident);
}
if fa.name == "input" {
let ident = ident_from_expr(fa.args[0].clone())?;
input = Some(ident);
}
}
let ret = DiffStructDef {
name: name.expect("Expect name for DiffStructDef"),
input: input.expect("Expect input for DiffStructDef"),
};
Ok(ret)
}
}
#[derive(Debug)]
struct StatsTreeDef {
stats_struct_defs: Vec<StatsStructDef>,
agg_defs: Vec<AggStructDef>,
diff_defs: Vec<DiffStructDef>,
}
impl syn::parse::Parse for StatsTreeDef {
fn parse(inp: ParseStream) -> syn::Result<Self> {
let k = inp.parse::<syn::ExprTuple>()?;
let mut a = Vec::new();
let mut agg_defs = Vec::new();
let mut diff_defs = Vec::new();
for k in k.elems {
let fa = FuncCallWithArgs::from_expr(k)?;
if fa.name == "stats_struct" {
let stats_struct_def = StatsStructDef::from_args(fa.args)?;
a.push(stats_struct_def);
} else if fa.name == "agg" {
let agg_def = AggStructDef::from_args(fa.args)?;
agg_defs.push(agg_def);
} else if fa.name == "diff" {
let diff_def = DiffStructDef::from_args(fa.args)?;
diff_defs.push(diff_def);
} else {
return Err(syn::Error::new(fa.name.span(), "Unexpected"));
}
}
let ret = StatsTreeDef {
stats_struct_defs: a,
agg_defs,
diff_defs,
};
Ok(ret)
}
}
#[proc_macro]
pub fn stats_struct(ts: TokenStream) -> TokenStream {
let mut def: StatsTreeDef = parse_macro_input!(ts);
for h in &mut def.agg_defs {
for k in &def.stats_struct_defs {
if k.name == h.parent {
// TODO factor this out..
h.stats = k.clone();
}
}
}
if false {
for j in &def.agg_defs {
let h = StatsStructDef {
name: j.name.clone(),
prefix: None,
counters: j.stats.counters.clone(),
values: Vec::new(),
histolog2s: Vec::new(),
};
def.stats_struct_defs.push(h);
}
}
let mut code = String::new();
let mut ts1 = TokenStream::new();
for k in &def.stats_struct_defs {
let s = stats_struct_decl_impl(k);
code.push_str(&s);
let ts2: TokenStream = s.parse().unwrap();
ts1.extend(ts2);
}
for k in &def.agg_defs {
for st in &def.stats_struct_defs {
if st.name == k.parent {
let s = agg_decl_impl(st, k);
code.push_str(&s);
let ts2: TokenStream = s.parse().unwrap();
ts1.extend(ts2);
}
}
}
for k in &def.diff_defs {
for j in &def.agg_defs {
if j.name == k.input {
// TODO currently, "j.stats" describes the input to the "agg", so that contains the wrong name.
let p = StatsStructDef {
name: k.input.clone(),
// TODO refactor
prefix: None,
counters: j.stats.counters.clone(),
// TODO compute values
values: Vec::new(),
// TODO not supported yet
histolog2s: Vec::new(),
};
let s = diff_decl_impl(k, &p);
code.push_str(&s);
let ts2: TokenStream = s.parse().unwrap();
ts1.extend(ts2);
}
}
for j in &def.stats_struct_defs {
if j.name == k.input {
let s = diff_decl_impl(k, j);
code.push_str(&s);
let ts2: TokenStream = s.parse().unwrap();
ts1.extend(ts2);
}
}
}
//panic!("CODE: {}", code);
let _ts3 = TokenStream::from(quote!(
mod asd {}
));
ts1
}