perf(runtime): use parking_lot for TCP stream registry (#11065)

Signed-off-by: jthomson04 <jwillthomson19@gmail.com>

Author: jthomson04

lib/runtime/src/pipeline/network/tcp/server.rs

@@ -9,7 +9,6 @@ use std::{
     sync::Arc,
     time::Duration,
 };
-use tokio::sync::Mutex;
 use tokio::time::Instant;
 
 /// Tombstone lifetime. Bridges the `register()` → `associate_instance()`
@@ -22,6 +21,7 @@ use bytes::Bytes;
 use derive_builder::Builder;
 use futures::{SinkExt, StreamExt};
 use local_ip_address::{Error, list_afinet_netifas, local_ip, local_ipv6};
+use parking_lot::Mutex;
 
 use serde::{Deserialize, Serialize};
 use tokio::{
@@ -259,7 +259,7 @@ impl TcpStreamServer {
         send_subject: Option<&str>,
         id: &EndpointInstanceId,
     ) -> bool {
-        let mut state = self.state.lock().await;
+        let mut state = self.state.lock();
         let now = Instant::now();
         prune_tombstones(&mut state.removed_instances, now);
         if state.removed_instances.contains_key(id) {
@@ -296,7 +296,7 @@ impl TcpStreamServer {
     /// Cancel one pending response-stream registration. Drops the
     /// `oneshot::Sender` so the waiting receiver resolves with `RecvError`.
     pub async fn cancel_recv_stream(&self, subject: &str) {
-        let mut state = self.state.lock().await;
+        let mut state = self.state.lock();
         state.rx_subjects.remove(subject);
         if let Some(key) = state.subject_instance.remove(subject)
             && let Some(subjects) = state.instance_subjects.get_mut(&key)
@@ -314,7 +314,7 @@ impl TcpStreamServer {
     /// `(StreamType::Request, _)` tag from `instance_subjects` so the per-
     /// instance bookkeeping stays consistent.
     pub async fn cancel_send_stream(&self, subject: &str) {
-        let mut state = self.state.lock().await;
+        let mut state = self.state.lock();
         state.tx_subjects.remove(subject);
         if let Some(key) = state.subject_instance.remove(subject)
             && let Some(subjects) = state.instance_subjects.get_mut(&key)
@@ -331,7 +331,7 @@ impl TcpStreamServer {
     /// `associate_instance` — and tombstone the id so any racing associate
     /// for the same id cancels too. Returns the number of streams cancelled.
     pub async fn cancel_instance_streams(&self, id: &EndpointInstanceId) -> usize {
-        let mut state = self.state.lock().await;
+        let mut state = self.state.lock();
         let now = Instant::now();
         prune_tombstones(&mut state.removed_instances, now);
         state.removed_instances.insert(id.clone(), now);
@@ -357,25 +357,63 @@ impl TcpStreamServer {
     /// Drop the tombstone for an instance that has reappeared in discovery,
     /// so future subjects for that identity are tracked normally.
     pub async fn clear_instance_tombstone(&self, id: &EndpointInstanceId) {
-        let mut state = self.state.lock().await;
+        let mut state = self.state.lock();
         state.removed_instances.remove(id);
     }
 
-    #[allow(clippy::await_holding_lock)]
     async fn start(local_ip: String, local_port: u16, state: Arc<Mutex<State>>) -> Result<u16> {
         let addr = format!("{}:{}", local_ip, local_port);
         let state_clone = state.clone();
-        let mut guard = state.lock().await;
-        if guard.handle.is_some() {
-            panic!("TcpStreamServer already started");
-        }
         let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<Result<u16>>();
-        let handle = tokio::spawn(tcp_listener(addr, state_clone, ready_tx));
-        guard.handle = Some(handle);
-        drop(guard);
+        {
+            let mut guard = state.lock();
+            if guard.handle.is_some() {
+                panic!("TcpStreamServer already started");
+            }
+            guard.handle = Some(tokio::spawn(tcp_listener(addr, state_clone, ready_tx)));
+        }
         let local_port = ready_rx.await??;
         Ok(local_port)
     }
+
+    fn insert_request_stream(&self, subject: String, connection: RequestedSendConnection) {
+        self.state.lock().tx_subjects.insert(subject, connection);
+    }
+
+    fn insert_response_stream(&self, subject: String, connection: RequestedRecvConnection) {
+        self.state.lock().rx_subjects.insert(subject, connection);
+    }
+
+    fn take_request_stream(state: &Mutex<State>, subject: &str) -> Option<RequestedSendConnection> {
+        let mut state = state.lock();
+        let connection = state.tx_subjects.remove(subject);
+        if let Some(key) = state.subject_instance.remove(subject)
+            && let Some(subjects) = state.instance_subjects.get_mut(&key)
+        {
+            subjects.remove(&(StreamType::Request, subject.to_string()));
+            if subjects.is_empty() {
+                state.instance_subjects.remove(&key);
+            }
+        }
+        connection
+    }
+
+    fn take_response_stream(
+        state: &Mutex<State>,
+        subject: &str,
+    ) -> Option<RequestedRecvConnection> {
+        let mut state = state.lock();
+        let connection = state.rx_subjects.remove(subject);
+        if let Some(key) = state.subject_instance.remove(subject)
+            && let Some(subjects) = state.instance_subjects.get_mut(&key)
+        {
+            subjects.remove(&(StreamType::Response, subject.to_string()));
+            if subjects.is_empty() {
+                state.instance_subjects.remove(&key);
+            }
+        }
+        connection
+    }
 }
 
 // todo - possible rename ResponseService to ResponseServer
@@ -409,6 +447,7 @@ impl ResponseService for TcpStreamServer {
 
         let send_stream = if options.enable_request_stream {
             let sender_subject = uuid::Uuid::new_v4().to_string();
+            let registry_subject = sender_subject.clone();
 
             let (pending_sender_tx, pending_sender_rx) = oneshot::channel();
 
@@ -418,11 +457,6 @@ impl ResponseService for TcpStreamServer {
                 send_buffer_count: options.send_buffer_count,
             };
 
-            let mut state = self.state.lock().await;
-            state
-                .tx_subjects
-                .insert(sender_subject.clone(), connection_info);
-
             let cleanup_subject = sender_subject.clone();
             let cleanup_state = self.state.clone();
             let registered_stream = RegisteredStream::new(
@@ -438,7 +472,7 @@ impl ResponseService for TcpStreamServer {
             .with_cleanup(move || {
                 // Drop is sync; fire-and-forget the lock acquisition.
                 tokio::spawn(async move {
-                    let mut state = cleanup_state.lock().await;
+                    let mut state = cleanup_state.lock();
                     state.tx_subjects.remove(&cleanup_subject);
                     if let Some(key) = state.subject_instance.remove(&cleanup_subject)
                         && let Some(subjects) = state.instance_subjects.get_mut(&key)
@@ -451,6 +485,8 @@ impl ResponseService for TcpStreamServer {
                 });
             });
 
+            self.insert_request_stream(registry_subject, connection_info);
+
             Some(registered_stream)
         } else {
             None
@@ -459,18 +495,14 @@ impl ResponseService for TcpStreamServer {
         let recv_stream = if options.enable_response_stream {
             let (pending_recver_tx, pending_recver_rx) = oneshot::channel();
             let receiver_subject = uuid::Uuid::new_v4().to_string();
+            let registry_subject = receiver_subject.clone();
 
             let connection_info = RequestedRecvConnection {
                 context: options.context.clone(),
                 connection: pending_recver_tx,
                 send_buffer_count: options.send_buffer_count,
             };
 
-            let mut state = self.state.lock().await;
-            state
-                .rx_subjects
-                .insert(receiver_subject.clone(), connection_info);
-
             let cleanup_subject = receiver_subject.clone();
             let cleanup_state = self.state.clone();
             let registered_stream = RegisteredStream::new(
@@ -486,7 +518,7 @@ impl ResponseService for TcpStreamServer {
             .with_cleanup(move || {
                 // Drop is sync; fire-and-forget the lock acquisition.
                 tokio::spawn(async move {
-                    let mut state = cleanup_state.lock().await;
+                    let mut state = cleanup_state.lock();
                     state.rx_subjects.remove(&cleanup_subject);
                     if let Some(key) = state.subject_instance.remove(&cleanup_subject)
                         && let Some(subjects) = state.instance_subjects.get_mut(&key)
@@ -499,6 +531,8 @@ impl ResponseService for TcpStreamServer {
                 });
             });
 
+            self.insert_response_stream(registry_subject, connection_info);
+
             Some(registered_stream)
         } else {
             None
@@ -653,22 +687,12 @@ async fn tcp_listener(
         // Request stream is unidirectional; we don't read from the downstream.
         drop(reader);
 
-        let request_stream = {
-            let mut guard = state.lock().await;
-            let conn = guard.tx_subjects.remove(&subject).ok_or(error!(
+        let request_stream = TcpStreamServer::take_request_stream(&state, &subject).ok_or_else(|| {
+            error!(
                 "Subject not found: {}; downstream subscriber specified a subject unknown to the upstream publisher",
                 subject
-            ))?;
-            if let Some(key) = guard.subject_instance.remove(&subject)
-                && let Some(subjects) = guard.instance_subjects.get_mut(&key)
-            {
-                subjects.remove(&(StreamType::Request, subject.clone()));
-                if subjects.is_empty() {
-                    guard.instance_subjects.remove(&key);
-                }
-            }
-            conn
-        };
+            )
+        })?;
 
         let RequestedSendConnection {
             context,
@@ -771,22 +795,9 @@ async fn tcp_listener(
         mut reader: FramedRead<tokio::io::ReadHalf<tokio::net::TcpStream>, TwoPartCodec>,
         writer: FramedWrite<tokio::io::WriteHalf<tokio::net::TcpStream>, TwoPartCodec>,
     ) -> Result<()> {
-        let response_stream = {
-            let mut guard = state.lock().await;
-            let conn = guard
-                .rx_subjects
-                .remove(&subject)
-                .ok_or(error!("Subject not found: {}; upstream publisher specified a subject unknown to the downsteam subscriber", subject))?;
-            if let Some(key) = guard.subject_instance.remove(&subject)
-                && let Some(subjects) = guard.instance_subjects.get_mut(&key)
-            {
-                subjects.remove(&(StreamType::Response, subject.clone()));
-                if subjects.is_empty() {
-                    guard.instance_subjects.remove(&key);
-                }
-            }
-            conn
-        };
+        let response_stream = TcpStreamServer::take_response_stream(&state, &subject).ok_or_else(|| {
+            error!("Subject not found: {}; upstream publisher specified a subject unknown to the downsteam subscriber", subject)
+        })?;
 
         // unwrap response_stream
         let RequestedRecvConnection {
@@ -1036,6 +1047,7 @@ mod tests {
     use crate::engine::AsyncEngineContextProvider;
     use crate::pipeline::Context;
     use crate::pipeline::network::DEFAULT_SEND_BUFFER_COUNT;
+    use crate::pipeline::network::tcp::client::TcpClient;
     use tokio::io::{AsyncWriteExt, ReadHalf, WriteHalf};
     use tokio::net::TcpStream;
 
@@ -1138,7 +1150,7 @@ mod tests {
 
         let _pending = server.register(options).await;
 
-        let state = server.state.lock().await;
+        let state = server.state.lock();
         assert_eq!(state.tx_subjects.len(), 1, "one request stream registered");
         assert_eq!(state.rx_subjects.len(), 1, "one response stream registered");
         assert!(
@@ -1439,7 +1451,7 @@ mod tests {
 
         // Verify it's in rx_subjects
         {
-            let state = server.state.lock().await;
+            let state = server.state.lock();
             assert!(state.rx_subjects.contains_key(&subject));
         }
 
@@ -1451,7 +1463,7 @@ mod tests {
 
         // Verify it's been removed from rx_subjects
         {
-            let state = server.state.lock().await;
+            let state = server.state.lock();
             assert!(
                 !state.rx_subjects.contains_key(&subject),
                 "RAII cleanup should have removed the rx_subjects entry"
@@ -1486,7 +1498,7 @@ mod tests {
 
         // The entry should still be in rx_subjects (cleanup was disarmed)
         {
-            let state = server.state.lock().await;
+            let state = server.state.lock();
             assert!(
                 state.rx_subjects.contains_key(&subject),
                 "into_parts() should disarm the RAII cleanup"
@@ -1643,7 +1655,7 @@ mod tests {
         // Tombstone the identity.
         server.cancel_instance_streams(&id).await;
         {
-            let state = server.state.lock().await;
+            let state = server.state.lock();
             assert!(state.removed_instances.contains_key(&id));
         }
 
@@ -1661,7 +1673,7 @@ mod tests {
         // The expired tombstone must have been pruned (lazy pruning fires on
         // every associate_instance/cancel_instance_streams call).
         {
-            let state = server.state.lock().await;
+            let state = server.state.lock();
             assert!(
                 !state.removed_instances.contains_key(&id),
                 "expired tombstone should be pruned, not retained"
@@ -1702,7 +1714,7 @@ mod tests {
         tokio::time::advance(TOMBSTONE_TTL + Duration::from_secs(1)).await;
         server.cancel_instance_streams(&id_new).await;
 
-        let state = server.state.lock().await;
+        let state = server.state.lock();
         assert!(
             !state.removed_instances.contains_key(&id_old),
             "old tombstone should be pruned by the next cancel_instance_streams call"
@@ -1728,7 +1740,7 @@ mod tests {
         server.cancel_instance_streams(&id_a).await;
         server.clear_instance_tombstone(&id_b).await;
 
-        let state = server.state.lock().await;
+        let state = server.state.lock();
         assert!(
             state.removed_instances.contains_key(&id_a),
             "clearing a different identity must not remove id_a's tombstone"
@@ -1957,6 +1969,63 @@ mod tests {
         }
     }
 
+    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+    async fn test_concurrent_response_registration_and_call_home() {
+        const STREAMS: usize = 128;
+
+        let result = time::timeout(Duration::from_secs(20), async {
+            let server = test_server().await;
+            let mut pending_streams = Vec::with_capacity(STREAMS);
+            let mut client_tasks = Vec::with_capacity(STREAMS);
+
+            for idx in 0..STREAMS {
+                let context = Context::new(());
+                let options = StreamOptions::builder()
+                    .context(context.context())
+                    .enable_request_stream(false)
+                    .enable_response_stream(true)
+                    .build()
+                    .unwrap();
+
+                let pending = server.register(options).await;
+                let registered_stream = pending.recv_stream.unwrap();
+                let (connection_info, stream_provider) = registered_stream.into_parts();
+                let client_context =
+                    Context::with_id_and_metadata((), context.id().to_string(), Default::default());
+                let payload = Bytes::from(format!("payload-{idx}"));
+
+                pending_streams.push((idx, payload.clone(), stream_provider));
+                client_tasks.push(tokio::spawn(async move {
+                    let mut sender = TcpClient::create_response_stream(
+                        client_context.context(),
+                        connection_info,
+                        None,
+                    )
+                    .await
+                    .unwrap();
+                    sender.send_prologue(None).await.unwrap();
+                    sender.send(payload).await.unwrap();
+                }));
+            }
+
+            for task in client_tasks {
+                task.await.unwrap();
+            }
+
+            for (idx, expected, stream_provider) in pending_streams {
+                let mut stream = stream_provider.await.unwrap().unwrap();
+                let actual = stream.rx.recv().await.unwrap();
+                assert_eq!(actual, expected, "payload mismatch for stream {idx}");
+            }
+        })
+        .await;
+
+        assert!(
+            result.is_ok(),
+            "concurrent response registration and call-home timed out"
+        );
+    }
+
     // ==================== request_stream_send_handler integration tests ====================
     //
     // These exercise the closing-message contract of `request_stream_send_handler`