From 0cdb9aaf2eb770ed6ac33559e8e5e72dc63a2004 Mon Sep 17 00:00:00 2001 From: Patrick Ferris Date: Fri, 5 Jun 2026 13:25:41 +0100 Subject: [PATCH] Yield in for and while loops We make sure to yield to other commands inside merry shell operations. Additionally, we make sure to close the stdin if we need to. --- src/lib/eval.ml | 137 ++++++++++++++++++++++++----------- src/lib/import.ml | 6 ++ src/lib/job.ml | 6 +- src/lib/posix/exec.ml | 1 + src/lib/posix/merry_posix.ml | 2 +- src/lib/types.ml | 5 +- test/nofork.t | 27 +++++++ 7 files changed, 136 insertions(+), 48 deletions(-) create mode 100644 test/nofork.t diff --git a/src/lib/eval.ml b/src/lib/eval.ml index 72d9289..2b7b4b6 100644 --- a/src/lib/eval.ml +++ b/src/lib/eval.ml @@ -153,11 +153,13 @@ module Make (S : Types.State) (E : Types.Exec) = struct stdin : Eio_unix.source_ty Eio.Flow.source; } - let stdout_for_pipeline ~sw (ctx : pipeline_ctx) = function - | [] -> (None, `Global ctx.stdout) + let std_for_pipeline ?(first = false) ~sw (ctx : pipeline_ctx) rest = + match (first, rest) with + | true, [] -> (`Global ctx.stdin, `Global ctx.stdout) + | false, [] -> (`Local ctx.stdin, `Global ctx.stdout) | _ -> let r, w = Safe_fd.pipe sw in - (Some r, `Local (w :> Eio_unix.sink_ty Eio.Flow.sink)) + (`Global r, `Local (w :> Eio_unix.sink_ty Eio.Flow.sink)) let fd_of_int ?(close_unix = true) ~sw (n : int) = Eio_unix.Fd.of_unix ~close_unix ~sw (Obj.magic n : Unix.file_descr) @@ -194,6 +196,18 @@ module Make (S : Types.State) (E : Types.Exec) = struct Option.iter (fun f -> f ()) ctx.exit_handler; exit code + let catch_execs_with_error fn = + try fn () + with Eio.Io (err, _) -> + Fmt.epr "%a\n%!" Eio.Exn.pp_err err; + Error 2 + + let catch_execs_with_exit ctx fn = + try fn () + with Eio.Io (err, _) -> + Fmt.epr "%a\n%!" Eio.Exn.pp_err err; + Exit.nonzero ~message:(Fmt.str "%a" Eio.Exn.pp_err err) ctx 2 + let rec handle_pipeline ~async initial_ctx p : ctx Exit.t = let set_last_background ~async process ctx = if async then begin @@ -216,18 +230,18 @@ module Make (S : Types.State) (E : Types.Exec) = struct | `Error p -> J.add_error p j | `Exit p -> J.add_exit p j in - let close_stdout ~is_global some_write = + let close_flow ~is_global flow = if not is_global then begin - Eio.Flow.close some_write + Eio.Flow.close flow end in - let update_stdin ~stdin ctx = - { ctx with stdin = Option.value ~default:ctx.stdin stdin } - in + let update_stdin ~stdin ctx = { ctx with stdin } in let exec_process ~sw ctx job ?fds ?stdin ~stdout ?pgid executable args = let pgid = match pgid with None -> 0 | Some p -> p in let reap = J.get_reaper job in - let mode = if async then Types.Async else Types.Switched sw in + let mode = + if async then Types.Switched ctx.async_switch else Types.Switched sw + in let fds = ctx.rdrs @ Option.value ~default:[] fds in let ctx, process = let hash, prog = @@ -241,7 +255,7 @@ module Make (S : Types.State) (E : Types.Exec) = struct Eio.Flow.copy_string (Fmt.str "msh: command not found: %s\n" executable) stdout; - (ctx, Error (127, `Not_found)) + (ctx, Error 127) | _, Some full_path -> Debug.Log.debug (fun f -> f "executing %a\n%a" @@ -250,13 +264,14 @@ module Make (S : Types.State) (E : Types.Exec) = struct Fmt.(list Types.pp_redirect) fds); ( ctx, + catch_execs_with_error @@ fun () -> E.exec ctx.executor ~delay_reap:(fst reap) ~fds ?stdin ~stdout ~pgid ~mode ~cwd:(cwd_of_ctx ctx) ~pipe:Safe_fd.pipe ~env:(get_env ~extra:ctx.local_state ctx) ~executable:full_path (executable :: args) ) in match process with - | Error (n, _) -> + | Error n -> let job = handle_job job (`Error n) in (on_process ~async ctx, job) | Ok process -> @@ -278,7 +293,7 @@ module Make (S : Types.State) (E : Types.Exec) = struct let ctx = collect_assignments initial_ctx prefix in let job = handle_job job (`Noop ctx) in loop pctx job rest - | Ast.SimpleCommand v :: rest -> ( + | Ast.SimpleCommand v :: rest as pipeline -> ( let ctx, executable, suffix = match v with | Prefixed (prefix, Some exec, suffix) -> @@ -323,13 +338,20 @@ module Make (S : Types.State) (E : Types.Exec) = struct let job = handle_job job (`Noop ctx) in loop pctx job rest | Exit.Zero ctx -> ( - let some_read, some_write = - stdout_for_pipeline ~sw:pipeline_switch pctx rest + let first = List.length p = List.length pipeline in + let std = + std_for_pipeline ~first ~sw:pipeline_switch pctx rest in - let is_global, some_write = - match some_write with - | `Global p -> (true, p) - | `Local p -> (false, p) + let is_stdin_global, is_stdout_global, some_read, some_write = + match std with + | `Local stdin, `Local stdout -> + (false, false, stdin, stdout) + | `Local stdin, `Global stdout -> + (false, true, stdin, stdout) + | `Global stdin, `Local stdout -> + (true, false, stdin, stdout) + | `Global stdin, `Global stdout -> + (true, true, stdin, stdout) in let rdrs = List.fold_left @@ -418,6 +440,7 @@ module Make (S : Types.State) (E : Types.Exec) = struct else job | ":" -> job | _ -> ( + (* TODO: Make concurrent *) let func_app = if is_command then None else @@ -427,19 +450,13 @@ module Make (S : Types.State) (E : Types.Exec) = struct in match func_app with | Some ctx -> - close_stdout ~is_global some_write; + close_flow ~is_global:is_stdout_global + some_write; (* TODO: Proper job stuff and redirects etc. *) let job = handle_job job (immediate_built_in ctx) in - let pctx = - { - pctx with - stdin = - Option.value ~default:pctx.stdin - some_read; - } - in + let pctx = { pctx with stdin = some_read } in loop pctx job rest | None -> ( match Built_ins.of_args command_args with @@ -455,10 +472,12 @@ module Make (S : Types.State) (E : Types.Exec) = struct @@ fun () -> Trace.name "built-in"; let bi = + catch_execs_with_exit ctx @@ fun () -> handle_built_in ~rdrs ~stdout:some_write ctx bi in - close_stdout ~is_global some_write; + close_flow ~is_global:is_stdout_global + some_write; bi in let ctx = @@ -528,7 +547,22 @@ module Make (S : Types.State) (E : Types.Exec) = struct ~pgid:(job_pgid job) executable args in - close_stdout ~is_global some_write; + close_flow ~is_global:is_stdout_global + some_write; + Eio.Fiber.fork + ~sw: + (if async then ctx.async_switch + else pipeline_switch) + (fun () -> + match + Option.map E.await + (J.last_process job) + with + | Some _ -> + close_flow + ~is_global:is_stdin_global + some_read + | None -> ()); let pctx = update_stdin ~stdin:some_read pctx in @@ -539,9 +573,10 @@ module Make (S : Types.State) (E : Types.Exec) = struct Fiber.fork_promise ~sw:pipeline_switch @@ fun () -> Trace.name "built-in"; let bi = + catch_execs_with_exit ctx @@ fun () -> handle_built_in ~rdrs ~stdout:some_write ctx bi in - close_stdout ~is_global some_write; + close_flow ~is_global:is_stdout_global some_write; bi in let ctx = @@ -549,7 +584,7 @@ module Make (S : Types.State) (E : Types.Exec) = struct (Exit.map ~f:clear_local_state) ctx in - close_stdout ~is_global some_write; + close_flow ~is_global:is_stdout_global some_write; let job = match bi with | Built_ins.Exit _ -> @@ -565,14 +600,15 @@ module Make (S : Types.State) (E : Types.Exec) = struct in let pctx = update_stdin ~stdin:some_read pctx in loop pctx job rest)))) - | CompoundCommand (c, rdrs) :: rest -> ( - let some_read, some_write = - stdout_for_pipeline ~sw:pipeline_switch pctx rest - in - let is_global, some_write = - match some_write with - | `Global p -> (true, p) - | `Local p -> (false, p) + | CompoundCommand (c, rdrs) :: rest as v -> ( + let first = List.length p = List.length v in + let std = std_for_pipeline ~first ~sw:pipeline_switch pctx rest in + let is_stdin_global, is_stdout_global, some_read, some_write = + match std with + | `Local stdin, `Local stdout -> (false, false, stdin, stdout) + | `Local stdin, `Global stdout -> (false, true, stdin, stdout) + | `Global stdin, `Local stdout -> (true, false, stdin, stdout) + | `Global stdin, `Global stdout -> (true, true, stdin, stdout) in match handle_redirections ~sw:pipeline_switch initial_ctx rdrs with | Error ctx -> handle_job job (`Rdr (Exit.nonzero ctx 1)) @@ -592,13 +628,19 @@ module Make (S : Types.State) (E : Types.Exec) = struct Trace.name "compound-command"; with_pipeline_scope ctx @@ fun ctx -> let ctx = handle_compound_command ctx c in - close_stdout ~is_global some_write; + close_flow ~is_global:is_stdout_global some_write; Exit.map ~f:(fun c -> { c with rdrs = saved_rdrs }) ctx in + (* Close stdin *) + Eio.Fiber.fork + ~sw: + (if async then initial_ctx.async_switch else pipeline_switch) + (fun () -> + match Option.map E.await (J.last_process job) with + | Some _ -> close_flow ~is_global:is_stdin_global some_read + | None -> ()); let job = handle_job job (`Built_in ctx) in - let pctx = - { pctx with stdin = Option.value ~default:pctx.stdin some_read } - in + let pctx = { pctx with stdin = some_read } in loop pctx job rest) | FunctionDefinition (name, (body, _rdrs)) :: rest -> let ctx = @@ -639,7 +681,11 @@ module Make (S : Types.State) (E : Types.Exec) = struct } end else begin - let last_process = Option.map string_of_int @@ J.last_process job in + (* We must allow async jobs to be reaped, J.await_exit does this for us. *) + J.reap job; + let last_process = + Option.map (fun p -> string_of_int (E.pid p)) @@ J.last_process job + in Exit.zero { ctx with @@ -1226,6 +1272,7 @@ module Make (S : Types.State) (E : Types.Exec) = struct try Nlist.fold_left (fun _ (_, words) -> + Fiber.yield (); List.fold_left (fun _ word -> update ctx ~param:name word.Ast.txt >>= fun ctx -> @@ -1327,6 +1374,8 @@ module Make (S : Types.State) (E : Types.Exec) = struct match exec running_ctx (term, Some sep) with | Exit.Nonzero _ -> exit_so_far (* TODO: Context? *) | Exit.Zero ctx -> + (* Before we loop, we yield to other pipeline tasks. *) + Fiber.yield (); loop (try exec ctx (term', Some sep') with | Continue (1, ctx) -> Exit.zero ctx diff --git a/src/lib/import.ml b/src/lib/import.ml index 3500f19..2e36619 100644 --- a/src/lib/import.ml +++ b/src/lib/import.ml @@ -6,6 +6,12 @@ module Trace = struct Eio.Private.Trace.name (Eio.Private.Fiber_context.tid ctx) s end +module Fmt = struct + include Fmt + + let lst pp ppf = (brackets @@ list ~sep:comma pp) ppf +end + module Promise = struct include Eio.Promise diff --git a/src/lib/job.ml b/src/lib/job.ml index 73cf7c3..fced84b 100644 --- a/src/lib/job.ml +++ b/src/lib/job.ml @@ -39,6 +39,10 @@ module Make (E : Types.Exec) = struct let set_id new_id t = { t with id = new_id } let get_reaper t = t.reap + let reap t = + let _, r = t.reap in + try Eio.Promise.resolve r () with Invalid_argument _ -> () + let make id processes = let reap = Eio.Promise.create ~label:("reap-" ^ string_of_int id) () in { id; processes; reap } @@ -88,6 +92,6 @@ module Make (E : Types.Exec) = struct match List.find_map (function `Process p -> Some p | _ -> None) t.processes with - | Some (_, p) -> Some (E.pid p) + | Some (_, p) -> Some p | None -> None end diff --git a/src/lib/posix/exec.ml b/src/lib/posix/exec.ml index 445fe92..c5f1ef2 100644 --- a/src/lib/posix/exec.ml +++ b/src/lib/posix/exec.ml @@ -100,6 +100,7 @@ module Process = struct Fiber.fork_daemon ~sw (fun () -> Merry.Import.Trace.name "reap-daemon"; Option.iter Eio.Promise.await delay_reap; + Merry.Debug.Log.info (fun f -> f "Delay reap past: %i" t.pid); reap t set_exit_status; Switch.remove_hook hook; `Stop_daemon)) diff --git a/src/lib/posix/merry_posix.ml b/src/lib/posix/merry_posix.ml index 1e379ab..123df3c 100644 --- a/src/lib/posix/merry_posix.ml +++ b/src/lib/posix/merry_posix.ml @@ -28,5 +28,5 @@ module Exec = struct ?stdout ?stderr ?env t ~executable args) with Eio.Io (Eio.Process.E (Eio.Process.Executable_not_found m), _ctx) -> Fmt.epr "msh: command not found: %s\n%!" m; - Error (127, `Not_found) + Error 127 end diff --git a/src/lib/types.ml b/src/lib/types.ml index 1c3d18f..5dedf2f 100644 --- a/src/lib/types.ml +++ b/src/lib/types.ml @@ -104,7 +104,7 @@ module type Exec = sig executable:string -> t -> string list -> - (process, int * [ `Not_found ]) result + (process, int) result (** Run a command in a child process *) val await : process -> unit Exit.t @@ -117,6 +117,7 @@ module type Job = sig type process val get_reaper : _ t -> unit Eio.Promise.t * unit Eio.Promise.u + val reap : _ t -> unit val make : int -> @@ -153,7 +154,7 @@ module type Job = sig (** Given a job, [await_exit] will wait for the job to finish and return the exit based on the various options passed in. *) - val last_process : 'a t -> int option + val last_process : 'a t -> process option (** The last process in the job, if any. *) end diff --git a/test/nofork.t b/test/nofork.t new file mode 100644 index 0000000..71a028d --- /dev/null +++ b/test/nofork.t @@ -0,0 +1,27 @@ +These tests exercise edge-cases where a shell should really fork itself. +We try our best to avoid having to fork at all. + +The first challenge is ensuring that separate pipeline commands to run with +some amount of concurrency. To test this, we make sure that certain commands +do not block the rest of the pipeline from running. + + $ cat > test.sh << EOF + > while echo "hello"; do + > : + > done | head -n 5 + > EOF + + $ sh test.sh + hello + hello + hello + hello + hello + + $ msh test.sh + hello + hello + hello + hello + hello + Net Connection_reset Unix_error (Broken pipe, "writev", "") -- 2.51.2