From c4468b5e41ccb9e94021fb3bb219f88f0dcc3e1f Mon Sep 17 00:00:00 2001
From: massimoalbarello <massimo00albarello@gmail.com>
Date: Tue, 23 Jan 2024 12:14:57 +0100
Subject: [PATCH] added canister id attribute to client session span

---
 Cargo.lock                                             | 2 +-
 src/ic-websocket-gateway/Cargo.toml                    | 2 +-
 src/ic-websocket-gateway/src/client_session_handler.rs | 7 ++++---
 src/ic-websocket-gateway/src/ws_listener.rs            | 2 +-
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6bc8433..29224b6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1343,7 +1343,7 @@ checksum = "576c539151d4769fb4d1a0c25c4108dd18facd04c5695b02cf2d226ab4e43aa5"
 
 [[package]]
 name = "ic_websocket_gateway"
-version = "1.3.0"
+version = "1.3.1"
 dependencies = [
  "async-trait",
  "candid",
diff --git a/src/ic-websocket-gateway/Cargo.toml b/src/ic-websocket-gateway/Cargo.toml
index ffde276..d0ec656 100644
--- a/src/ic-websocket-gateway/Cargo.toml
+++ b/src/ic-websocket-gateway/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ic_websocket_gateway"
-version = "1.3.0"
+version = "1.3.1"
 edition.workspace = true
 rust-version.workspace = true
 repository.workspace = true
diff --git a/src/ic-websocket-gateway/src/client_session_handler.rs b/src/ic-websocket-gateway/src/client_session_handler.rs
index 6d1b5c9..e808f8e 100644
--- a/src/ic-websocket-gateway/src/client_session_handler.rs
+++ b/src/ic-websocket-gateway/src/client_session_handler.rs
@@ -13,7 +13,7 @@ use tokio::{
     sync::mpsc::{self, Receiver, Sender},
 };
 use tokio_tungstenite::accept_async;
-use tracing::{debug, info, span, warn, Instrument, Level, Span};
+use tracing::{debug, field, info, span, warn, Instrument, Level, Span};
 
 /// Handler of a client IC WS session
 pub struct ClientSessionHandler {
@@ -63,8 +63,7 @@ impl ClientSessionHandler {
                     Receiver<IcWsCanisterMessage>,
                 ) = mpsc::channel(100);
 
-                let client_session_span =
-                    span!(parent: &Span::current(), Level::TRACE, "Client Session");
+                let client_session_span = span!(parent: &Span::current(), Level::TRACE, "Client Session", canister_id = field::Empty);
 
                 let client_session = ClientSession::init(
                     self.id,
@@ -142,6 +141,8 @@ impl ClientSessionHandler {
                             client_session_span.clone(),
                         );
 
+                    client_session_span.record("canister_id", canister_id.to_string());
+
                     // TODO: figure out if it is guaranteed that all threads see the updated state of the gateway
                     //       before relaying the message to the IC
                     client_session
diff --git a/src/ic-websocket-gateway/src/ws_listener.rs b/src/ic-websocket-gateway/src/ws_listener.rs
index 6a59287..629e443 100644
--- a/src/ic-websocket-gateway/src/ws_listener.rs
+++ b/src/ic-websocket-gateway/src/ws_listener.rs
@@ -145,7 +145,7 @@ impl WsListener {
         let accept_client_connection_span = span!(
             Level::DEBUG,
             "Accept Connection",
-            ?client_addr,
+            client_addr = ?client_addr.ip(),
             client_id = self.next_client_id
         );
         let client_id = self.next_client_id;