Merge pull request #122 from talex5/setup-flags

Add Uring.Setup_flags

Author: talex5

lib/uring/heap.mli

@@ -28,7 +28,7 @@ type ptr = private int
 
 val ptr : 'a entry -> ptr
 (** [ptr e] is the index of [e].
-    @raise Invalid_arg if [e] has already been freed. *)
+    @raise Invalid_argument if [e] has already been freed. *)
 
 val alloc : 'a t -> 'a -> extra_data:'b -> 'a entry
 (** [alloc t a ~extra_data] adds the value [a] to [t] and returns a

lib/uring/include/discover.ml

@@ -99,7 +99,24 @@ let uring_ops = [
   "IORING_OP_URING_CMD";
 ]
 
-let ops c =
+let uring_setup_flags = [
+  "IORING_SETUP_IOPOLL";
+  (* "IORING_SETUP_SQPOLL" *)  (* Enabled by passing a polling timeout instead *)
+  (* "IORING_SETUP_SQ_AFF"; *) (* todo: requires setting sq_thread_cpu too *)
+  (* "IORING_SETUP_CQSIZE"; *) (* todo: requires setting cq_entries *)
+  "IORING_SETUP_CLAMP";
+  (* "IORING_SETUP_ATTACH_WQ"; *) (* todo: requires setting wq_fd *)
+  "IORING_SETUP_R_DISABLED";
+  "IORING_SETUP_SUBMIT_ALL";
+  "IORING_SETUP_COOP_TASKRUN";
+  "IORING_SETUP_TASKRUN_FLAG";
+  "IORING_SETUP_SQE128";
+  "IORING_SETUP_CQE32";
+  "IORING_SETUP_SINGLE_ISSUER";
+  "IORING_SETUP_DEFER_TASKRUN";
+]
+
+let uring_defs c =
   Gen.abstract_module "Op" (
     C.C_define.import c (List.map (fun name -> name, C.C_define.Type.Int) uring_ops)
       ~c_flags:["-D_GNU_SOURCE"; "-I"; include_dir]
@@ -111,6 +128,18 @@ let ops c =
           (ocaml_name, v)
         | _ -> assert false
       )
+  ) @
+  Gen.hex_module "Ioring_setup" (
+    C.C_define.import c (List.map (fun name -> name, C.C_define.Type.Int) uring_setup_flags)
+      ~c_flags:["-D_GNU_SOURCE"; "-I"; include_dir]
+      ~includes:["liburing.h"]
+    |> List.map (function
+        | name, C.C_define.Value.Int v ->
+          let prefix_len = String.length "IORING_SETUP_" in
+          let ocaml_name = String.sub name prefix_len (String.length name - prefix_len) |> String.lowercase_ascii in
+          (ocaml_name, v)
+        | _ -> assert false
+      )
   )
 
 let stat_flags c =
@@ -199,7 +228,7 @@ let () =
   C.main ~name:"discover" (fun c ->
       C.Flags.write_lines "config.ml" @@ List.flatten @@ List.map (fun f -> f c) [
         toplevel_defs;
-        ops;
+        uring_defs;
         stat_flags;
       ]
     )

lib/uring/primitives.h

@@ -57,7 +57,7 @@ CAMLprim value ocaml_uring_set_iovec(value, value);
 CAMLprim value ocaml_uring_set_string(value, value);
 CAMLprim value ocaml_uring_make_msghdr(value, value, value);
 CAMLprim value ocaml_uring_get_msghdr_fds(value);
-CAMLprim value ocaml_uring_setup(value, value);
+CAMLprim value ocaml_uring_setup(value, value, value);
 CAMLprim value ocaml_uring_exit(value);
 CAMLprim value ocaml_uring_unregister_buffers(value);
 CAMLprim value ocaml_uring_register_ba(value, value);

lib/uring/uring.ml

@@ -86,6 +86,11 @@ module Poll_mask = struct
   let pollhup = Config.pollhup
 end
 
+module Setup_flags = struct
+  include Flags
+  include Config.Ioring_setup
+end
+
 module Statx = struct
   type t
 
@@ -299,7 +304,7 @@ type probe
 module Uring = struct
   type t
 
-  external create : int -> int option -> t = "ocaml_uring_setup"
+  external create : int -> int option -> Setup_flags.t -> t = "ocaml_uring_setup"
   external exit : t -> unit = "ocaml_uring_exit"
 
   external unregister_buffers : t -> unit = "ocaml_uring_unregister_buffers"
@@ -394,9 +399,9 @@ let register_gc_root t =
 let unregister_gc_root t =
   update_gc_roots (Ring_set.remove (Generic_ring.T t))
 
-let create ?polling_timeout ~queue_depth () =
+let create ?(flags=Setup_flags.empty) ?polling_timeout ~queue_depth () =
   if queue_depth < 1 then Fmt.invalid_arg "Non-positive queue depth: %d" queue_depth;
-  let uring = Uring.create queue_depth polling_timeout in
+  let uring = Uring.create queue_depth polling_timeout flags in
   let data = Heap.create queue_depth in
   let id = object end in
   let fixed_iobuf = Cstruct.empty.buffer in

lib/uring/uring.mli

@@ -23,14 +23,65 @@
 
 module Region = Region
 
+(** Type of flags that can be combined. *)
+module type FLAGS = sig
+  type t = private int
+  (** A set of flags. *)
+
+  val empty : t
+
+  val of_int : int -> t
+
+  val ( + ) : t -> t -> t
+  (** [a + b] is the union of the sets. *)
+
+  val mem : t -> t -> bool
+  (** [mem x flags] is [true] iff [x] is a subset of [flags]. *)
+end
+
+(** Flags that can be passed to {!create}. *)
+module Setup_flags : sig
+  include FLAGS
+
+  val iopoll : t
+  (** io_context is polled *)
+
+  val clamp : t
+  (** Clamp SQ/CQ ring sizes *)
+
+  val r_disabled : t
+  (** Start with ring disabled *)
+
+  val submit_all : t
+  (** Continue submit on error *)
+
+  val coop_taskrun : t
+  (** Cooperative task running *)
+
+  val taskrun_flag : t
+  (** Get notified if task work is available *)
+
+  val sqe128 : t
+  (** SQEs are 128 byte *)
+
+  val cqe32 : t
+  (** CQEs are 32 byte *)
+
+  val single_issuer : t
+  (** Only one task is allowed to submit requests *)
+
+  val defer_taskrun : t
+  (** Defer running task work to get events *)
+end
+
 type 'a t
 (** ['a t] is a reference to an Io_uring structure. *)
 
 type 'a job
 (** A handle for a submitted job, which can be used to cancel it.
     If an operation returns [None], this means that submission failed because the ring is full. *)
 
-val create : ?polling_timeout:int -> queue_depth:int -> unit -> 'a t
+val create : ?flags:Setup_flags.t -> ?polling_timeout:int -> queue_depth:int -> unit -> 'a t
 (** [create ~queue_depth] will return a fresh Io_uring structure [t].
     Initially, [t] has no fixed buffer. Use {!set_fixed_buffer} if you want one.
     @param polling_timeout If given, use polling mode with the given idle timeout (in ms).
@@ -86,21 +137,6 @@ val timeout: ?absolute:bool -> 'a t -> clock -> int64 -> 'a -> 'a job option
 
     [ns] is the timeout time in nanoseconds *)
 
-module type FLAGS = sig
-  type t = private int
-  (** A set of flags. *)
-
-  val empty : t
-
-  val of_int : int -> t
-
-  val ( + ) : t -> t -> t
-  (** [a + b] is the union of the sets. *)
-
-  val mem : t -> t -> bool
-  (** [mem x flags] is [true] iff [x] is a subset of [flags]. *)
-end
-
 (** Flags that can be passed to {!openat2}. *)
 module Open_flags : sig
   include FLAGS

lib/uring/uring_stubs.c

@@ -77,7 +77,7 @@ static struct custom_operations ring_ops = {
   custom_fixed_length_default
 };
 
-value ocaml_uring_setup(value entries, value polling_timeout) {
+value ocaml_uring_setup(value entries, value polling_timeout, value v_flags) {
   CAMLparam1(entries);
   CAMLlocal1(v_uring);
   struct io_uring_params params;
@@ -90,6 +90,7 @@ value ocaml_uring_setup(value entries, value polling_timeout) {
   Ring_val(v_uring) = ring;
 
   memset(&params, 0, sizeof(params));
+  params.flags = Int_val(v_flags);
   if (Is_some(polling_timeout)) {
     params.flags |= IORING_SETUP_SQPOLL;
     params.sq_thread_idle = Int_val(Some_val(polling_timeout));