refactor: resolve dns using a happyballs algorithm

server/utils/dnsCacheManager.ts

@@ -1,7 +1,8 @@
-import { getSettings } from '@server/lib/settings';
 import logger from '@server/logger';
 import { LRUCache } from 'lru-cache';
+import type { RecordWithTtl } from 'node:dns';
 import dns from 'node:dns';
+import net from 'node:net';
 
 interface DnsCache {
   address: string;
@@ -15,15 +16,31 @@ interface CacheStats {
   misses: number;
 }
 
+interface ConnectionResult {
+  address: string;
+  family: number;
+  duration: number;
+  success: boolean;
+}
+
 class DnsCacheManager {
   private cache: LRUCache<string, DnsCache>;
   private lookupAsync: typeof dns.promises.lookup;
   private resolver: dns.promises.Resolver;
+  private connectionCache: Map<
+    string,
+    {
+      preferredFamily: number;
+      lastSuccess: number;
+    }
+  >;
   private stats: CacheStats = {
     hits: 0,
     misses: 0,
   };
   private hardTtlMs: number;
+  private readonly IPV_DELAY = 50;
+  private readonly CONNECTION_TIMEOUT = 2000;
 
   constructor(maxSize = 500, hardTtlMs = 300000) {
     this.cache = new LRUCache<string, DnsCache>({
@@ -33,6 +50,175 @@ class DnsCacheManager {
     this.hardTtlMs = hardTtlMs;
     this.lookupAsync = dns.promises.lookup;
     this.resolver = new dns.promises.Resolver();
+    this.connectionCache = new Map();
+  }
+
+  private async attemptConnection(
+    address: string,
+    port: number,
+    family: number
+  ): Promise<ConnectionResult> {
+    const startTime = Date.now();
+
+    return new Promise((resolve) => {
+      const socket = new net.Socket();
+
+      const timeoutId = setTimeout(() => {
+        socket.destroy();
+        resolve({
+          address,
+          family,
+          duration: Date.now() - startTime,
+          success: false,
+        });
+      }, this.CONNECTION_TIMEOUT);
+
+      socket.once('error', () => {
+        clearTimeout(timeoutId);
+        socket.destroy();
+        resolve({
+          address,
+          family,
+          duration: Date.now() - startTime,
+          success: false,
+        });
+      });
+
+      socket.connect(port, address, () => {
+        clearTimeout(timeoutId);
+        socket.destroy();
+        resolve({
+          address,
+          family,
+          duration: Date.now() - startTime,
+          success: true,
+        });
+      });
+    });
+  }
+
+  private async resolveAddresses(hostname: string): Promise<{
+    ipv4Records: RecordWithTtl[];
+    ipv6Records: RecordWithTtl[];
+  }> {
+    const [ipv4Result, ipv6Result] = await Promise.allSettled([
+      this.resolver.resolve4(hostname, { ttl: true }),
+      this.resolver.resolve6(hostname, { ttl: true }),
+    ]);
+
+    return {
+      ipv4Records: ipv4Result.status === 'fulfilled' ? ipv4Result.value : [],
+      ipv6Records: ipv6Result.status === 'fulfilled' ? ipv6Result.value : [],
+    };
+  }
+
+  private sortAddresses(
+    hostname: string,
+    ipv4Records: RecordWithTtl[],
+    ipv6Records: RecordWithTtl[]
+  ): { ipv4Sorted: RecordWithTtl[]; ipv6Sorted: RecordWithTtl[] } {
+    const cached = this.connectionCache.get(hostname);
+
+    if (cached) {
+      const sortBySuccess = (a: RecordWithTtl, b: RecordWithTtl) => {
+        const aSuccess = this.connectionCache.get(a.address)?.lastSuccess || 0;
+        const bSuccess = this.connectionCache.get(b.address)?.lastSuccess || 0;
+        return bSuccess - aSuccess;
+      };
+
+      if (cached.preferredFamily === 6) {
+        ipv6Records.sort(sortBySuccess);
+      } else {
+        ipv4Records.sort(sortBySuccess);
+      }
+    }
+
+    return {
+      ipv4Sorted: ipv4Records,
+      ipv6Sorted: ipv6Records,
+    };
+  }
+
+  private async resolveWithTtl(
+    hostname: string,
+    port = 443
+  ): Promise<{ address: string; family: number; ttl: number }> {
+    if (!this.resolver) {
+      throw new Error('Resolver is not initialized');
+    }
+
+    try {
+      const { ipv4Records, ipv6Records } = await this.resolveAddresses(
+        hostname
+      );
+
+      if (!ipv4Records.length && !ipv6Records.length) {
+        throw new Error(`No DNS records found for ${hostname}`);
+      }
+
+      const { ipv4Sorted, ipv6Sorted } = this.sortAddresses(
+        hostname,
+        ipv4Records,
+        ipv6Records
+      );
+
+      let ipv6Promise: Promise<ConnectionResult> | null = null;
+      if (ipv6Sorted.length) {
+        ipv6Promise = this.attemptConnection(ipv6Sorted[0].address, port, 6);
+      }
+
+      await new Promise((resolve) => setTimeout(resolve, this.IPV_DELAY));
+
+      let ipv4Promise: Promise<ConnectionResult> | null = null;
+      if (ipv4Sorted.length) {
+        ipv4Promise = this.attemptConnection(ipv4Sorted[0].address, port, 4);
+      }
+
+      const attempts: Promise<ConnectionResult>[] = [];
+      if (ipv6Promise) {
+        attempts.push(ipv6Promise);
+      }
+      if (ipv4Promise) {
+        attempts.push(ipv4Promise);
+      }
+
+      const result = await Promise.race(attempts);
+
+      if (!result.success) {
+        this.connectionCache.set(hostname, {
+          preferredFamily: result.family,
+          lastSuccess: Date.now(),
+        });
+        this.connectionCache.set(result.address, {
+          preferredFamily: result.family,
+          lastSuccess: Date.now(),
+        });
+      }
+
+      const record =
+        result.family === 6
+          ? ipv6Sorted.find((r) => r.address === result.address)
+          : ipv4Sorted.find((r) => r.address === result.address);
+
+      // const ttl = record.ttl > 0 ? record.ttl * 1000 : 30000;
+      const ttl = record && record.ttl > 0 ? record.ttl * 1000 : 30000;
+      logger.debug(
+        `Resolved ${hostname} with TTL: ${record?.ttl} (Original), ${ttl} (Applied)`,
+        {
+          label: 'DNSCache',
+          address: result.address,
+          family: result.family,
+          duration: result.duration,
+        }
+      );
+
+      return { address: result.address, family: result.family, ttl };
+    } catch (error) {
+      logger.error(`Failed to resolve ${hostname} with TTL: ${error.message}`, {
+        label: 'DNSCache',
+      });
+      throw error;
+    }
   }
 
   async lookup(hostname: string): Promise<DnsCache> {
@@ -128,77 +314,6 @@ class DnsCacheManager {
     }
   }
 
-  private async resolveWithTtl(
-    hostname: string
-  ): Promise<{ address: string; family: number; ttl: number }> {
-    if (
-      !this.resolver ||
-      typeof this.resolver.resolve4 !== 'function' ||
-      typeof this.resolver.resolve6 !== 'function'
-    ) {
-      throw new Error('Resolver is not initialized');
-    }
-
-    try {
-      const [ipv4Records, ipv6Records] = await Promise.allSettled([
-        this.resolver.resolve4(hostname, { ttl: true }),
-        this.resolver.resolve6(hostname, { ttl: true }),
-      ]);
-
-      let record: { address: string; ttl: number } | null = null;
-      let family = 4;
-
-      const settings = getSettings();
-      const preferIpv6 = settings.main.forceIpv4First ? false : true;
-
-      if (preferIpv6) {
-        if (
-          ipv6Records.status === 'fulfilled' &&
-          ipv6Records.value.length > 0
-        ) {
-          record = ipv6Records.value[0];
-          family = 6;
-        } else if (
-          ipv4Records.status === 'fulfilled' &&
-          ipv4Records.value.length > 0
-        ) {
-          record = ipv4Records.value[0];
-          family = 4;
-        }
-      } else {
-        if (
-          ipv4Records.status === 'fulfilled' &&
-          ipv4Records.value.length > 0
-        ) {
-          record = ipv4Records.value[0];
-          family = 4;
-        } else if (
-          ipv6Records.status === 'fulfilled' &&
-          ipv6Records.value.length > 0
-        ) {
-          record = ipv6Records.value[0];
-          family = 6;
-        }
-      }
-
-      if (!record) {
-        throw new Error('No DNS records found for ${hostname');
-      }
-
-      const ttl = record.ttl > 0 ? record.ttl * 1000 : 30000;
-      logger.debug(
-        `Resolved ${hostname} with TTL: ${record.ttl} (Original), ${ttl} (Applied)`
-      );
-
-      return { address: record.address, family, ttl };
-    } catch (error) {
-      logger.error(`Failed to resolve ${hostname} with TTL: ${error.message}`, {
-        label: 'DNSCache',
-      });
-      throw error;
-    }
-  }
-
   getStats() {
     const entries = [...this.cache.entries()];
     return {
@@ -252,6 +367,7 @@ class DnsCacheManager {
 
   clear() {
     this.cache.clear();
+    this.connectionCache.clear();
     this.stats.hits = 0;
     this.stats.misses = 0;
     logger.debug('DNS cache cleared', { label: 'DNSCache' });