Merge pull request #321 from djs55/async-dns

Mac: make the DNS more scalable

Author: djs55

src/bin/main.ml

@@ -352,6 +352,11 @@ let hvsock_addr_of_uri ~default_serviceid uri =
       if debug || env_debug then Some Logs.Debug else Some Logs.Info in
     Logging.setup log_destination level;
 
+    if Sys.os_type = "Unix" then begin
+      Log.info (fun f -> f "Increasing preemptive thread pool size to 1024 threads");
+      Uwt_preemptive.set_bounds (0, 1024);
+    end;
+
     let host_names = List.map Dns.Name.of_string @@ Astring.String.cuts ~sep:"," host_names in
     let dns_path, resolver = match dns with
     | None -> None, Configuration.default_resolver

src/hostnet/host.ml

@@ -1133,11 +1133,66 @@ module Dns = struct
 
   let resolve_dnssd question =
     let open Dns.Packet in
-    let query name ty =
-      Uwt_preemptive.detach
-        (fun () ->
-           Dnssd.query (Dns.Name.to_string name) ty
-        ) () in
+    let query_one name ty =
+      let query = Dnssd.LowLevel.query (Dns.Name.to_string name) ty in
+      let socket = Dnssd.LowLevel.socket query in
+      let t, u = Lwt.task () in
+      match Uwt.Poll.start socket [ Uwt.Poll.Readable ]
+        ~cb:(fun _poll events ->
+          match events with
+          | Error error ->
+            Log.err (fun f -> f "Uwt.Poll callback failed with %s" (Uwt.strerror error))
+          | Ok events ->
+            List.iter (fun event ->
+              if event = Uwt.Poll.Readable then Lwt.wakeup_later u ()
+            ) events
+        ) with
+      | Error error ->
+        Log.err (fun f -> f "Uwt.Poll.start failed with %s" (Uwt.strerror error));
+        Lwt.return (Ok [])
+      | Ok poll ->
+        t >>= fun () ->
+        let result = Uwt.Poll.close poll in
+        if not (Uwt.Int_result.is_ok result) then begin
+          let error = Uwt.Int_result.to_error result in
+          Log.err (fun f -> f "Uwt.Poll.close failed with %s" (Uwt.strerror error));
+          Lwt.return (Ok [])
+        end else begin
+          Uwt_preemptive.detach
+            (fun () ->
+              Dnssd.LowLevel.response query
+            ) ()
+        end in
+
+    let query requested_name ty =
+      (* The DNSServiceRef API will return CNAMEs first, without resolving to
+        A/AAAA/... This function recursively resolves the CNAMES while avoiding
+        returning duplicate records. *)
+      (* NB we only return NoSuchRecord if we find no records. This is because it
+        is possible to query a CNAME which exists, but which points to a non-existent
+        record. *)
+      let open Dnssd in
+      let rec loop acc name ty =
+        query_one name ty
+        >>= function
+        (* When we're recursing, ignore the NoSuchRecord error *)
+        | Error NoSuchRecord when name <> requested_name -> Lwt.return (Ok acc)
+        | Error e -> Lwt.return (Error e)
+        | Ok rrs ->
+          let not_seen_before = List.filter (fun x -> not (List.mem x acc)) rrs in
+          (* If there are any CNAMEs, resolve these too *)
+          let cnames = List.rev @@ List.fold_left (fun acc rr ->
+            match rr.Dns.Packet.rdata with
+            | CNAME name -> name :: acc
+            | _ -> acc
+          ) [] not_seen_before in
+          Lwt_list.fold_left_s
+            (fun acc name -> match acc with
+              | Error e -> Lwt.return (Error e)
+              | Ok acc ->
+                loop acc name ty
+            ) (Ok (acc @ not_seen_before)) cnames in
+      loop [] requested_name ty in
 
     begin match question with
     | { q_class = Q_IN; q_name; _ } when q_name = localhost_local ->