(*****************************************************************************)
(* *)
(* Open Source License *)
(* Copyright (c) 2023 Nomadic Labs <contact@nomadic-labs.com> *)
(* *)
(* Permission is hereby granted, free of charge, to any person obtaining a *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense, *)
(* and/or sell copies of the Software, and to permit persons to whom the *)
(* Software is furnished to do so, subject to the following conditions: *)
(* *)
(* The above copyright notice and this permission notice shall be included *)
(* in all copies or substantial portions of the Software. *)
(* *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER *)
(* DEALINGS IN THE SOFTWARE. *)
(* *)
(*****************************************************************************)
module Parameters = struct
type persistent_state = {
name : Agent_name.t;
local_octogram_path : string;
mutable pending_ready : unit option Lwt.u list;
runner : Runner.t;
}
type session_state = {mutable ready : bool}
let base_default_name = "octogram-prometheus"
let default_colors = Log.Color.[|FG.magenta|]
end
open Parameters
include Daemon.Make (Parameters)
let wait agent =
match agent.status with
| Not_running ->
Test.fail
"%s is not running, cannot wait for it to terminate"
(name agent)
| Running {process; _} -> Process.wait process
let wait_for_promise ?timeout ?where agent name promise =
let promise = Lwt.map Result.ok promise in
let* result =
match timeout with
| None -> promise
| Some timeout ->
Lwt.pick
[
promise;
(let* () = Lwt_unix.sleep timeout in
Lwt.return_error ());
]
in
match result with
| Ok (Some x) -> return x
| Ok None ->
raise (Terminated_before_event {daemon = agent.name; event = name; where})
| Error () ->
Format.ksprintf
failwith
"Timeout waiting for event %s of %s"
name
agent.name
let trigger_ready agent value =
let pending = agent.persistent_state.pending_ready in
agent.persistent_state.pending_ready <- [] ;
List.iter (fun pending -> Lwt.wakeup_later pending value) pending
let set_ready agent =
(match agent.status with
| Not_running -> ()
| Running status -> status.session_state.ready <- true) ;
trigger_ready agent (Some ())
let wait_for_ready agent =
match agent.status with
| Running {session_state = {ready = true; _}; _} -> unit
| Not_running | Running {session_state = {ready = false; _}; _} ->
let promise, resolver = Lwt.task () in
agent.persistent_state.pending_ready <-
resolver :: agent.persistent_state.pending_ready ;
wait_for_promise agent "Prometheus is ready" promise
let handle_raw_stdout agent line =
if line =~ rex Prometheus_agent.greetings then set_ready agent
let on_stdout agent handler =
agent.stdout_handlers <- handler :: agent.stdout_handlers
let create ~runner ~(name : Agent_name.t) ?path ?color ?event_pipe () =
let local_octogram_path = Option.value ~default:"./octogram" path in
let path = "/tmp" // sf "octogram-%s" (name :> string) in
let agent =
create
~path
~name:(name :> string)
~runner
?color
?event_pipe
{pending_ready = []; local_octogram_path; runner; name}
in
on_stdout agent (handle_raw_stdout agent) ;
agent
let run ?event_level ?event_sections_levels agent =
if agent.status <> Not_running then Test.fail "Remote agent is not running" ;
let on_terminate _ =
trigger_ready agent None ;
unit
in
let* () =
Helpers.deploy
~r:true
~for_runner:agent.persistent_state.runner
agent.persistent_state.local_octogram_path
agent.path
in
(Background.register
@@ let* () = wait_for_ready agent in
Log.debug "Agent is ready, removing the binary" ;
Helpers.exec ~runner:agent.persistent_state.runner "rm" [agent.path]) ;
run
~runner:agent.persistent_state.runner
?event_level
?event_sections_levels
agent
{ready = false}
~on_terminate
[
Format.(
sprintf "--%s" @@ Cli.get_string ~default:"info" "remote_verbosity");
"prometheus";
"-a";
sf "agent_name=%s" (agent.persistent_state.name :> string);
]
let with_session_status state k =
match state.status with
| Running status -> k status
| _ -> Test.fail "Remote prometheus is not running"
let record_metrics_source agent source =
with_session_status agent @@ fun status ->
let str =
Helpers.to_json_string Prometheus_agent.input_encoding (Record source)
in
Log.debug "%s" str ;
Lwt_io.write status.stdin Format.(sprintf "%s\n%!" str)
let quit agent =
with_session_status agent @@ fun status ->
let str = Helpers.to_json_string Prometheus_agent.input_encoding Quit in
Log.debug "%s" str ;
Lwt_io.write status.stdin Format.(sprintf "%s\n%!" str)