From 24d04f2b1f5586b33103e3916bf1cdebda9d6a1f Mon Sep 17 00:00:00 2001
From: _ <>
Date: Sat, 17 Jul 2021 21:46:17 +0000
Subject: [PATCH] :recycle: refactor: move Mutex inside RelayState

---
 Cargo.lock                                    |  7 +++++++
 prototypes/quic_demo/Cargo.toml               |  1 +
 .../src/bin/quic_demo_relay_server.rs         | 20 +++++++++----------
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1ce5b35..fde7a14 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -319,6 +319,12 @@ version = "0.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10"
 
+[[package]]
+name = "drop_guard"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78bb5d73478dbcb7c404cd86336e90c67425497ca94b6f7352c8ea7deb9098e2"
+
 [[package]]
 name = "either"
 version = "1.6.1"
@@ -1351,6 +1357,7 @@ name = "quic_demo"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "drop_guard",
  "futures-util",
  "quinn",
  "rcgen",
diff --git a/prototypes/quic_demo/Cargo.toml b/prototypes/quic_demo/Cargo.toml
index f7d4d10..0b2f25c 100644
--- a/prototypes/quic_demo/Cargo.toml
+++ b/prototypes/quic_demo/Cargo.toml
@@ -8,6 +8,7 @@ edition = "2018"
 
 [dependencies]
 anyhow = "1.0.38"
+drop_guard = "0.2.1"
 futures-util = "0.3.9"
 quinn = "0.7.2"
 rcgen = "0.8.11"
diff --git a/prototypes/quic_demo/src/bin/quic_demo_relay_server.rs b/prototypes/quic_demo/src/bin/quic_demo_relay_server.rs
index b6a76d5..beb8fb4 100644
--- a/prototypes/quic_demo/src/bin/quic_demo_relay_server.rs
+++ b/prototypes/quic_demo/src/bin/quic_demo_relay_server.rs
@@ -9,7 +9,7 @@ async fn main () -> anyhow::Result <()> {
 	tokio::fs::write ("quic_server.crt", &server_cert).await?;
 	
 	let relay_state = RelayState::default ();
-	let relay_state = Arc::new (Mutex::new (relay_state));
+	let relay_state = Arc::new (relay_state);
 	
 	while let Some (conn) = incoming.next ().await {
 		let relay_state = Arc::clone (&relay_state);
@@ -25,7 +25,7 @@ async fn main () -> anyhow::Result <()> {
 
 #[derive (Default)]
 struct RelayState {
-	p4_server_proxies: HashMap <u8, P4State>,
+	p4_server_proxies: Mutex <HashMap <u8, P4State>>,
 }
 
 struct P4State {
@@ -104,7 +104,7 @@ impl PtthNewConnection {
 }
 
 async fn handle_quic_connection (
-	relay_state: Arc <Mutex <RelayState>>,
+	relay_state: Arc <RelayState>,
 	conn: quinn::Connecting,
 ) -> anyhow::Result <()> 
 {
@@ -146,7 +146,7 @@ async fn handle_quic_connection (
 }
 
 async fn handle_p2_connection (
-	relay_state: Arc <Mutex <RelayState>>,
+	relay_state: Arc <RelayState>,
 	client_id: u8,
 	conn: quinn::NewConnection,
 ) -> anyhow::Result <()>
@@ -186,7 +186,7 @@ async fn handle_p2_connection (
 }
 
 async fn handle_request_p2_to_p4 (
-	relay_state: Arc <Mutex <RelayState>>,
+	relay_state: Arc <RelayState>,
 	client_id: u8,
 	server_id: u8,
 	mut client_send: quinn::SendStream,
@@ -201,8 +201,8 @@ async fn handle_request_p2_to_p4 (
 	client_send.write_all (&resp_buf).await?;
 	
 	{
-		let relay_state = relay_state.lock ().await;
-		match relay_state.p4_server_proxies.get (&server_id) {
+		let p4_server_proxies = relay_state.p4_server_proxies.lock ().await;
+		match p4_server_proxies.get (&server_id) {
 			Some (p4_state) => {
 				p4_state.req_channel.send (RequestP2ToP4 {
 					client_send,
@@ -218,7 +218,7 @@ async fn handle_request_p2_to_p4 (
 }
 
 async fn handle_p4_connection (
-	relay_state: Arc <Mutex <RelayState>>,
+	relay_state: Arc <RelayState>,
 	server_id: u8,
 	conn: quinn::NewConnection,
 ) -> anyhow::Result <()>
@@ -234,8 +234,8 @@ async fn handle_p4_connection (
 	};
 	
 	{
-		let mut relay_state = relay_state.lock ().await;
-		relay_state.p4_server_proxies.insert (server_id, p4_state);
+		let mut p4_server_proxies = relay_state.p4_server_proxies.lock ().await;
+		p4_server_proxies.insert (server_id, p4_state);
 	}
 	
 	while let Some (req) = rx.recv ().await {