unikernel/duniverse/happy-eyeballs/src/happy_eyeballs.ml

@@ -0,0 +1,497 @@
+let src = Logs.Src.create "happy-eyeballs" ~doc:"Happy Eyeballs"
+module Log = (val Logs.src_log src : Logs.LOG)
+
+type conn_state =
+  | Resolving of int64
+  | Waiting_for_aaaa of int64 * Ipaddr.V4.Set.t (* TODO ensure non-empty set *)
+  | Connecting of int64 * (Ipaddr.t * int) list * (Ipaddr.t * int) list
+
+type timeouts = {
+  created : int64 ;
+  aaaa_timeout : int64 ;
+  connect_delay : int64 ;
+  connect_timeout : int64 ;
+  resolve_timeout : int64 ;
+}
+
+type connection = {
+  state : conn_state ;
+  ports : int list ;
+  resolved : [ `none | `v4 | `v6 | `both ] ;
+  resolve_left : int ;
+  attempt : int ;
+  time : timeouts ;
+}
+
+let resolve st ev = match st, ev with
+  | `none, `v4 -> `v4
+  | `none, `v6 -> `v6
+  | `v6, `v4 -> `both
+  | `v4, `v6 -> `both
+  | x, _ -> x
+
+module IM = Map.Make(Int)
+
+type t = {
+  resolve_retries : int ;
+  counter : int ;
+  conns : connection IM.t Domain_name.Host_map.t ;
+  time : timeouts ;
+}
+
+let resolve_timeout t = t.time.resolve_timeout
+
+type id = int
+
+type action =
+  | Resolve_a of [`host] Domain_name.t
+  | Resolve_aaaa of [`host] Domain_name.t
+  | Connect of [`host] Domain_name.t * id * int * (Ipaddr.t * int)
+  | Connect_failed of [`host] Domain_name.t * id * string
+
+let host_or_ip v =
+  match Ipaddr.of_domain_name v with
+  | None -> Domain_name.to_string v
+  | Some ip -> Ipaddr.to_string ip
+
+let pp_action ppf = function
+  | Resolve_a host -> Fmt.pf ppf "resolve A %a" Domain_name.pp host
+  | Resolve_aaaa host -> Fmt.pf ppf "resolve AAAA %a" Domain_name.pp host
+  | Connect (host, id, attempt, (ip, port)) ->
+    Fmt.pf ppf "%u connect %s (using %a:%u), attempt %u" id (host_or_ip host)
+         Ipaddr.pp ip port attempt
+  | Connect_failed (host, id, reason) ->
+    Fmt.pf ppf "%u connect failed %s: %s" id (host_or_ip host) reason
+
+type event =
+  | Resolved_a of [`host] Domain_name.t * Ipaddr.V4.Set.t
+  | Resolved_aaaa of [`host] Domain_name.t * Ipaddr.V6.Set.t
+  | Resolved_a_failed of [`host] Domain_name.t * string
+  | Resolved_aaaa_failed of [`host] Domain_name.t * string
+  | Connection_failed of [`host] Domain_name.t * id * (Ipaddr.t * int) * string
+  | Connected of [`host] Domain_name.t * id * (Ipaddr.t * int)
+
+let pp_event ppf = function
+  | Resolved_a (host, ips) ->
+    Fmt.pf ppf "resolved A %a: %a" Domain_name.pp host
+      Fmt.(list ~sep:(any ", ") Ipaddr.V4.pp)
+      (Ipaddr.V4.Set.elements ips)
+  | Resolved_aaaa (host, ips) ->
+    Fmt.pf ppf "resolved AAAA %a: %a" Domain_name.pp host
+      Fmt.(list ~sep:(any ", ") Ipaddr.V6.pp)
+      (Ipaddr.V6.Set.elements ips)
+  | Resolved_a_failed (host, reason) ->
+    Fmt.pf ppf "resolve A failed for %a: %s" Domain_name.pp host reason
+  | Resolved_aaaa_failed (host, reason) ->
+    Fmt.pf ppf "resolve AAAA failed for %a: %s" Domain_name.pp host reason
+  | Connection_failed (host, id, (ip, port), reason) ->
+    Fmt.pf ppf "%u connection to %s failed %a:%d: %s" id (host_or_ip host)
+      Ipaddr.pp ip port reason
+  | Connected (host, id, (ip, port)) ->
+    Fmt.pf ppf "%u connected to %s (using %a:%d)" id (host_or_ip host)
+      Ipaddr.pp ip port
+
+let ctr = ref 0
+
+let create
+    ?(aaaa_timeout = Duration.of_ms 50)
+    ?(connect_delay = Duration.of_ms 50)
+    ?(connect_timeout = Duration.of_sec 10)
+    ?(resolve_timeout = Duration.of_sec 1)
+    ?(resolve_retries = 3)
+    created =
+  incr ctr;
+  let time = {
+    aaaa_timeout ;
+    connect_delay ;
+    connect_timeout ;
+    resolve_timeout ;
+    created ;
+  } in
+  {
+    resolve_retries ;
+    counter = !ctr ;
+    conns = Domain_name.Host_map.empty ;
+    time ;
+  }
+
+let add_conn host id conn c =
+  Domain_name.Host_map.update host
+    (function
+      | None -> Some (IM.singleton id conn)
+      | Some cs -> Some (IM.add id conn cs))
+    c
+
+let expand_list ips ports =
+  List.flatten (List.map (fun ip -> List.map (fun p -> (ip, p)) ports) ips)
+
+(* all input has been verified that ips and ports are non-empty. *)
+let expand_list_split ips ports =
+  match expand_list ips ports with
+  | hd :: tl -> hd, tl
+  | _ -> failwith "ips or ports are empty"
+
+let tick now host id conn =
+  match conn.state with
+  | Resolving ts when Int64.sub now ts > conn.time.resolve_timeout ->
+    begin
+      let ok actions =
+        Ok ({ conn with resolve_left = conn.resolve_left - 1 },
+            actions)
+      in
+      match conn.resolve_left <= 1, conn.resolved with
+      | true, _ | _, `both -> Error ()
+      | false, `none -> ok [ Resolve_a host ; Resolve_aaaa host ]
+      | false, `v4 -> ok [ Resolve_aaaa host ]
+      | false, `v6 -> ok [ Resolve_a host ]
+    end
+  | Waiting_for_aaaa (started, ips) when Int64.sub now started > conn.time.aaaa_timeout ->
+    let ips = List.map (fun ip -> Ipaddr.V4 ip) (Ipaddr.V4.Set.elements ips) in
+    let dst, dsts = expand_list_split ips conn.ports in
+    let state = Connecting (now, [ dst ], dsts)
+    and attempt = conn.attempt + 1
+    in
+    Ok ({ conn with state ; attempt }, [ Connect (host, id, conn.attempt, dst) ])
+  | Connecting (last_conn, active_conns, dsts) ->
+    (* if there are further IP addresses, and there was no activity within
+       connect_delay, start the next connection. *)
+    if Int64.sub now conn.time.created > conn.time.connect_timeout then
+      Error ()
+    else if Int64.sub now last_conn > conn.time.connect_delay then
+      (match dsts with
+       | [] -> Ok (conn, [])
+       | dst :: dsts ->
+         let state = Connecting (now, dst :: active_conns, dsts)
+         and attempt = conn.attempt + 1
+         in
+         Ok ({ conn with state ; attempt }, [ Connect (host, id, conn.attempt, dst) ]))
+    else
+      Ok (conn, [])
+  | _ -> Ok (conn, [])
+
+let timer (t : t) now =
+  let conns, actions =
+    Domain_name.Host_map.fold (fun host v (dm, actions) ->
+        let v, actions = IM.fold (fun id conn (acc, actions) ->
+            match tick now host id conn with
+            | Ok (conn, action) -> IM.add id conn acc, action @ actions
+            | Error () -> acc, Connect_failed (host, id, "timeout") :: actions)
+            v (IM.empty, actions)
+        in
+        let dm =
+          if IM.cardinal v = 0 then dm else Domain_name.Host_map.add host v dm
+        in
+        dm, actions) t.conns (Domain_name.Host_map.empty, [])
+  in
+  (match actions with
+   | [] when not (Domain_name.Host_map.is_empty conns) -> ()
+   | _ ->
+     Log.debug (fun m -> m "[%u] timer continue %B, %d actions: %a"
+                   t.counter (not (Domain_name.Host_map.is_empty conns))
+                   (List.length actions)
+                   Fmt.(list ~sep:(any "@.") pp_action) actions));
+  { t with conns },
+  (if Domain_name.Host_map.is_empty conns then `Suspend else `Act),
+  actions
+
+let connect t now ?aaaa_timeout ?connect_delay ?connect_timeout ?resolve_timeout ?resolve_retries ~id host ports =
+  Log.debug (fun m -> m "[%u] connect: id %d host %a" t.counter id
+                Domain_name.pp host);
+  if ports = [] then failwith "empty port list not supported";
+  let tt = t.time in
+  let time = {
+    created = now ;
+    aaaa_timeout = Option.value ~default:tt.aaaa_timeout aaaa_timeout ;
+    connect_delay = Option.value ~default:tt.connect_delay connect_delay ;
+    connect_timeout = Option.value ~default:tt.connect_timeout connect_timeout ;
+    resolve_timeout = Option.value ~default:tt.resolve_timeout resolve_timeout ;
+  }
+  in
+  let resolve_left = Option.value ~default:t.resolve_retries resolve_retries in
+  let conn = {
+    ports ;
+    state = Resolving now ;
+    resolved = `none ;
+    resolve_left ;
+    attempt = 0 ;
+    time ;
+  } in
+  let actions = [ Resolve_aaaa host ; Resolve_a host ] in
+  Log.debug (fun m -> m "[%u] actions: %a" t.counter
+                Fmt.(list ~sep:(any "@.") pp_action) actions);
+  { t with conns = add_conn host id conn t.conns }, actions
+
+let merge ?(ipv4 = Ipaddr.V4.Set.empty) ?(ipv6 = Ipaddr.V6.Set.empty) ips =
+  List.fold_left (fun (ipv4, ipv6) -> function
+      | Ipaddr.V4 ip -> Ipaddr.V4.Set.add ip ipv4, ipv6
+      | Ipaddr.V6 ip -> ipv4, Ipaddr.V6.Set.add ip ipv6)
+    (ipv4, ipv6) ips
+
+let shuffle ?first v4 v6 =
+  match List.length v4, List.length v6 with
+  | 0, _ -> v6
+  | _, 0 -> v4
+  | v4l, v6l ->
+    let rec shuffle a b = function
+      | 0 -> []
+      | n -> match a, b with
+        | [], _ -> shuffle v4 b n
+        | _, [] -> shuffle a v6 n
+        | hd :: tl, hd' :: tl' ->
+          match first with
+          | Some Ipaddr.V6 _ -> hd :: hd' :: shuffle tl tl' (pred n)
+          | None | Some Ipaddr.V4 _ -> hd' :: hd :: shuffle tl tl' (pred n)
+    in
+    shuffle v4 v6 (max v4l v6l)
+
+(* the idea is to first separate into V4 and V6 addresses, and then mix them *)
+let mix ?first ?ipv4 ?ipv6 ips =
+  let ipv4, ipv6 = merge ?ipv4 ?ipv6 ips in
+  let v4, v6 =
+    Ipaddr.V4.Set.fold (fun ip acc -> Ipaddr.V4 ip :: acc) ipv4 [],
+    Ipaddr.V6.Set.fold (fun ip acc -> Ipaddr.V6 ip :: acc) ipv6 []
+  in
+  shuffle ?first v4 v6
+
+let mix_dsts ?(ipv4 = Ipaddr.V4.Set.empty) ?(ipv6 = Ipaddr.V6.Set.empty) ports dst dsts =
+  let v4_present, v6_present = merge (List.map fst (dst @ dsts)) in
+  let ipv4 = Ipaddr.V4.Set.diff ipv4 v4_present
+  and ipv6 = Ipaddr.V6.Set.diff ipv6 v6_present
+  in
+  let v4_dsts, v6_dsts =
+    List.fold_left (fun (ipv4, ipv6) -> function
+        | Ipaddr.V4 _, _ as a -> a :: ipv4, ipv6
+        | Ipaddr.V6 _, _ as a -> ipv4, a :: ipv6)
+      ([], []) dsts
+  in
+  let v4s =
+    expand_list
+      (Ipaddr.V4.Set.fold (fun ip acc -> Ipaddr.V4 ip :: acc) ipv4 [])
+      ports
+  and v6s =
+    expand_list
+      (Ipaddr.V6.Set.fold (fun ip acc -> Ipaddr.V6 ip :: acc) ipv6 [])
+      ports
+  in
+  let first = match dst with [] -> None | (ip, _) :: _ -> Some ip in
+  shuffle ?first (List.rev v4_dsts @ v4s) (List.rev v6_dsts @ v6s)
+
+let connect_ip t now ?aaaa_timeout ?connect_delay ?connect_timeout ~id dsts =
+  Log.debug (fun m -> m "[%u] connect_ip id %d dsts %a" t.counter id
+                Fmt.(list ~sep:(any ", ") (pair ~sep:(any ":") Ipaddr.pp int))
+                dsts);
+  let dst, dsts = match dsts with
+    | dst :: dsts -> dst, dsts
+    | [] -> failwith "addresses are empty"
+  in
+  let state = Connecting (now, [ dst ], dsts) in
+  let tt = t.time in
+  let time = {
+    created = now ;
+    aaaa_timeout = Option.value ~default:tt.aaaa_timeout aaaa_timeout ;
+    connect_delay = Option.value ~default:tt.connect_delay connect_delay ;
+    connect_timeout = Option.value ~default:tt.connect_timeout connect_timeout ;
+    resolve_timeout = 0L
+  }
+  in
+  let conn = {
+    ports = [] ;
+    state ;
+    resolved = `both ;
+    resolve_left = 0 ;
+    attempt = 1 ;
+    time ;
+  } in
+  let host = Ipaddr.to_domain_name (fst dst) in
+  let actions = [ Connect (host, id, 0, dst) ] in
+  Log.debug (fun m -> m "[%u] actions: %a" t.counter
+                Fmt.(list ~sep:(any "@.") pp_action) actions);
+  { t with conns = add_conn host id conn t.conns }, actions
+
+let event t now e =
+  Log.debug (fun m -> m "[%u] received event %a" t.counter pp_event e);
+  let t, actions =
+    match e with
+    | Resolved_a (name, ips) ->
+      let conns, actions =
+        match Domain_name.Host_map.find name t.conns with
+        | None -> t.conns, []
+        | Some cs ->
+          let cs, actions = IM.fold (fun id c (cs, actions) ->
+              let resolved = resolve c.resolved `v4 in
+              let state, attempt, actions = match c.state with
+                | Resolving _ts when resolved = `both ->
+                  let ips =
+                    List.map (fun ip -> Ipaddr.V4 ip) (Ipaddr.V4.Set.elements ips)
+                  in
+                  let dst, dsts = expand_list_split ips c.ports in
+                  Connecting (now, [ dst ], dsts), c.attempt + 1,
+                  Connect (name, id, c.attempt, dst) :: actions
+                | Resolving _ts -> Waiting_for_aaaa (now, ips), c.attempt, actions
+                | Waiting_for_aaaa (ts, ips') ->
+                  Log.debug (fun m -> m "%a already waiting for AAAA with %a"
+                               Domain_name.pp name
+                               Fmt.(list ~sep:(any ", ") Ipaddr.V4.pp)
+                               (Ipaddr.V4.Set.elements ips'));
+                  Waiting_for_aaaa (ts, Ipaddr.V4.Set.union ips' ips), c.attempt,
+                  actions
+                | Connecting (ts, dst, dsts) ->
+                  let dsts = mix_dsts ~ipv4:ips c.ports dst dsts in
+                  Connecting (ts, dst, dsts), c.attempt, actions
+              in
+              IM.add id { c with state ; resolved ; attempt } cs, actions)
+              cs (IM.empty, [])
+          in
+          Domain_name.Host_map.add name cs t.conns, actions
+      in
+      { t with conns }, actions
+    | Resolved_a_failed (name, reason) ->
+      let conns, actions =
+        match Domain_name.Host_map.find name t.conns with
+        | None -> t.conns, []
+        | Some cs ->
+          let cs, actions = IM.fold (fun id c (cs, actions) ->
+              let resolved = resolve c.resolved `v4 in
+              match c.state with
+              | Resolving _ts when resolved = `both ->
+                cs, Connect_failed (name, id, reason) :: actions
+              | _ -> IM.add id { c with resolved } cs, actions)
+              cs (IM.empty, [])
+          in
+          (if IM.is_empty cs then
+             Domain_name.Host_map.remove name t.conns
+           else
+             Domain_name.Host_map.add name cs t.conns), actions
+      in
+      { t with conns }, actions
+    | Resolved_aaaa (name, ips) ->
+      let conns, actions =
+        match Domain_name.Host_map.find name t.conns with
+        | None -> t.conns, []
+        | Some cs ->
+          let cs, actions = IM.fold (fun id c (cs, actions) ->
+              let resolved = resolve c.resolved `v6 in
+              let state, attempt, actions' = match c.state with
+                | Resolving _ts ->
+                  let ips = mix ~ipv6:ips [] in
+                  let dst, dsts = expand_list_split ips c.ports in
+                  Connecting (now, [ dst ], dsts), c.attempt + 1,
+                  [ Connect (name, id, c.attempt, dst) ]
+                | Waiting_for_aaaa (_ts, ips') ->
+                  let ips = mix ~ipv4:ips' ~ipv6:ips [] in
+                  let dst, dsts = expand_list_split ips c.ports in
+                  Connecting (now, [ dst ], dsts), c.attempt + 1,
+                  [ Connect (name, id, c.attempt, dst) ]
+                | Connecting (ts, dst, dsts) ->
+                  let dsts = mix_dsts ~ipv6:ips c.ports dst dsts in
+                  Connecting (ts, dst, dsts), c.attempt, []
+              in
+              IM.add id { c with state ; resolved ; attempt } cs, actions @ actions')
+              cs (IM.empty, [])
+          in
+          Domain_name.Host_map.add name cs t.conns, actions
+      in
+      { t with conns }, actions
+    | Resolved_aaaa_failed (name, reason) ->
+      let conns, actions =
+        match Domain_name.Host_map.find name t.conns with
+        | None -> t.conns, []
+        | Some cs ->
+          let cs, actions = IM.fold (fun id c (cs, actions) ->
+              let resolved = resolve c.resolved `v6 in
+              match c.state with
+              | Resolving _ts when resolved = `both ->
+                cs, Connect_failed (name, id, reason) :: actions
+              | Waiting_for_aaaa (_ts, ips) ->
+                let ips =
+                  List.map (fun ip -> Ipaddr.V4 ip) (Ipaddr.V4.Set.elements ips)
+                in
+                let dst, dsts = expand_list_split ips c.ports in
+                let state = Connecting (now, [ dst ], dsts) in
+                let attempt = c.attempt + 1 in
+                IM.add id { c with state ; resolved ; attempt } cs,
+                Connect (name, id, c.attempt, dst) :: actions
+              | _ -> IM.add id { c with resolved } cs, actions)
+              cs (IM.empty, [])
+          in
+          (if IM.is_empty cs then
+             Domain_name.Host_map.remove name t.conns
+           else
+             Domain_name.Host_map.add name cs t.conns), actions
+      in
+      { t with conns }, actions
+    | Connection_failed (name, id, (ip, port), reason) ->
+      let conns, actions =
+        match Domain_name.Host_map.find name t.conns with
+        | None ->
+          Log.warn (fun m -> m "[%u] connection failed to %s: %s; no entry in conns"
+                       t.counter (host_or_ip name) reason);
+          t.conns, []
+        | Some cs ->
+          match IM.find_opt id cs with
+          | None ->
+            Log.warn (fun m -> m "[%u] %u connection failed to %s: %s; no entry in IM"
+                         t.counter id (host_or_ip name) reason);
+            t.conns, []
+          | Some c ->
+            let not_failed (ip', port') = not (Ipaddr.compare ip ip' = 0 && port = port') in
+            match c.state with
+            | Connecting (ts, dst, []) ->
+              let dst' = List.filter not_failed dst in
+              begin match dst', c.resolved with
+                | [], `both ->
+                  let cs = IM.remove id cs in
+                  Domain_name.Host_map.add name cs t.conns,
+                  [ Connect_failed (name, id, reason) ]
+                | [], _ ->
+                  let state = Resolving now in
+                  let cs = IM.add id { c with state } cs in
+                  Domain_name.Host_map.add name cs t.conns, []
+                | dst', _ ->
+                  let state = Connecting (ts, dst', []) in
+                  let cs = IM.add id { c with state } cs in
+                  Domain_name.Host_map.add name cs t.conns, []
+              end
+            | Connecting (_ts, dst, ndst :: dsts) ->
+              let dst' = List.filter not_failed dst in
+              let state = Connecting (now, ndst :: dst', dsts) in
+              let attempt = c.attempt + 1 in
+              let cs = IM.add id { c with state ; attempt } cs in
+              Domain_name.Host_map.add name cs t.conns,
+              [ Connect (name, id, c.attempt, ndst) ]
+            | _ -> t.conns, []
+      in
+      { t with conns }, actions
+    | Connected (name, id, (_ip, _port)) ->
+      let conns =
+        Domain_name.Host_map.update name (function
+            | None ->
+              Log.warn (fun m -> m "[%u] connected to an unexpected domain: %a"
+                           t.counter Domain_name.pp name);
+              None
+            | Some xs ->
+              let m = IM.remove id xs in
+              if IM.cardinal m = 0 then None else Some m)
+          t.conns
+      in
+      { t with conns }, []
+  in
+  Log.debug (fun m -> m "[%u] actions: %a" t.counter
+                Fmt.(list ~sep:(any "@.") pp_action) actions);
+  t, actions
+
+module Waiter_map = struct
+  include Map.Make(Int)
+
+  let _id = ref 0
+
+  let register v t =
+    incr _id;
+    let id = !_id in
+    add id v t, id
+
+  let find_and_remove id t =
+    match find_opt id t with
+    | None -> t, None
+    | Some x -> remove id t, Some x
+end