chore: optimized tracing

makefile

@@ -1,31 +1,35 @@
 GO=go
+CARGO=cargo
 
 GOFLAGS=-ldflags="-s -w"
 ifeq ($(shell uname -s),Darwin)
  GOFLAGS+=-buildmode=pie
 endif
 
-COMPONENTS=$(shell find components -type d -depth 1 -exec basename {} \;)
+GO_COMPONENTS=$(shell find components -type d -depth 1 -exec basename {} \;)
+RUST_COMPONENTS=$(shell find rust-components -type d -depth 1 -exec basename {} \;)
 
-build-%:
+build-go-%:
  $(GO) build $(GOFLAGS) -o bin/$* components/$*/cmd/main.go
 
+build-rust-%:
+ $(CARGO) build --release --manifest-path=rust-components/$*/Cargo.toml
+
 build:
- @for component in $(COMPONENTS); do \
- $(MAKE) build-$$component; \
+ @for component in $(GO_COMPONENTS); do \
+ $(MAKE) build-go-$$component; \
  done
 
-run-%:
- $(GO) run components/$*/cmd/main.go
+ @for component in $(RUST_COMPONENTS); do \
+ $(MAKE) build-rust-$$component; \
+ done
 
 test:
  $(GO) test -v ./...
-
-deps:
- $(GO) mod download
+ $(CARGO) test
 
 work-init:
  $(GO) work init
  $(GO) work use . components/*
 
-.PHONY: build-% build run-% test deps work-init
+.PHONY: build-go-% build-rust-% build test deps work-init

rust-common/src/bot/channel.rs

@@ -1,7 +1,9 @@
+use std::pin::Pin;
+
 use futures::{FutureExt, StreamExt};
 use lapin::options::BasicCancelOptions;
 use lapin::protocol::constants::REPLY_SUCCESS;
-use opentelemetry::trace::Span;
+use opentelemetry::trace::TraceContextExt;
 use teloxide::prelude::Update;
 use teloxide::stop::{mk_stop_token, StopFlag, StopToken};
 use teloxide::update_listeners::{AsUpdateStream, UpdateListener};
@@ -20,11 +22,11 @@ pub struct MqUpdateListener {
 impl<'a> AsUpdateStream<'a> for MqUpdateListener {
  type StreamErr = lapin::Error;
  type Stream =
- Box<dyn futures::Stream<Item = Result<Update, Self::StreamErr>> + Unpin + Send + 'a>;
+ Pin<Box<dyn futures::Stream<Item = Result<Update, Self::StreamErr>> + Unpin + Send + 'a>>;
 
  fn as_stream(&'a mut self) -> Self::Stream {
  let flag = self.flag.clone();
- let stream = self.consumer.clone().filter_map(move |delivery| {
+ Box::pin(self.consumer.clone().filter_map(move |delivery| {
  assert!(!flag.is_stopped(), "Update listener stopped");
  if self.consumer.state() != lapin::ConsumerState::Active
  && self.consumer.state() != lapin::ConsumerState::ActiveWithDelegate
@@ -34,27 +36,29 @@ impl<'a> AsUpdateStream<'a> for MqUpdateListener {
 
  async move {
  match delivery {
- Ok(delivery) => match serde_json::from_slice::<Update>(&delivery.data) {
- Ok(mut update) => {
- let cx = extract_span_from_delivery(&delivery);
- update.cx = Some(cx);
- Some(Ok(update))
- }
- Err(e) => {
- log::error!("Error deserializing message: {}", e);
- None
+ Ok(delivery) => {
+ let cx = extract_span_from_delivery(&delivery);
+
+ match serde_json::from_slice::<Update>(&delivery.data) {
+ Ok(mut update) => {
+ update.cx = Some(cx);
+ Some(Ok(update))
+ }
+ Err(e) => {
+ log::error!("Error deserializing message: {}", e);
+ cx.span().record_error(&e);
+ None
+ }
  }
- },
+ }
  Err(e) => {
  log::error!("Error receiving message: {}", e);
  None
  }
  }
  }
  .boxed()
- });
-
- Box::new(stream)
+ }))
  }
 }
 

rust-components/network-functions-handler/src/handlers.rs

@@ -1,11 +1,11 @@
 use std::time::Duration;
 
 use fast_qr::convert::Builder;
-use fastping_rs::PingResult::{Idle, Receive};
 use fastping_rs::Pinger;
+use fastping_rs::PingResult::{Idle, Receive};
 use moka::future::Cache;
-use opentelemetry::trace::{Span, SpanKind, Tracer};
 use opentelemetry::{global, KeyValue};
+use opentelemetry::trace::{SpanKind, TraceContextExt, Tracer};
 use teloxide::prelude::*;
 use teloxide::types::InputFile;
 use teloxide::utils::command::BotCommands;
@@ -50,15 +50,17 @@ pub(crate) async fn qrcode_handler(
  cache: Cache<String, Vec<u8>>,
 ) -> anyhow::Result<()> {
  let tracer = global::tracer("pegasus/rust-components/network-functions-handler/handlers");
- let mut span = tracer
+ let parent_cx = update.cx.unwrap_or_default();
+ let span = tracer
  .span_builder("qrcode_handler")
  .with_kind(SpanKind::Internal)
- .start_with_context(&tracer, &update.cx.unwrap_or_default());
+ .start_with_context(&tracer, &parent_cx);
+ let cx = parent_cx.with_span(span);
 
  if text.is_empty() {
  let err = anyhow::anyhow!("Text is empty");
  send_error_message!(bot, message, "Text is empty");
- span.record_error(err.as_ref());
+ cx.span().record_error(err.as_ref());
  return Err(err);
  }
 
@@ -74,7 +76,7 @@ pub(crate) async fn qrcode_handler(
  .send()
  .await?;
 
- span.set_attribute(KeyValue::new("cache_hit", true));
+ cx.span().set_attribute(KeyValue::new("cache_hit", true));
 
  return Ok(());
  }
@@ -119,20 +121,22 @@ pub(crate) async fn ping_handler(
  target: String,
 ) -> anyhow::Result<()> {
  let tracer = global::tracer("pegasus/rust-components/network-functions-handler/handlers");
- let mut span = tracer
+ let parent_tx = update.cx.unwrap_or_default();
+ let span = tracer
  .span_builder("ping_handler")
  .with_kind(SpanKind::Internal)
- .start_with_context(&tracer, &update.cx.unwrap_or_default());
+ .start_with_context(&tracer, &parent_tx);
+ let cx = parent_tx.with_span(span);
 
  if target.is_empty() {
  let err = anyhow::anyhow!("Target is empty");
  send_error_message!(bot, message, "Usage: /ping <target>");
- span.record_error(err.as_ref());
+ cx.span().record_error(err.as_ref());
  return Err(err);
  }
 
  let target_ip = match_error!(
- parse_target(&target).await,
+ parse_target(cx.clone(), &target).await,
  bot,
  message,
  "Failed to parse target: {}"
@@ -153,7 +157,8 @@ pub(crate) async fn ping_handler(
  Idle { addr } => {
  let err = format!("Failed to ping target: {}", addr);
  send_error_message!(bot, message, &err);
- span.record_error(anyhow::anyhow!("Failed to ping target: {}", addr).as_ref());
+ cx.span()
+ .record_error(anyhow::anyhow!("Failed to ping target: {}", addr).as_ref());
  return Err(anyhow::anyhow!("Failed to ping target: {}", err));
  }
  Receive { addr, rtt } => {
@@ -169,7 +174,8 @@ pub(crate) async fn ping_handler(
  Err(e) => {
  let err = format!("Failed to receive result: {}", e);
  send_error_message!(bot, message, err.clone());
- span.record_error(anyhow::anyhow!(err.clone()).as_ref());
+ cx.span()
+ .record_error(anyhow::anyhow!(err.clone()).as_ref());
  return Err(anyhow::anyhow!(err));
  }
  }

rust-components/network-functions-handler/src/utils.rs

@@ -1,6 +1,8 @@
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 
 use lazy_static::lazy_static;
+use opentelemetry::{Context, global, KeyValue};
+use opentelemetry::trace::{TraceContextExt, Tracer};
 use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
 use trust_dns_resolver::TokioAsyncResolver;
 
@@ -9,7 +11,17 @@ lazy_static! {
  TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default());
 }
 
-pub(crate) async fn parse_target(target: &str) -> anyhow::Result<IpAddr> {
+pub(crate) async fn parse_target(parent_cx: Context, target: &str) -> anyhow::Result<IpAddr> {
+ let tracer = global::tracer("pegasus/rust-components/network-functions-handler/utils");
+ let cx = parent_cx.with_span(
+ tracer
+ .span_builder("parse_target")
+ .start_with_context(&tracer, &parent_cx),
+ );
+
+ cx.span()
+ .set_attribute(KeyValue::new("target", target.to_string()));
+
  if let Ok(ip_addr) = target.parse::<Ipv4Addr>() {
  Ok(IpAddr::V4(ip_addr))
  } else if let Ok(ip_addr) = target.parse::<Ipv6Addr>() {
@@ -39,15 +51,17 @@ mod tests {
  #[tokio::test]
  async fn test_parse_target() {
  assert_eq!(
- parse_target("127.0.0.1").await.unwrap(),
+ parse_target(Context::default(), "127.0.0.1").await.unwrap(),
  IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))
  );
 
  assert_eq!(
- parse_target("::1").await.unwrap(),
+ parse_target(Context::default(), "::1").await.unwrap(),
  IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1))
  );
 
- assert!(parse_target("example.com").await.is_ok());
+ assert!(parse_target(Context::default(), "example.com")
+ .await
+ .is_ok());
  }
 }