use protosocket_rpc::{Message, ProtosocketControlCode};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum WireFieldValue {
U64(u64),
I64(i64),
F64(f64),
Bool(bool),
Str(String),
}
impl WireFieldValue {
pub fn to_string_value(&self) -> String {
match self {
WireFieldValue::U64(v) => v.to_string(),
WireFieldValue::I64(v) => v.to_string(),
WireFieldValue::F64(v) => v.to_string(),
WireFieldValue::Bool(v) => v.to_string(),
WireFieldValue::Str(s) => s.clone(),
}
}
pub fn contains(&self, needle: &str) -> bool {
match self {
WireFieldValue::Str(s) => s.contains(needle),
other => other.to_string_value().contains(needle),
}
}
}
impl std::fmt::Display for WireFieldValue {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
WireFieldValue::U64(v) => write!(f, "{v}"),
WireFieldValue::I64(v) => write!(f, "{v}"),
WireFieldValue::F64(v) => write!(f, "{v}"),
WireFieldValue::Bool(v) => write!(f, "{v}"),
WireFieldValue::Str(s) => f.write_str(s),
}
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum WireLevel {
Trace,
Debug,
Info,
Warn,
Error,
}
impl WireLevel {
pub fn from_tracing(level: &tracing::Level) -> Self {
match *level {
tracing::Level::TRACE => WireLevel::Trace,
tracing::Level::DEBUG => WireLevel::Debug,
tracing::Level::INFO => WireLevel::Info,
tracing::Level::WARN => WireLevel::Warn,
tracing::Level::ERROR => WireLevel::Error,
}
}
pub fn to_tracing(self) -> tracing::Level {
match self {
WireLevel::Trace => tracing::Level::TRACE,
WireLevel::Debug => tracing::Level::DEBUG,
WireLevel::Info => tracing::Level::INFO,
WireLevel::Warn => tracing::Level::WARN,
WireLevel::Error => tracing::Level::ERROR,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WireEvent {
pub name: String,
pub level: WireLevel,
pub fields: Vec<(String, WireFieldValue)>,
pub recorded_at_ns: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WireSpan {
pub id: u64,
pub parent_id: Option<u64>,
pub name: String,
pub target: String,
pub level: WireLevel,
pub fields: Vec<(String, WireFieldValue)>,
pub events: Vec<WireEvent>,
pub opened_at_ns: u64,
pub closed_at_ns: Option<u64>,
}
impl WireSpan {
pub fn field(&self, name: &str) -> Option<&WireFieldValue> {
self.fields.iter().find(|(k, _)| k == name).map(|(_, v)| v)
}
}
impl WireEvent {
pub fn field(&self, name: &str) -> Option<&WireFieldValue> {
self.fields.iter().find(|(k, _)| k == name).map(|(_, v)| v)
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum WireLevelFilter {
Off,
Error,
Warn,
Info,
Debug,
Trace,
}
impl WireLevelFilter {
pub fn from_tracing(filter: tracing::metadata::LevelFilter) -> Self {
use tracing::metadata::LevelFilter as L;
if filter == L::OFF {
WireLevelFilter::Off
} else if filter == L::ERROR {
WireLevelFilter::Error
} else if filter == L::WARN {
WireLevelFilter::Warn
} else if filter == L::INFO {
WireLevelFilter::Info
} else if filter == L::DEBUG {
WireLevelFilter::Debug
} else {
WireLevelFilter::Trace
}
}
pub fn to_tracing(self) -> tracing::metadata::LevelFilter {
use tracing::metadata::LevelFilter as L;
match self {
WireLevelFilter::Off => L::OFF,
WireLevelFilter::Error => L::ERROR,
WireLevelFilter::Warn => L::WARN,
WireLevelFilter::Info => L::INFO,
WireLevelFilter::Debug => L::DEBUG,
WireLevelFilter::Trace => L::TRACE,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RequestBody {
Noop,
StartStream,
StopStream,
SetLevel(WireLevel),
SetCacheLevel(WireLevelFilter),
SetCacheChance(f64),
SetSamplingRate(f64),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Request {
pub id: u64,
pub control: u8,
pub body: RequestBody,
}
impl Request {
pub fn new(body: RequestBody) -> Self {
Self {
id: 0,
control: ProtosocketControlCode::Normal.as_u8(),
body,
}
}
}
impl Message for Request {
fn message_id(&self) -> u64 {
self.id
}
fn control_code(&self) -> ProtosocketControlCode {
ProtosocketControlCode::from_u8(self.control)
}
fn set_message_id(&mut self, id: u64) {
self.id = id
}
fn cancelled(id: u64) -> Self {
Self {
id,
control: ProtosocketControlCode::Cancel.as_u8(),
body: RequestBody::Noop,
}
}
fn ended(id: u64) -> Self {
Self {
id,
control: ProtosocketControlCode::End.as_u8(),
body: RequestBody::Noop,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct WireServerInfo {
pub version: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ResponseBody {
Noop,
ServerInfo(WireServerInfo),
Span(WireSpan),
CacheLevel(WireLevelFilter),
CacheChance(f64),
Ack,
Error(String),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Response {
pub id: u64,
pub control: u8,
pub body: ResponseBody,
}
impl Response {
pub fn new(body: ResponseBody) -> Self {
Self {
id: 0,
control: ProtosocketControlCode::Normal.as_u8(),
body,
}
}
pub fn ack() -> Self {
Self::new(ResponseBody::Ack)
}
pub fn error(msg: impl Into<String>) -> Self {
Self::new(ResponseBody::Error(msg.into()))
}
pub fn span(s: WireSpan) -> Self {
Self::new(ResponseBody::Span(s))
}
pub fn cache_level(level: WireLevelFilter) -> Self {
Self::new(ResponseBody::CacheLevel(level))
}
pub fn cache_chance(pct: f64) -> Self {
Self::new(ResponseBody::CacheChance(pct))
}
pub fn server_info(version: impl Into<String>) -> Self {
Self::new(ResponseBody::ServerInfo(WireServerInfo {
version: version.into(),
}))
}
pub fn with_id(mut self, id: u64) -> Self {
self.id = id;
self
}
}
impl Message for Response {
fn message_id(&self) -> u64 {
self.id
}
fn control_code(&self) -> ProtosocketControlCode {
ProtosocketControlCode::from_u8(self.control)
}
fn set_message_id(&mut self, id: u64) {
self.id = id
}
fn cancelled(id: u64) -> Self {
Self {
id,
control: ProtosocketControlCode::Cancel.as_u8(),
body: ResponseBody::Noop,
}
}
fn ended(id: u64) -> Self {
Self {
id,
control: ProtosocketControlCode::End.as_u8(),
body: ResponseBody::Noop,
}
}
}