mte/unikernel/duniverse/happy-eyeballs/src/happy_eyeballs.ml
2025-11-11 02:07:51 +01:00

497 lines
18 KiB
OCaml

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