From 686bd8b061bc47c874a056b912a02be71a374a8e Mon Sep 17 00:00:00 2001
From: Keuin <keuinx@gmail.com>
Date: Sat, 29 Jul 2023 01:21:40 +0800
Subject: cache DNS resolve result to reduce ~15% latency

---
 src/main.rs | 55 +++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 41 insertions(+), 14 deletions(-)

(limited to 'src/main.rs')

diff --git a/src/main.rs b/src/main.rs
index 27051f1..d52d3b9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,12 +1,17 @@
 use std::io;
 use std::net::{SocketAddr, ToSocketAddrs};
 
+use std::time::{Duration, SystemTime, SystemTimeError};
+
 use clap::{arg, Parser};
 use tokio::net::UdpSocket;
 use tracing::{debug, error};
 
+
 use crate::Destination::{Address, SockAddr};
 
+const TTL: Duration = Duration::from_secs(300);
+
 /// Simple program to greet a person
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -51,6 +56,8 @@ async fn main() {
         args_sendto = s;
         sendto = Address(&args_sendto);
     }
+    let mut addr_sendto: Option<SocketAddr> = None;
+    let mut addr_sendto_t = SystemTime::UNIX_EPOCH;
     loop {
         match sock.recv_from(&mut buf).await {
             Ok((len, addr)) => {
@@ -71,24 +78,44 @@ async fn main() {
                 }
                 let sendto = match sendto {
                     Destination::None => Destination::None,
-                    Address(addr) => match resolve_addr(&addr) {
-                        Ok(addr) => SockAddr(addr),
-                        Err(why) => {
-                            error!("Failed to resolve address {}: {}", addr, why);
-                            Destination::None
+                    Address(addr) => {
+                        let now = SystemTime::now();
+                        if now.duration_since(addr_sendto_t)
+                            .or::<SystemTimeError>(Ok(Duration::from_secs(0))).unwrap() > TTL {
+                            // update
+                            debug!("resolving address: {}", &addr);
+                            match resolve_addr(&addr) {
+                                Ok(addr) => {
+                                    addr_sendto = Some(addr);
+                                    addr_sendto_t = SystemTime::now();
+                                    SockAddr(addr)
+                                }
+                                Err(why) => {
+                                    error!("Failed to resolve address {}: {}", addr, why);
+                                    Destination::None
+                                }
+                            }
+                        } else {
+                            // cache is still valid
+                            match addr_sendto {
+                                None => Destination::None,
+                                Some(addr) => SockAddr(addr),
+                            }
                         }
-                    },
+                    }
                     SockAddr(addr) => SockAddr(addr),
                 };
-                let target = if let SockAddr(sendto_addr) = sendto {
-                    if addr == sendto_addr {
-                        // remote -> local
-                        Some(local)
-                    } else {
-                        Some(sendto_addr)
+                let target = match sendto {
+                    Destination::None => None,
+                    Address(_) => panic!("unresolved address"),
+                    SockAddr(sendto_addr) => {
+                        if addr == sendto_addr {
+                            // remote -> local
+                            Some(local)
+                        } else {
+                            Some(sendto_addr)
+                        }
                     }
-                } else {
-                    None
                 };
                 match target {
                     None => {
-- 
cgit v1.2.3