Source file keepalive.ml

(*
 * Copyright (c) 2017 Docker Inc
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *)

type action = [
  | `SendProbe
  | `Wait of Duration.t
  | `Close
]

type state = {
  probes_sent: int
}

let alive = {
  probes_sent = 0;
}

let next ~configuration ~ns state =
  let open Tcpip.Tcp.Keepalive in
  let after_ns = configuration.after in
  (* Wait until [time] has gone past *)
  if after_ns > ns
  then `Wait (Int64.sub after_ns ns), alive
  else begin
    let sending_probes_for_ns = Int64.sub ns after_ns in
    let interval_ns = configuration.interval in
    let should_have_sent = Int64.(to_int (div sending_probes_for_ns interval_ns)) in
    if should_have_sent > configuration.probes
    then `Close, state
    else
      if should_have_sent > state.probes_sent
      then `SendProbe, { probes_sent = should_have_sent } (* we don't want to send back-to-back probes *)
      else begin
        let since_last_probe_ns = Int64.rem sending_probes_for_ns interval_ns in
        `Wait (Int64.sub interval_ns since_last_probe_ns), state
      end
  end

type t = {
  configuration: Tcpip.Tcp.Keepalive.t;
  callback: ([ `SendProbe | `Close ] -> unit Lwt.t);
  mutable state: state;
  mutable timer: unit Lwt.t;
  mutable start: int64;
}
(** A keep-alive timer *)

let rec restart t =
  let open Lwt.Infix in
  let ns = Int64.sub (Mirage_mtime.elapsed_ns ()) t.start in
  match next ~configuration:t.configuration ~ns t.state with
  | `Wait ns, state ->
    Mirage_sleep.ns ns >>= fun () ->
    t.state <- state;
    restart t
  | `SendProbe, state ->
    t.callback `SendProbe >>= fun () ->
    t.state <- state;
    restart t
  | `Close, _ ->
    t.callback `Close >>= fun () ->
    Lwt.return_unit

let create configuration callback =
  let state = alive in
  let timer = Lwt.return_unit in
  let start = Mirage_mtime.elapsed_ns () in
  let t = { configuration; callback; state; timer; start } in
  t.timer <- restart t;
  t

let refresh t =
  t.start <- Mirage_mtime.elapsed_ns ();
  t.state <- alive;
  Lwt.cancel t.timer;
  t.timer <- restart t