chore: add tracing

containers/docker-compose.yml

@@ -1,6 +1,7 @@
 version: "3.9"
 services:
   mqtt:
-    build: ./mosquitto/
+    image: docker.io/emqx/emqx
     ports:
       - "1883:1883"
+      - "18083:18083"

crates/mqtt-client/src/client.rs

@@ -2,7 +2,7 @@ pub(crate) mod command;
 pub(crate) mod inner;
 pub(crate) mod subscription;
 
-use self::{inner::InnerClient, subscription::SubscriptionToken};
+use self::subscription::SubscriptionToken;
 use crate::{entity::EntityTopicBuilder, HassMqttOptions};
 use futures::Stream;
 use hass_dyn_error::DynError;
@@ -14,12 +14,14 @@ use std::{
 	task::{Context, Poll},
 };
 use thiserror::Error;
+use tracing::{field, instrument, span, Level, Span};
 
 #[derive(Clone)]
 pub struct Message {
 	pub topic: Arc<str>,
 	pub payload: Arc<[u8]>,
 	pub retained: bool,
+	pub span: Span,
 }
 
 impl Message {
@@ -34,6 +36,10 @@ impl Message {
 	pub fn retained(&self) -> bool {
 		self.retained
 	}
+
+	pub fn span(&self) -> &Span {
+		&self.span
+	}
 }
 
 #[derive(Clone)]
@@ -56,6 +62,7 @@ impl Stream for Subscription {
 
 #[derive(Clone)]
 pub struct HassMqttClient {
+	client_id: Arc<str>,
 	sender: flume::Sender<command::Command>,
 }
 
@@ -65,8 +72,9 @@ impl HassMqttClient {
 		T: command::ClientCommand,
 		command::Command: command::FromClientCommand<T>,
 	{
+		let span = Span::current();
 		let cmd = Arc::new(cmd);
-		let (msg, receiver) = command::Command::from_command(cmd.clone());
+		let (msg, receiver) = command::Command::from_command(cmd.clone(), span);
 		self
 			.sender
 			.send_async(msg)
@@ -96,11 +104,20 @@ impl ConnectError {
 }
 
 impl HassMqttClient {
+	#[instrument(
+		level = Level::DEBUG,
+		name = "HassMqttClient::new",
+		skip_all,
+		fields(
+			provider.name = T::NAME,
+		)
+		err,
+	)]
 	pub async fn new<T: MqttProvider>(options: HassMqttOptions) -> Result<Self, ConnectError> {
-		let sender = InnerClient::spawn::<T>(options)
+		let (sender, client_id) = inner::spawn::<T>(options)
 			.await
 			.map_err(ConnectError::new)?;
-		Ok(Self { sender })
+		Ok(Self { sender, client_id })
 	}
 }
 
@@ -110,7 +127,12 @@ impl HassMqttClient {
 		domain: impl Into<Arc<str>>,
 		entity_id: impl Into<Arc<str>>,
 	) -> EntityTopicBuilder {
-		EntityTopicBuilder::new(self, domain, entity_id)
+		self._entity(domain.into(), entity_id.into())
+	}
+
+	fn _entity(&self, domain: Arc<str>, entity_id: Arc<str>) -> EntityTopicBuilder {
+		let span = span!(Level::DEBUG, "HassMqttClient::entity", client.id = %self.client_id, entity.domain = %domain, entity.id = %entity_id, entity.topic = field::Empty);
+		EntityTopicBuilder::new(self, domain, entity_id, span)
 	}
 }
 
@@ -125,16 +147,24 @@ pub struct PublishMessageError {
 }
 
 impl HassMqttClient {
+	#[instrument(
+		level = Level::DEBUG,
+		name = "HassMqttClient::publish_message",
+		skip_all,
+		fields(
+			client.id = %self.client_id,
+			message.topic = %topic,
+			message.retained = retained,
+			message.qos = %qos,
+			message.payload.len = payload.len(),
+		))]
 	pub(crate) async fn publish_message(
 		&self,
-		topic: impl Into<Arc<str>>,
-		payload: impl Into<Arc<[u8]>>,
+		topic: Arc<str>,
+		payload: Arc<[u8]>,
 		retained: bool,
 		qos: QosLevel,
 	) -> Result<(), PublishMessageError> {
-		let topic = topic.into();
-		let payload = payload.into();
-
 		self
 			.command(command::publish(topic.clone(), payload, retained, qos))
 			.await
@@ -159,12 +189,20 @@ pub struct SubscribeError {
 }
 
 impl HassMqttClient {
+	#[instrument(
+		level = Level::DEBUG,
+		name = "HassMqttClient::subscribe",
+		skip_all,
+		fields(
+			client.id = %self.client_id,
+			subscription.topic = %topic,
+			subscription.qos,
+		))]
 	pub(crate) async fn subscribe(
 		&self,
-		topic: impl Into<Arc<str>>,
+		topic: Arc<str>,
 		qos: QosLevel,
 	) -> Result<Subscription, SubscribeError> {
-		let topic = topic.into();
 		let result = self
 			.command(command::subscribe(topic.clone(), qos))
 			.await

crates/mqtt-client/src/client/command.rs

@@ -7,6 +7,7 @@ use async_trait::async_trait;
 use hass_mqtt_provider::MqttClient;
 use std::sync::Arc;
 use tokio::sync::oneshot;
+use tracing::{Instrument, Span};
 
 pub(super) use entity::EntityCommand;
 pub(super) use publish::PublishCommand;
@@ -19,8 +20,7 @@ pub(crate) trait ClientCommand {
 
 	async fn run<T: MqttClient>(
 		&self,
-		client: &mut InnerClient,
-		mqtt: &T,
+		client: &mut InnerClient<T>,
 	) -> Result<Self::Result, Self::Error>;
 
 	fn create_error(&self, source: impl std::error::Error + Send + Sync + 'static) -> Self::Error;
@@ -32,7 +32,7 @@ pub(crate) type CommandResultSender<T> = oneshot::Sender<CommandResult<T>>;
 pub(crate) type CommandResultReceiver<T> = oneshot::Receiver<CommandResult<T>>;
 
 pub(crate) trait FromClientCommand<T: ClientCommand>: Sized {
-	fn from_command(command: Arc<T>) -> (Self, CommandResultReceiver<T>);
+	fn from_command(command: Arc<T>, span: Span) -> (Self, CommandResultReceiver<T>);
 }
 
 macro_rules! commands {
@@ -41,38 +41,36 @@ macro_rules! commands {
 	}) => {
 		#[allow(clippy::enum_variant_names)]
 		$vis enum $name {
-			$($variant(Arc<$variant>, CommandResultSender<$variant>),)*
+			$($variant(Arc<$variant>, CommandResultSender<$variant>, Span),)*
 		}
 
 		$(
 			impl FromClientCommand<$variant> for $name {
-				fn from_command(command: Arc<$variant>) -> (Self, CommandResultReceiver<$variant>) {
+				fn from_command(command: Arc<$variant>, span: Span) -> (Self, CommandResultReceiver<$variant>) {
 					let (tx, rx) = oneshot::channel();
 
-					(Self::$variant(command, tx), rx)
+					(Self::$variant(command, tx, span), rx)
 				}
 			}
 		)*
 
 		impl $name {
-			pub(super) fn from_command<T>(command: Arc<T>) -> (Self, CommandResultReceiver<T>)
+			pub(super) fn from_command<T>(command: Arc<T>, span: Span) -> (Self, CommandResultReceiver<T>)
 			where
 				T: ClientCommand,
 				Self: FromClientCommand<T>,
 			{
-				<Self as FromClientCommand<T>>::from_command(command)
+				<Self as FromClientCommand<T>>::from_command(command, span)
 			}
 
 			pub(super) async fn run<T: MqttClient>(
 				self,
-				client: &mut InnerClient,
-				mqtt: &T,
+				client: &mut InnerClient<T>,
 			) {
 				match self {
 					$(
-						Self::$variant(command, tx) => {
-							// TODO: tracing
-							let result = command.run(client, mqtt).await;
+						Self::$variant(command, tx, span) => {
+							let result = command.run(client).instrument(span).await;
 
 							if let Err(_err) = tx.send(result) {
 								// TODO: log

crates/mqtt-client/src/client/command/entity.rs

@@ -42,8 +42,7 @@ impl ClientCommand for EntityCommand {
 
 	async fn run<T: MqttClient>(
 		&self,
-		client: &mut InnerClient,
-		_mqtt: &T,
+		client: &mut InnerClient<T>,
 	) -> Result<Self::Result, Self::Error> {
 		let topics_config = client
 			.topics

crates/mqtt-client/src/client/command/publish.rs

@@ -2,7 +2,7 @@ use super::{ClientCommand, InnerClient};
 use crate::client::QosLevel;
 use async_trait::async_trait;
 use hass_dyn_error::DynError;
-use hass_mqtt_provider::{MqttClient, MqttMessage, MqttMessageBuilder};
+use hass_mqtt_provider::{MqttBuildableMessage, MqttClient, MqttMessageBuilder};
 use std::sync::Arc;
 use thiserror::Error;
 
@@ -41,18 +41,18 @@ impl ClientCommand for PublishCommand {
 
 	async fn run<T: MqttClient>(
 		&self,
-		_client: &mut InnerClient,
-		mqtt: &T,
+		client: &mut InnerClient<T>,
 	) -> Result<Self::Result, Self::Error> {
-		let msg = <T::Message as MqttMessage>::builder()
+		let msg = <T::Message as MqttBuildableMessage>::builder()
 			.topic(&*self.topic)
 			.payload(&*self.payload)
 			.retain(self.retained)
 			.qos(self.qos)
 			.build()
 			.map_err(|source| self.create_error(source))?;
 
-		mqtt
+		client
+			.client
 			.publish(msg)
 			.await
 			.map_err(|source| self.create_error(source))

crates/mqtt-client/src/client/command/subscribe.rs

@@ -1,5 +1,8 @@
 use super::{ClientCommand, InnerClient};
-use crate::client::{subscription::SubscriptionToken, Message, QosLevel};
+use crate::{
+	client::{subscription::SubscriptionToken, Message, QosLevel},
+	router::RouterEntry,
+};
 use async_trait::async_trait;
 use hass_dyn_error::DynError;
 use hass_mqtt_provider::MqttClient;
@@ -38,20 +41,23 @@ impl ClientCommand for SubscribeCommand {
 
 	async fn run<T: MqttClient>(
 		&self,
-		client: &mut InnerClient,
-		mqtt: &T,
+		client: &mut InnerClient<T>,
 	) -> Result<Self::Result, Self::Error> {
 		let (sender, receiver) = flume::unbounded();
-		let route_id = client.router.insert(&self.topic, sender);
-		let token = client.subscriptions.insert(route_id);
+		let route_id = match client.router.entry(self.topic.clone()) {
+			RouterEntry::Occupied(entry) => entry.insert(sender),
+			RouterEntry::Vacant(entry) => {
+				let key = client
+					.client
+					.subscribe(self.topic.clone(), self.qos)
+					.await
+					.map_err(|source| self.create_error(source))?;
 
-		// Note: if the subscription fails, the token gets dropped,
-		// which in turn will clean up the route in the router.
-		mqtt
-			.subscribe(&*self.topic, self.qos)
-			.await
-			.map_err(|source| self.create_error(source))?;
+				entry.insert(key, sender)
+			}
+		};
 
+		let token = client.subscriptions.insert(route_id);
 		Ok(SubscribeCommandResult { token, receiver })
 	}