diff --git a/lib/pk_core/Cargo.toml b/lib/pk_core/Cargo.toml new file mode 100644 index 00000000..71072412 --- /dev/null +++ b/lib/pk_core/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "pk_core" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +chrono = "0.4.19" +reqwest = { version = "0.11.11", default-features = false, features = ["blocking", "rustls-tls"] } +tokio = { version = "1", features = ["full"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["json","std","env-filter"] } +tracing-appender = "0.2" \ No newline at end of file diff --git a/lib/pk_core/src/lib.rs b/lib/pk_core/src/lib.rs new file mode 100644 index 00000000..d53791ac --- /dev/null +++ b/lib/pk_core/src/lib.rs @@ -0,0 +1,2 @@ +pub mod tracing; +pub mod util; \ No newline at end of file diff --git a/lib/pk_core/src/tracing.rs b/lib/pk_core/src/tracing.rs new file mode 100644 index 00000000..bac79efe --- /dev/null +++ b/lib/pk_core/src/tracing.rs @@ -0,0 +1,105 @@ +use reqwest::blocking::{Client, ClientBuilder}; + +use tracing::*; +use tracing_subscriber::{fmt, layer::SubscriberExt, EnvFilter}; +use tracing_appender::non_blocking::WorkerGuard; + +use std::io::{Error, ErrorKind}; + +#[macro_export] +macro_rules! init { + ($e:expr) => { + let _guard1 = pk_core::tracing::init_inner(None, None, None).await; + let _guard2 = span!(Level::ERROR, "service_name", service_name = $e).entered(); + }; + ($e:expr, $f:expr) => { + let _guard1 = pk_core::tracing::init_inner($f, None, None).await; + let _guard2 = span!(Level::ERROR, "service_name", service_name = $e).entered(); + }; + ($e:expr, $f:expr, $g:expr, $h:expr) => { + let _guard1 = pk_core::tracing::init_inner($f, $g, $h).await; + let _guard2 = span!(Level::ERROR, "service_name", service_name = $e).entered(); + }; +} + +pub use init; + +pub async fn init_inner(url: Option, user: Option<&str>, password: Option<&str>) -> Option { + let stdout_subscriber = tracing_subscriber::fmt().with_max_level(Level::INFO).finish(); + + if url.is_none() { + tracing::subscriber::set_global_default(stdout_subscriber) + .expect("Unable to set a global collector"); + return None; + } + + let client = tokio::task::spawn_blocking(|| { + let client = ClientBuilder::new() + .danger_accept_invalid_certs(true); + client.build() + }).await.unwrap().unwrap(); + + let elastic = ElasticLogWriter { + user: user.map(|v| v.to_string()), + password: password.map(|v| v.to_string()), + url: url.unwrap(), + client + }; + + let (non_blocking, _guard) = tracing_appender::non_blocking::NonBlockingBuilder::default() + .lossy(true) + .buffered_lines_limit(2000) + .finish(elastic); + + let subscriber = stdout_subscriber + .with(EnvFilter::from_default_env().add_directive(tracing::Level::TRACE.into())) + .with( + fmt::Layer::new() + .json() + .flatten_event(true) + .with_writer(non_blocking) + ); + + tracing::subscriber::set_global_default(subscriber) + .expect("Unable to set a global collector"); + + return Some(_guard); +} + +pub struct ElasticLogWriter { + user: Option, + password: Option, + url: String, + client: Client, +} + +impl std::io::Write for ElasticLogWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let data = String::from_utf8(buf.to_vec()); + if data.is_err() { + return Err(Error::new(ErrorKind::Other, data.err().unwrap())); + } + + let orig_json = data.unwrap(); + + let cur_date = chrono::Utc::now().format("%Y-%m-%d"); + let mut builder = self.client.post(format!("{}/pluralkit-logs-{}/_doc", self.url.clone(), cur_date)) + .header("content-type", "application/json") + .body(orig_json); + + if self.user.is_some() { + builder = builder.basic_auth(self.user.as_ref().unwrap(), self.password.as_ref()); + } + + let res = builder.send(); + + match res { + Ok(_) => Ok(buf.len()), + Err(err) => Err(Error::new(ErrorKind::Other, err)), + } + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} diff --git a/lib/pk_core/src/util.rs b/lib/pk_core/src/util.rs new file mode 100644 index 00000000..a3defeaa --- /dev/null +++ b/lib/pk_core/src/util.rs @@ -0,0 +1,7 @@ +/// get_env gets an env variable as Option. +pub fn get_env(key: &str) -> Option { + match std::env::var(key) { + Ok(val) => { Some(val) } + Err(_) => None + } +} \ No newline at end of file