some cleanup

googleforgames · Aug 24, 2023 · a1ac3c5 · a1ac3c5
1 parent 2af4534
commit a1ac3c5
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 46 deletions.
diff --git a/src/endpoint/locality.rs b/src/endpoint/locality.rs
@@ -306,11 +306,6 @@ impl LocalitySet {
                 Entry::Occupied(entry) => {
                     let (key, original_locality) = entry.remove_entry();
                     let mut value = value.clone();
-                    let new_set_addresses = value
-                        .endpoints
-                        .iter()
-                        .map(|endpoint| endpoint.address.clone())
-                        .collect::<BTreeSet<_>>();
 
                     if tracing::enabled!(tracing::Level::INFO) {
                         for endpoint in value.endpoints.iter() {
@@ -326,8 +321,7 @@ impl LocalitySet {
                         .endpoints
                         .into_iter()
                         .partition(|endpoint| {
-                            !new_set_addresses.contains(&endpoint.address)
-                                && endpoint.sessions.load(std::sync::atomic::Ordering::SeqCst) != 0
+                            endpoint.sessions.load(std::sync::atomic::Ordering::SeqCst) != 0
                         });
 
                     if tracing::enabled!(tracing::Level::INFO) {

diff --git a/src/maxmind_db.rs b/src/maxmind_db.rs
@@ -68,19 +68,7 @@ impl MaxmindDb {
         };
 
         match mmdb.lookup::<IpNetEntry>(ip) {
-            Ok(asn) => {
-                tracing::info!(
-                    number = asn.r#as,
-                    organization = asn.as_name,
-                    country_code = asn.as_cc,
-                    prefix = asn.prefix,
-                    prefix_entity = asn.prefix_entity,
-                    prefix_name = asn.prefix_name,
-                    "maxmind information"
-                );
-
-                Some(asn)
-            }
+            Ok(asn) => Some(asn),
             Err(error) => {
                 tracing::warn!(%ip, %error, "ip not found in maxmind database");
                 None

diff --git a/src/proxy.rs b/src/proxy.rs
@@ -70,33 +70,31 @@ impl DownstreamReceiveWorkerConfig {
                     "Awaiting packet"
                 );
 
-                tokio::select! {
-                    result = socket.recv_from(&mut buf) => {
-                        match result {
-                            Ok((size, source)) => {
-                                let packet = DownstreamPacket {
-                                    received_at: chrono::Utc::now().timestamp_nanos(),
-                                    asn_info: crate::maxmind_db::MaxmindDb::lookup(source.ip()),
-                                    contents: buf[..size].to_vec(),
-                                    source,
-                                };
-
-                                if let Some(last_received_at) = last_received_at {
-                                    crate::metrics::packet_jitter(
-                                        crate::metrics::READ,
-                                        packet.asn_info.as_ref(),
-                                    )
-                                        .set(packet.received_at - last_received_at);
-                                }
-                                last_received_at = Some(packet.received_at);
-
-                                Self::spawn_process_task(packet, source, worker_id, &socket, &config, &sessions)
-                            }
-                            Err(error) => {
-                                tracing::error!(%error, "error receiving packet");
-                                return;
-                            }
+                match socket.recv_from(&mut buf).await {
+                    Ok((size, source)) => {
+                        let packet = DownstreamPacket {
+                            received_at: chrono::Utc::now().timestamp_nanos(),
+                            asn_info: crate::maxmind_db::MaxmindDb::lookup(source.ip()),
+                            contents: buf[..size].to_vec(),
+                            source,
+                        };
+
+                        if let Some(last_received_at) = last_received_at {
+                            crate::metrics::packet_jitter(
+                                crate::metrics::READ,
+                                packet.asn_info.as_ref(),
+                            )
+                            .set(packet.received_at - last_received_at);
                         }
+                        last_received_at = Some(packet.received_at);
+
+                        Self::spawn_process_task(
+                            packet, source, worker_id, &socket, &config, &sessions,
+                        )
+                    }
+                    Err(error) => {
+                        tracing::error!(%error, "error receiving packet");
+                        return;
                     }
                 }
             }

diff --git a/src/proxy/sessions.rs b/src/proxy/sessions.rs
@@ -109,6 +109,17 @@ impl Session {
         };
 
         tracing::debug!(source = %s.source, dest = ?s.dest, "Session created");
+        if let Some(asn) = &s.asn_info {
+            tracing::info!(
+                number = asn.r#as,
+                organization = asn.as_name,
+                country_code = asn.as_cc,
+                prefix = asn.prefix,
+                prefix_entity = asn.prefix_entity,
+                prefix_name = asn.prefix_name,
+                "maxmind information"
+            );
+        }
 
         self::metrics::total_sessions().inc();
         s.active_session_metric().inc();