Statically enabled probes (#1388)

* Support [~enabled_at_init] argument of %probe

The following new constructs are accepted:
[%probe "name" ~enabled_at_init:true handler]
[%probe "name" ~enabled_at_init:false handler]
The default is
[%probe "name" handler]
and implies [~enabled_at_init:false]

* Propagate [enabled_at_init] to Emit

* emit call instead of jmp opcode for enabled_at_init probes

* Add test

* Initialize semaphore according to enabled_at_init

Check that all probes with the same name in the current compilation
unit have consistent [enabled_at_init].

* Improve error message

* Add a test for inconsistent enabled_at_init for %probe

* Refactor parsing

* Add a module for Probe in Flambda2

* Rename variables [probe_name] to [probe] in middle_end/flambda2 code

* Improve code style and make [desc] private

Author: gretay-js

backend/amd64/emit.mlp

@@ -732,15 +732,32 @@ let reset_probes () =
   probes := [];
   probe_semaphores := String.Map.empty
 
-let find_or_add_semaphore name =
+let find_or_add_semaphore name enabled_at_init dbg =
   match String.Map.find_opt name !probe_semaphores with
-  | Some label -> label
+  | Some (label, e) ->
+    (match e, enabled_at_init with
+     | None, None -> ()
+     | None, Some _ ->
+       let d = (label, enabled_at_init) in
+       probe_semaphores :=
+         String.Map.remove name !probe_semaphores |>
+         String.Map.add name d
+     | Some _, None ->
+       (*  [find_or_add_semaphore] is called with None for Iprobe_is_enabled
+           during code emission only. [find_or_add_semaphore] us called with
+           Some to emit probe notes only after all code is emitted. *)
+       assert false
+     | Some b, Some b' ->
+       if not (Bool.equal b b') then
+         raise (Error (Inconsistent_probe_init (name,dbg))));
+    label
   | None ->
     let sym = "caml_probes_semaphore_"^name in
-    probe_semaphores := String.Map.add name sym !probe_semaphores;
+    let d = (sym, enabled_at_init) in
+    probe_semaphores := String.Map.add name d !probe_semaphores;
     sym
 
-let emit_call_probe_handler_wrapper i ~probe_label =
+let emit_call_probe_handler_wrapper i ~enabled_at_init ~probe_label =
   assert !frame_required;
   let wrap_label = probe_handler_wrapper_name probe_label in
   (* We emit a cmp instruction that is effectively a nop:
@@ -756,8 +773,12 @@ let emit_call_probe_handler_wrapper i ~probe_label =
        result, instead of relying on an assembler that might choose a different
        encoding which produces an incorrect relocation and changes the meaning
        of the program. *)
-    (* Emit the required encoding of "cmp $0, %eax" directly, using .byte *)
-    D.byte (Const 0x3dL);
+    (* Emit the required encoding of "cmp $0, %eax" or "call" directly,
+       using .byte *)
+    if enabled_at_init then
+      D.byte (Const 0xe8L)
+    else
+      D.byte (Const 0x3dL);
     D.byte (Const 0L);
     D.byte (Const 0L);
     D.byte (Const 0L);
@@ -1222,7 +1243,7 @@ let emit_instr fallthrough i =
   | Lop(Iendregion) ->
       I.mov (arg i 0) (domain_field Domainstate.Domain_local_sp)
   | Lop (Iname_for_debugger _) -> ()
-  | Lop (Iprobe _) ->
+  | Lop (Iprobe { enabled_at_init; _ }) ->
     let probe_label = new_label () in
     let probe =
       { probe_label;
@@ -1238,9 +1259,9 @@ let emit_instr fallthrough i =
        intervening wrapper that deals with getting the arguments to the probe
        in the correct place and managing the spilling/reloading of live
        registers.  See [emit_probe_handler_wrapper] below. *)
-    emit_call_probe_handler_wrapper i ~probe_label
+    emit_call_probe_handler_wrapper i ~enabled_at_init ~probe_label
   | Lop (Iprobe_is_enabled { name; }) ->
-    let semaphore_sym = find_or_add_semaphore name in
+    let semaphore_sym = find_or_add_semaphore name None i.dbg in
     (* Load unsigned 2-byte integer value of the semaphore.
        According to the documentation [1],
        semaphores are of type unsigned short.
@@ -1584,7 +1605,8 @@ let emit_probe_handler_wrapper p =
   let wrap_label = probe_handler_wrapper_name p.probe_label in
   let probe_name, handler_code_sym =
     match p.probe_insn.desc with
-    | Lop (Iprobe { name; handler_code_sym; }) -> name, handler_code_sym
+    | Lop (Iprobe { name; handler_code_sym; enabled_at_init = _; }) ->
+      name, handler_code_sym
     | _ -> assert false
   in
   (* Restore stack frame info as it was at the probe site, so we can easily
@@ -1739,9 +1761,10 @@ let emit_probe_notes0 () =
     Printf.sprintf "%d@%s" (Reg.size_of_contents_in_bytes arg) arg_name
   in
   let describe_one_probe p =
-    let probe_name =
+    let probe_name, enabled_at_init =
       match p.probe_insn.desc with
-      | Lop (Iprobe { name; _; }) -> name
+      | Lop (Iprobe { name; enabled_at_init; handler_code_sym=_; }) ->
+        name, enabled_at_init
       | _ -> assert false
     in
     let args =
@@ -1750,7 +1773,10 @@ let emit_probe_notes0 () =
         []
       |> String.concat " "
     in
-    let semaphore_label = emit_symbol (find_or_add_semaphore probe_name) in
+    let semsym =
+      find_or_add_semaphore probe_name (Some enabled_at_init) p.probe_insn.dbg
+    in
+    let semaphore_label = emit_symbol semsym in
     let emit_desc () =
       D.qword (ConstLabel (emit_label p.probe_label));
       (match is_macosx_system with
@@ -1771,12 +1797,15 @@ let emit_probe_notes0 () =
    | true ->
      D.section ["__TEXT";"__probes"] None ["regular"]);
   D.align ~data:true 2;
-  String.Map.iter (fun _ label ->
+  String.Map.iter (fun _ (label, enabled_at_init) ->
+      (* Unresolved weak symbols have a zero value regardless
+         of the following initialization. *)
+      let enabled_at_init = Option.value enabled_at_init ~default:false in
       D.weak label;
       D.hidden label;
       _label label;
       D.word (const 0); (* for systemtap probes *)
-      D.word (const 0); (* for ocaml probes *)
+      D.word (const (Bool.to_int enabled_at_init)); (* for ocaml probes *)
       add_def_symbol label)
     !probe_semaphores
 

backend/cfg/cfg.ml

@@ -368,8 +368,8 @@ let dump_terminator' ?(print_reg = Printmach.reg) ?(res = [||]) ?(args = [||])
       | Alloc { bytes; dbginfo; mode } -> Mach.Ialloc { bytes; dbginfo; mode }
       | Checkbound { immediate = Some x } -> Mach.Iintop_imm (Icheckbound, x)
       | Checkbound { immediate = None } -> Mach.Iintop Icheckbound
-      | Probe { name; handler_code_sym } ->
-        Mach.Iprobe { name; handler_code_sym });
+      | Probe { name; handler_code_sym; enabled_at_init } ->
+        Mach.Iprobe { name; handler_code_sym; enabled_at_init });
     Format.fprintf ppf "%sgoto %d" sep label_after
   | Specific_can_raise { op; label_after } ->
     Format.fprintf ppf "%a" specific_can_raise op;

backend/cfg/cfg_equivalence.ml

@@ -294,10 +294,18 @@ let check_prim_call_operation :
     when Option.equal Int.equal expected_immediate result_immediate ->
     ()
   | ( Probe
-        { name = expected_name; handler_code_sym = expected_handler_code_sym },
-      Probe { name = result_name; handler_code_sym = result_handler_code_sym } )
+        { name = expected_name;
+          handler_code_sym = expected_handler_code_sym;
+          enabled_at_init = expected_enabled_at_init
+        },
+      Probe
+        { name = result_name;
+          handler_code_sym = result_handler_code_sym;
+          enabled_at_init = result_enabled_at_init
+        } )
     when String.equal expected_name result_name
-         && String.equal expected_handler_code_sym result_handler_code_sym ->
+         && String.equal expected_handler_code_sym result_handler_code_sym
+         && Bool.equal expected_enabled_at_init result_enabled_at_init ->
     ()
   | _ -> different location "primitive call operation"
  [@@ocaml.warning "-4"]

backend/cfg/cfg_intf.ml

@@ -51,7 +51,8 @@ module S = struct
     | Checkbound of { immediate : int option }
     | Probe of
         { name : string;
-          handler_code_sym : string
+          handler_code_sym : string;
+          enabled_at_init : bool
         }
 
   type operation =

backend/cfg/cfg_to_linear.ml

@@ -171,7 +171,8 @@ let linearize_terminator cfg_with_layout (func : string) start
         | Checkbound { immediate = None } -> Iintop Icheckbound
         | Checkbound { immediate = Some i } -> Iintop_imm (Icheckbound, i)
         | Alloc { bytes; dbginfo; mode } -> Ialloc { bytes; dbginfo; mode }
-        | Probe { name; handler_code_sym } -> Iprobe { name; handler_code_sym }
+        | Probe { name; handler_code_sym; enabled_at_init } ->
+          Iprobe { name; handler_code_sym; enabled_at_init }
       in
       branch_or_fallthrough [L.Lop op] label_after, None
     | Specific_can_raise { op; label_after } ->

backend/cfg/cfgize.ml

@@ -215,10 +215,13 @@ let basic_or_terminator_of_operation :
     Misc.fatal_error
       "Cfgize.basic_or_terminator_of_operation: \"the Iname_for_debugger\" \
        instruction is currently not supported "
-  | Iprobe { name; handler_code_sym } ->
+  | Iprobe { name; handler_code_sym; enabled_at_init } ->
     With_next_label
       (fun label_after ->
-        Prim { op = Probe { name; handler_code_sym }; label_after })
+        Prim
+          { op = Probe { name; handler_code_sym; enabled_at_init };
+            label_after
+          })
   | Iprobe_is_enabled { name } -> Basic (Op (Probe_is_enabled { name }))
   | Ibeginregion -> Basic (Op Begin_region)
   | Iendregion -> Basic (Op End_region)

backend/cfg/linear_to_cfg.ml

@@ -582,8 +582,8 @@ let rec create_blocks (t : t) (i : L.instruction) (block : C.basic_block)
       terminator_prim (Checkbound { immediate = Some i })
     | Ialloc { bytes; dbginfo; mode } ->
       terminator_prim (Alloc { bytes; dbginfo; mode })
-    | Iprobe { name; handler_code_sym } ->
-      terminator_prim (Probe { name; handler_code_sym })
+    | Iprobe { name; handler_code_sym; enabled_at_init } ->
+      terminator_prim (Probe { name; handler_code_sym; enabled_at_init })
     | Istackoffset bytes ->
       let desc = C.Op (C.Stackoffset bytes) in
       DLL.add_end block.body (create_instruction t desc i ~stack_offset);

backend/checkmach.ml

@@ -699,7 +699,7 @@ end = struct
       assert (not (Mach.operation_can_raise op));
       let r = Value.transform next in
       check t r "heap allocation" dbg
-    | Iprobe { name; handler_code_sym } ->
+    | Iprobe { name; handler_code_sym; enabled_at_init = __ } ->
       let desc = Printf.sprintf "probe %s handler %s" name handler_code_sym in
       transform_call t ~next ~exn handler_code_sym ~desc dbg
     | Icall_ind -> transform t ~next ~exn ~effect:Value.top "indirect call" dbg

backend/cmm.ml

@@ -213,7 +213,7 @@ and operation =
   | Ccmpf of float_comparison
   | Craise of Lambda.raise_kind
   | Ccheckbound
-  | Cprobe of { name: string; handler_code_sym: string; }
+  | Cprobe of { name: string; handler_code_sym: string; enabled_at_init: bool }
   | Cprobe_is_enabled of { name: string }
   | Copaque
   | Cbeginregion | Cendregion

backend/cmm.mli

@@ -215,7 +215,7 @@ and operation =
                    then the index.
                    It results in a bounds error if the index is greater than
                    or equal to the bound. *)
-  | Cprobe of { name: string; handler_code_sym: string; }
+  | Cprobe of { name: string; handler_code_sym: string; enabled_at_init: bool }
   | Cprobe_is_enabled of { name: string }
   | Copaque (* Sys.opaque_identity *)
   | Cbeginregion | Cendregion

backend/cmm_helpers.ml

@@ -4028,8 +4028,12 @@ let beginregion ~dbg = Cop (Cbeginregion, [], dbg)
 
 let endregion ~dbg region = Cop (Cendregion, [region], dbg)
 
-let probe ~dbg ~name ~handler_code_linkage_name ~args =
-  Cop (Cprobe { name; handler_code_sym = handler_code_linkage_name }, args, dbg)
+let probe ~dbg ~name ~handler_code_linkage_name ~enabled_at_init ~args =
+  Cop
+    ( Cprobe
+        { name; handler_code_sym = handler_code_linkage_name; enabled_at_init },
+      args,
+      dbg )
 
 let load ~dbg kind mut ~addr = Cop (Cload (kind, mut), [addr], dbg)
 

backend/cmm_helpers.mli

@@ -1131,6 +1131,7 @@ val probe :
   dbg:Debuginfo.t ->
   name:string ->
   handler_code_linkage_name:string ->
+  enabled_at_init:bool ->
   args:expression list ->
   expression
 

backend/cmmgen.ml

@@ -511,10 +511,10 @@ let rec transl env e =
       let ptr = transl env arg in
       let dbg = Debuginfo.none in
       ptr_offset ptr offset dbg
-  | Udirect_apply(handler_code_sym, args, Some { name; }, _, _, dbg) ->
+  | Udirect_apply(handler_code_sym, args, Some { name; enabled_at_init }, _, _, dbg) ->
       let args = List.map (transl env) args in
       return_unit dbg
-        (Cop(Cprobe { name; handler_code_sym; }, args, dbg))
+        (Cop(Cprobe { name; handler_code_sym; enabled_at_init }, args, dbg))
   | Udirect_apply(lbl, args, None, result_layout, kind, dbg) ->
     let args = List.map (transl env) args in
     let sym =

backend/emitaux.ml

@@ -18,6 +18,7 @@
 type error =
   | Stack_frame_too_large of int
   | Stack_frame_way_too_large of int
+  | Inconsistent_probe_init of string * Debuginfo.t
 
 exception Error of error
 
@@ -506,3 +507,7 @@ let report_error ppf = function
                         Use -long-frames compiler flag." n
   | Stack_frame_way_too_large n ->
     Format.fprintf ppf "stack frame too large (%d bytes)." n
+  | Inconsistent_probe_init (name, dbg) ->
+    Format.fprintf ppf "Inconsistent use of ~enabled_at_init in [%%probe %s ..] at %a"
+      name Debuginfo.print_compact dbg
+

backend/emitaux.mli

@@ -106,6 +106,7 @@ val reduce_heap_size : reset:(unit -> unit) -> unit
 type error =
   | Stack_frame_too_large of int
   | Stack_frame_way_too_large of int
+  | Inconsistent_probe_init of string * Debuginfo.t
 
 module Dwarf_helpers : sig
   val init: disable_dwarf:bool -> string -> unit

backend/mach.ml

@@ -79,7 +79,7 @@ type operation =
   | Ipoll of { return_label: Cmm.label option }
   | Iname_for_debugger of { ident : Backend_var.t; which_parameter : int option;
       provenance : unit option; is_assignment : bool; }
-  | Iprobe of { name: string; handler_code_sym: string; }
+  | Iprobe of { name: string; handler_code_sym: string; enabled_at_init: bool; }
   | Iprobe_is_enabled of { name: string }
   | Ibeginregion | Iendregion
 

backend/mach.mli

@@ -89,7 +89,7 @@ type operation =
         (b) If [is_assignment] is [true], any information about other [Reg.t]s
             that have been previously deemed to hold the value of that
             identifier is forgotten. *)
-  | Iprobe of { name: string; handler_code_sym: string; }
+  | Iprobe of { name: string; handler_code_sym: string; enabled_at_init: bool; }
   | Iprobe_is_enabled of { name: string }
   | Ibeginregion | Iendregion
 

backend/printcmm.ml

@@ -226,8 +226,9 @@ let operation d = function
   | Ccmpf c -> Printf.sprintf "%sf" (float_comparison c)
   | Craise k -> Lambda.raise_kind k ^ location d
   | Ccheckbound -> "checkbound" ^ location d
-  | Cprobe { name; handler_code_sym } ->
-    Printf.sprintf "probe[%s %s]" name handler_code_sym
+  | Cprobe { name; handler_code_sym; enabled_at_init; } ->
+    Printf.sprintf "probe[%s %s%s]" name handler_code_sym
+      (if enabled_at_init then " enabled_at_init" else "")
   | Cprobe_is_enabled {name} -> Printf.sprintf "probe_is_enabled[%s]" name
   | Cprefetch { is_write; locality; } ->
     Printf.sprintf "prefetch is_write=%b prefetch_temporal_locality_hint=%s"

backend/selectgen.ml

@@ -690,8 +690,8 @@ method select_operation op args _dbg =
     (Iintop_atomic { op = Compare_and_swap; size; addr }, [compare_with; set_to; eloc])
   | (Ccheckbound, _) ->
     self#select_arith Icheckbound args
-  | (Cprobe { name; handler_code_sym; }, _) ->
-    Iprobe { name; handler_code_sym; }, args
+  | (Cprobe { name; handler_code_sym; enabled_at_init; }, _) ->
+    Iprobe { name; handler_code_sym; enabled_at_init; }, args
   | (Cprobe_is_enabled {name}, _) -> Iprobe_is_enabled {name}, []
   | (Cbeginregion, _) -> Ibeginregion, []
   | (Cendregion, _) -> Iendregion, args

middle_end/flambda2/compare/compare.ml

@@ -438,8 +438,8 @@ and subst_apply env apply =
   let args_arity = Apply_expr.args_arity apply in
   let return_arity = Apply_expr.return_arity apply in
   Apply_expr.create ~callee ~continuation exn_continuation ~args ~call_kind dbg
-    ~inlined ~inlining_state ~probe_name:None ~position ~relative_history
-    ~region ~args_arity ~return_arity
+    ~inlined ~inlining_state ~probe:None ~position ~relative_history ~region
+    ~args_arity ~return_arity
   |> Expr.create_apply
 
 and subst_apply_cont env apply_cont =
@@ -1002,7 +1002,7 @@ let apply_exprs env apply1 apply2 : Expr.t Comparison.t =
             ~args:args1' ~call_kind:call_kind1' (Apply.dbg apply1)
             ~inlined:(Apply.inlined apply1)
             ~inlining_state:(Apply.inlining_state apply1)
-            ~probe_name:None ~position:(Apply.position apply1)
+            ~probe:None ~position:(Apply.position apply1)
             ~relative_history:(Apply_expr.relative_history apply1)
             ~region:(Apply_expr.region apply1)
             ~args_arity:(Apply_expr.args_arity apply1)