improve tracing

amqp_transport/src/connection.rs

@@ -9,7 +9,7 @@ use amqp_core::{
         Channel, ChannelInner, ChannelNum, ConEventReceiver, ConEventSender, Connection,
         ConnectionEvent, ConnectionId, ContentHeader,
     },
-    message::{MessageId, RawMessage, RoutingInformation},
+    message::{MessageId, MessageInner, RoutingInformation},
     methods::{
         BasicPublish, ChannelClose, ChannelCloseOk, ChannelOpenOk, ConnectionClose,
         ConnectionCloseOk, ConnectionOpen, ConnectionOpenOk, ConnectionStart, ConnectionStartOk,
@@ -198,9 +198,8 @@ impl TransportConnection {
         Ok(())
     }
 
+    #[tracing::instrument(skip(self), level = "trace")]
     async fn send_method(&mut self, channel: ChannelNum, method: &Method) -> Result<()> {
-        trace!(%channel, ?method, "Sending method");
-
         let mut payload = Vec::with_capacity(64);
         methods::write::write_method(method, &mut payload)?;
         frame::write_frame(&mut self.stream, FrameType::Method, channel, &payload).await
@@ -324,6 +323,7 @@ impl TransportConnection {
         }
     }
 
+    #[tracing::instrument(skip(self), level = "debug")]
     async fn handle_frame(&mut self, frame: Frame) -> Result<()> {
         let channel = frame.channel;
         self.reset_timeout();
@@ -358,9 +358,9 @@ impl TransportConnection {
         }
     }
 
+    #[tracing::instrument(skip(self, frame), level = "trace")]
     async fn dispatch_method(&mut self, frame: Frame) -> Result<()> {
         let method = methods::parse_method(&frame.payload)?;
-        debug!(?method, "Received method");
 
         // Sending a method implicitly cancels the content frames that might be ongoing
         self.channels
@@ -483,7 +483,7 @@ impl TransportConnection {
             ..
         }) = method
         {
-            let message = RawMessage {
+            let message = MessageInner {
                 id: MessageId::random(),
                 header,
                 routing: RoutingInformation {
@@ -498,12 +498,7 @@ impl TransportConnection {
 
             let channel = self.channels.get(&channel).ok_or(ConException::Todo)?;
 
-            // Spawn the handler for the publish. The connection task goes back to handling
-            // just the connection.
-            tokio::spawn(amqp_messaging::methods::handle_basic_publish(
-                channel.global_chan.clone(),
-                message,
-            ));
+            amqp_messaging::methods::handle_basic_publish(channel.global_chan.clone(), message);
             Ok(())
         } else {
             Err(ConException::Todo.into())

amqp_transport/src/frame.rs

@@ -252,6 +252,7 @@ impl Debug for MaxFrameSize {
     }
 }
 
+#[tracing::instrument(skip(w), level = "trace")]
 pub async fn write_frame<W>(
     mut w: W,
     kind: FrameType,
@@ -261,8 +262,6 @@ pub async fn write_frame<W>(
 where
     W: AsyncWriteExt + Unpin + Send,
 {
-    trace!(?kind, ?channel, ?payload, "Sending frame");
-
     w.write_u8(kind as u8).await?;
     w.write_u16(channel.num()).await?;
     w.write_u32(u32::try_from(payload.len()).context("frame size too big")?)

amqp_transport/src/lib.rs

@@ -11,46 +11,92 @@ mod tests;
 // TODO: handle big types
 
 use crate::connection::TransportConnection;
-use amqp_core::GlobalData;
-use tokio::net;
-use tracing::{info, trace_span, Instrument};
+use amqp_core::{connection::ConnectionEvent, queue::QueueEvent, GlobalData};
+use anyhow::Context;
+use std::{future::Future, net::SocketAddr};
+use tokio::{net, net::TcpStream, select};
+use tracing::{info, info_span, Instrument};
 
-pub async fn do_thing_i_guess(global_data: GlobalData) -> anyhow::Result<()> {
+pub async fn do_thing_i_guess(
+    global_data: GlobalData,
+    terminate: impl Future + Send,
+) -> anyhow::Result<()> {
+    select! {
+        res = accept_cons(global_data.clone()) => {
+            res
+        }
+        _ = terminate => {
+            handle_shutdown(global_data).await
+        }
+    }
+}
+
+async fn accept_cons(global_data: GlobalData) -> anyhow::Result<()> {
     info!("Binding TCP listener...");
     let listener = net::TcpListener::bind(("127.0.0.1", 5672)).await?;
     info!(addr = ?listener.local_addr()?, "Successfully bound TCP listener");
 
     loop {
-        let (stream, peer_addr) = listener.accept().await?;
-
-        let id = rand::random();
-
-        info!(local_addr = ?stream.local_addr(), %id, "Accepted new connection");
-        let span = trace_span!("client-connection", %id);
-
-        let (method_send, method_recv) = tokio::sync::mpsc::channel(10);
-
-        let connection_handle = amqp_core::connection::ConnectionInner::new(
-            id,
-            peer_addr,
-            global_data.clone(),
-            method_send.clone(),
-        );
-
-        let mut global_data_guard = global_data.lock();
-        global_data_guard
-            .connections
-            .insert(id, connection_handle.clone());
-
-        let connection = TransportConnection::new(
-            id,
-            stream,
-            connection_handle,
-            global_data.clone(),
-            method_send,
-            method_recv,
-        );
-
-        tokio::spawn(connection.start_connection_processing().instrument(span));
+        let connection = listener.accept().await?;
+        handle_con(global_data.clone(), connection);
     }
 }
+
+fn handle_con(global_data: GlobalData, connection: (TcpStream, SocketAddr)) {
+    let (stream, peer_addr) = connection;
+    let id = rand::random();
+
+    info!(local_addr = ?stream.local_addr(), %id, "Accepted new connection");
+    let span = info_span!("client-connection", %id);
+
+    let (method_send, method_recv) = tokio::sync::mpsc::channel(10);
+
+    let connection_handle = amqp_core::connection::ConnectionInner::new(
+        id,
+        peer_addr,
+        global_data.clone(),
+        method_send.clone(),
+    );
+
+    let mut global_data_guard = global_data.lock();
+    global_data_guard
+        .connections
+        .insert(id, connection_handle.clone());
+
+    let connection = TransportConnection::new(
+        id,
+        stream,
+        connection_handle,
+        global_data.clone(),
+        method_send,
+        method_recv,
+    );
+
+    tokio::spawn(connection.start_connection_processing().instrument(span));
+}
+
+async fn handle_shutdown(global_data: GlobalData) -> anyhow::Result<()> {
+    info!("Shutting down...");
+
+    let lock = global_data.lock();
+
+    for con in lock.connections.values() {
+        con.event_sender
+            .try_send(ConnectionEvent::Shutdown)
+            .context("failed to stop connection")?;
+    }
+
+    for queue in lock.queues.values() {
+        queue
+            .event_send
+            .try_send(QueueEvent::Shutdown)
+            .context("failed to stop queue worker")?;
+    }
+
+    // todo: here we should wait for everything to close
+    // https://github.com/tokio-rs/mini-redis/blob/4b4ecf0310e6bca43d336dde90a06d9dcad00d6c/src/server.rs#L51
+
+    info!("Finished shutdown");
+
+    Ok(())
+}