use crate::seqstring::global_string;
use crate::stack::{Stack, pop, push};
use crate::value::{Value, VariantData};
use std::collections::HashMap;
use std::net::ToSocketAddrs;
use std::sync::mpsc as std_mpsc;
use std::sync::{Arc, LazyLock, Mutex};
use std::thread;
use std::time::{Duration, Instant};
const DEFAULT_WORKERS: usize = 8;
const MAX_WORKERS: usize = 64;
const CACHE_TTL: Duration = Duration::from_secs(60);
const CACHE_MAX: usize = 256;
struct CacheEntry {
addrs: Vec<String>,
expires: Instant,
}
struct DnsCache {
entries: HashMap<String, CacheEntry>,
}
impl DnsCache {
fn new() -> Self {
Self {
entries: HashMap::new(),
}
}
fn get(&mut self, host: &str) -> Option<Vec<String>> {
if let Some(entry) = self.entries.get(host)
&& Instant::now() < entry.expires
{
return Some(entry.addrs.clone());
}
self.entries.remove(host);
None
}
fn put(&mut self, host: String, addrs: Vec<String>) {
if self.entries.len() >= CACHE_MAX {
if let Some(k) = self.entries.keys().next().cloned() {
self.entries.remove(&k);
}
}
self.entries.insert(
host,
CacheEntry {
addrs,
expires: Instant::now() + CACHE_TTL,
},
);
}
}
static CACHE: LazyLock<Mutex<DnsCache>> = LazyLock::new(|| Mutex::new(DnsCache::new()));
struct Job {
hostname: String,
reply: may::sync::mpsc::Sender<Vec<String>>,
}
static JOB_QUEUE: LazyLock<Option<std_mpsc::Sender<Job>>> = LazyLock::new(|| {
let (tx, rx) = std_mpsc::channel::<Job>();
let rx = Arc::new(Mutex::new(rx));
let workers = std::env::var("SEQ_DNS_WORKERS")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|n| *n > 0)
.unwrap_or(DEFAULT_WORKERS)
.min(MAX_WORKERS);
let mut spawned = 0usize;
for i in 0..workers {
let rx = rx.clone();
if thread::Builder::new()
.name(format!("seq-dns-{i}"))
.spawn(move || worker_loop(rx))
.is_ok()
{
spawned += 1;
}
}
if spawned == 0 { None } else { Some(tx) }
});
fn worker_loop(rx: Arc<Mutex<std_mpsc::Receiver<Job>>>) {
loop {
let job = match rx.lock().unwrap().recv() {
Ok(j) => j,
Err(_) => return, };
let addrs = resolve_blocking(&job.hostname);
if !addrs.is_empty() {
CACHE
.lock()
.unwrap()
.put(job.hostname.clone(), addrs.clone());
}
let _ = job.reply.send(addrs); }
}
fn resolve_blocking(hostname: &str) -> Vec<String> {
match (hostname, 0u16).to_socket_addrs() {
Ok(iter) => {
let mut seen = Vec::new();
for sa in iter {
let ip = sa.ip().to_string();
if !seen.contains(&ip) {
seen.push(ip);
}
}
seen
}
Err(_) => Vec::new(),
}
}
pub fn resolve(hostname: &str) -> Vec<String> {
if hostname.is_empty() {
return Vec::new();
}
if let Some(addrs) = CACHE.lock().unwrap().get(hostname) {
return addrs;
}
let sender = match JOB_QUEUE.as_ref() {
Some(s) => s,
None => return Vec::new(),
};
let (reply_tx, reply_rx) = may::sync::mpsc::channel::<Vec<String>>();
let job = Job {
hostname: hostname.to_string(),
reply: reply_tx,
};
if sender.send(job).is_err() {
return Vec::new();
}
reply_rx.recv().unwrap_or_default()
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn patch_seq_dns_resolve(stack: Stack) -> Stack {
unsafe {
let (stack, host_val) = pop(stack);
let host = match host_val {
Value::String(s) => s,
_ => return push_failure(stack),
};
let hostname = host.as_str_or_empty().to_string();
let addrs = resolve(&hostname);
push_result(stack, addrs)
}
}
unsafe fn push_result(stack: Stack, addrs: Vec<String>) -> Stack {
unsafe {
if addrs.is_empty() {
return push_failure(stack);
}
let fields = addrs
.into_iter()
.map(|s| Value::String(global_string(s)))
.collect();
let list = Value::Variant(Arc::new(VariantData::new(
global_string("List".to_string()),
fields,
)));
let stack = push(stack, list);
push(stack, Value::Bool(true))
}
}
unsafe fn push_failure(stack: Stack) -> Stack {
unsafe {
let empty = Value::Variant(Arc::new(VariantData::new(
global_string("List".to_string()),
vec![],
)));
let stack = push(stack, empty);
push(stack, Value::Bool(false))
}
}