[Libguestfs] [libnbd PATCH v4 07/11] rust: async: Create an async friendly handle type

Create another handle type: `AsyncHandle`, which makes use of Rust's builtin asynchronous functions (see <https://doc.rust-lang.org/std/keyword.async.html>) and runs on top of the Tokio runtime (see <https://docs.rs/tokio>). For every asynchronous command, like `aio_connect()`, a corresponding `async` method is created on the handle. In this case it would be: async fn connect(...) -> Result<(), ...> When called, it will poll the file descriptor until the command is complete, and then return with a result. All the synchronous counterparts (like `nbd_connect()`) are excluded from this handle type as they are unnecessary and since they might interfear with the polling made by the Tokio runtime. For more details about how the asynchronous commands are executed, please see the comments in rust/src/async_handle.rs. --- generator/Rust.ml | 240 +++++++++++++++++++++++++++++++++++ generator/Rust.mli | 2 + generator/generator.ml | 1 + rust/Cargo.toml | 7 +- rust/Makefile.am | 2 + rust/src/async_handle.rs | 268 +++++++++++++++++++++++++++++++++++++++ rust/src/lib.rs | 8 ++ 7 files changed, 527 insertions(+), 1 deletion(-) create mode 100644 rust/src/async_handle.rs diff --git a/generator/Rust.ml b/generator/Rust.ml index d3225eb..dc82d46 100644 --- a/generator/Rust.ml +++ b/generator/Rust.ml @@ -558,3 +558,243 @@ let generate_rust_bindings () = pr "impl Handle {\n"; List.iter print_rust_handle_method handle_calls; pr "}\n\n" + +(*********************************************************) +(* The rest of the file conserns the asynchronous API. *) +(* *) +(* See the comments in rust/src/async_handle.rs for more *) +(* information about how it works. *) +(*********************************************************) + +let excluded_handle_calls : NameSet.t = + NameSet.of_list + [ + "aio_get_fd"; + "aio_get_direction"; + "aio_notify_read"; + "aio_notify_write"; + "clear_debug_callback"; + "get_debug"; + "poll"; + "poll2"; + "set_debug"; + "set_debug_callback"; + ] + +(* A mapping with names as keys. *) +module NameMap = Map.Make (String) + +(* Strip "aio_" from the beginning of a string. *) +let strip_aio name : string = + if String.starts_with ~prefix:"aio_" name then + String.sub name 4 (String.length name - 4) + else failwithf "Asynchronous call %s must begin with aio_" name + +(* A map with all asynchronous handle calls. The keys are names with "aio_" + stripped, the values are a tuple with the actual name (with "aio_"), the + [call] and the [async_kind]. *) +let async_handle_calls : ((string * call) * async_kind) NameMap.t = + handle_calls + |> List.filter (fun (n, _) -> not (NameSet.mem n excluded_handle_calls)) + |> List.filter_map (fun (name, call) -> + call.async_kind + |> Option.map (fun async_kind -> + (strip_aio name, ((name, call), async_kind)))) + |> List.to_seq |> NameMap.of_seq + +(* A mapping with all synchronous (not asynchronous) handle calls. Excluded + are also all synchronous calls that has an asynchronous counterpart. So if + "foo" is the name of a handle call and an asynchronous call "aio_foo" + exists, then "foo" will not b in this map. *) +let sync_handle_calls : call NameMap.t = + handle_calls + |> List.filter (fun (n, _) -> not (NameSet.mem n excluded_handle_calls)) + |> List.filter (fun (name, _) -> + (not (NameMap.mem name async_handle_calls)) + && not + (String.starts_with ~prefix:"aio_" name + && NameMap.mem (strip_aio name) async_handle_calls)) + |> List.to_seq |> NameMap.of_seq + +(* Get the Rust type for an argument in the asynchronous API. Like + [rust_arg_type] but no static lifetime on some closures and buffers. *) +let rust_async_arg_type : arg -> string = function + | Closure { cbargs; cbcount; cblifetime } -> + let lifetime = + match cblifetime with CBCommand -> None | CBHandle -> Some "'static" + in + "impl " ^ rust_closure_trait ~lifetime cbargs cbcount + | BytesPersistIn _ -> "&[u8]" + | BytesPersistOut _ -> "&mut [u8]" + | x -> rust_arg_type x + +(* Get the Rust type for an optional argument in the asynchronous API. Like + [rust_optarg_type] but no static lifetime on some closures. *) +let rust_async_optarg_type : optarg -> string = function + | OClosure x -> sprintf "Option<%s>" (rust_async_arg_type (Closure x)) + | x -> rust_optarg_type x + +(* A string of the argument list for a method on the handle, with both + mandotory and optional arguments. *) +let rust_async_handle_call_args { args; optargs } : string = + let rust_args_names = + List.map rust_arg_name args @ List.map rust_optarg_name optargs + and rust_args_types = + List.map rust_async_arg_type args + @ List.map rust_async_optarg_type optargs + in + String.concat ", " + (List.map2 (sprintf "%s: %s") rust_args_names rust_args_types) + +(* Print the Rust function for a not asynchronous handle call. *) +let print_rust_sync_handle_call (name : string) (call : call) = + print_rust_handle_call_comment call; + pr "pub fn %s(&self, %s) -> %s\n" name + (rust_async_handle_call_args call) + (rust_ret_type call); + print_ffi_call name "self.data.handle.handle" call; + pr "\n" + +(* Print the Rust function for an asynchronous handle call with a completion + callback. (Note that "callback" might be abbreviated with "cb" in the + following code. *) +let print_rust_async_handle_call_with_completion_cb name (aio_name, call) = + (* An array of all optional arguments. Useful because we need to deel with + the index of the completion callback. *) + let optargs = Array.of_list call.optargs in + (* The index of the completion callback in [optargs] *) + let completion_cb_index = + Array.find_map + (fun (i, optarg) -> + match optarg with + | OClosure { cbname } -> + if cbname = "completion" then Some i else None + | _ -> None) + (Array.mapi (fun x y -> (x, y)) optargs) + in + let completion_cb_index = + match completion_cb_index with + | Some x -> x + | None -> + failwithf + "The handle call %s is claimed to have a completion callback among \ + its optional arguments by the async_kind field, but so does not \ + seem to be the case." + aio_name + in + let optargs_before_completion_cb = + Array.to_list (Array.sub optargs 0 completion_cb_index) + and optargs_after_completion_cb = + Array.to_list + (Array.sub optargs (completion_cb_index + 1) + (Array.length optargs - (completion_cb_index + 1))) + in + (* All optional arguments excluding the completion callback. *) + let optargs_without_completion_cb = + optargs_before_completion_cb @ optargs_after_completion_cb + in + print_rust_handle_call_comment call; + pr "pub async fn %s(&self, %s) -> SharedResult<()> {\n" name + (rust_async_handle_call_args + { call with optargs = optargs_without_completion_cb }); + pr " // A oneshot channel to notify when the call is completed.\n"; + pr " let (ret_tx, ret_rx) = oneshot::channel::<SharedResult<()>>();\n"; + pr " let (ccb_tx, mut ccb_rx) = oneshot::channel::<c_int>();\n"; + (* Completion callback: *) + pr " let %s = Some(|err: &mut i32| {\n" + (rust_optarg_name (Array.get optargs completion_cb_index)); + pr " ccb_tx.send(*err).ok();\n"; + pr " 1\n"; + pr " });\n"; + (* End of completion callback. *) + print_ffi_call aio_name "self.data.handle.handle" call; + pr "?;\n"; + pr " let mut ret_tx = Some(ret_tx);\n"; + pr " let completion_predicate = \n"; + pr " move |_handle: &Handle, res: &SharedResult<()>| {\n"; + pr " let ret = if res.as_ref().is_err_and(|e| e.is_fatal()) {\n"; + pr " res.clone()\n"; + pr " } else {\n"; + pr " let Ok(errno) = ccb_rx.try_recv() else { return false; };\n"; + pr " if errno == 0 {\n"; + pr " Ok(())\n"; + pr " } else {\n"; + pr " if let Err(e) = res {\n"; + pr " Err(e.clone())\n"; + pr " } else {\n"; + pr " Err(Arc::new("; + pr " Error::Recoverable(ErrorKind::from_errno(errno))))\n"; + pr " }\n"; + pr " }\n"; + pr " };\n"; + pr " ret_tx.take().unwrap().send(ret).ok();\n"; + pr " true\n"; + pr " };\n"; + pr " self.add_command(completion_predicate)?;\n"; + pr " ret_rx.await.unwrap()\n"; + pr "}\n\n" + +(* Print a Rust function for an asynchronous handle call which signals + completion by changing state. The predicate is a call like + "aio_is_connecting" which should get the value (like false) for the call to + be complete. *) +let print_rust_async_handle_call_changing_state name (aio_name, call) + (predicate, value) = + let value = if value then "true" else "false" in + print_rust_handle_call_comment call; + pr "pub async fn %s(&self, %s) -> SharedResult<()>\n" name + (rust_async_handle_call_args call); + pr "{\n"; + print_ffi_call aio_name "self.data.handle.handle" call; + pr "?;\n"; + pr " let (ret_tx, ret_rx) = oneshot::channel::<SharedResult<()>>();\n"; + pr " let mut ret_tx = Some(ret_tx);\n"; + pr " let completion_predicate = \n"; + pr " move |handle: &Handle, res: &SharedResult<()>| {\n"; + pr " let ret = if let Err(_) = res {\n"; + pr " res.clone()\n"; + pr " } else {\n"; + pr " if handle.%s() != %s { return false; }\n" predicate value; + pr " else { Ok(()) }\n"; + pr " };\n"; + pr " ret_tx.take().unwrap().send(ret).ok();\n"; + pr " true\n"; + pr " };\n"; + pr " self.add_command(completion_predicate)?;\n"; + pr " ret_rx.await.unwrap()\n"; + pr "}\n\n" + +(* Print an impl with all handle calls. *) +let print_rust_async_handle_impls () = + pr "impl AsyncHandle {\n"; + NameMap.iter print_rust_sync_handle_call sync_handle_calls; + NameMap.iter + (fun name (call, async_kind) -> + match async_kind with + | WithCompletionCallback -> + print_rust_async_handle_call_with_completion_cb name call + | ChangesState (predicate, value) -> + print_rust_async_handle_call_changing_state name call + (predicate, value)) + async_handle_calls; + pr "}\n\n" + +let print_rust_async_imports () = + pr "use crate::{*, types::*};\n"; + pr "use os_str_bytes::OsStringBytes as _;\n"; + pr "use os_socketaddr::OsSocketAddr;\n"; + pr "use std::ffi::*;\n"; + pr "use std::mem;\n"; + pr "use std::net::SocketAddr;\n"; + pr "use std::os::fd::{AsRawFd, OwnedFd};\n"; + pr "use std::path::PathBuf;\n"; + pr "use std::ptr;\n"; + pr "use std::sync::Arc;\n"; + pr "use tokio::sync::oneshot;\n"; + pr "\n" + +let generate_rust_async_bindings () = + generate_header CStyle ~copyright:"Tage Johansson"; + pr "\n"; + print_rust_async_imports (); + print_rust_async_handle_impls () diff --git a/generator/Rust.mli b/generator/Rust.mli index 450e4ca..0960170 100644 --- a/generator/Rust.mli +++ b/generator/Rust.mli @@ -18,3 +18,5 @@ (* Print all flag-structs, enums, constants and handle calls in Rust code. *) val generate_rust_bindings : unit -> unit + +val generate_rust_async_bindings : unit -> unit diff --git a/generator/generator.ml b/generator/generator.ml index f5ef7cc..2118446 100644 --- a/generator/generator.ml +++ b/generator/generator.ml @@ -64,3 +64,4 @@ let () = output_to ~formatter:(Some Rustfmt) "rust/libnbd-sys/src/lib.rs" RustSys.generate_rust_sys_bindings; output_to ~formatter:(Some Rustfmt) "rust/src/bindings.rs" Rust.generate_rust_bindings; + output_to ~formatter:(Some Rustfmt) "rust/src/async_bindings.rs" Rust.generate_rust_async_bindings; diff --git a/rust/Cargo.toml b/rust/Cargo.toml index f74c3ac..c7e461a 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -45,10 +45,15 @@ thiserror = "1.0.40" log = { version = "0.4.19", optional = true } libc = "0.2.147" byte-strings = "0.3.1" +tokio = { optional = true, version = "1.29.1", default-features = false, features = ["rt", "sync", "net"] } +epoll = "4.3.3" [features] -default = ["log"] +default = ["log", "tokio"] [dev-dependencies] once_cell = "1.18.0" +pretty-hex = "0.3.0" +rand = { version = "0.8.5", default-features = false, features = ["small_rng", "min_const_gen"] } tempfile = "3.6.0" +tokio = { version = "1.29.1", default-features = false, features = ["macros", "rt-multi-thread"] } diff --git a/rust/Makefile.am b/rust/Makefile.am index 19dbf02..b954b22 100644 --- a/rust/Makefile.am +++ b/rust/Makefile.am @@ -19,6 +19,7 @@ include $(top_srcdir)/subdir-rules.mk generator_built = \ libnbd-sys/src/lib.rs \ + src/async_bindings.rs \ src/bindings.rs \ $(NULL) @@ -30,6 +31,7 @@ source_files = \ src/handle.rs \ src/types.rs \ src/utils.rs \ + src/async_handle.rs \ libnbd-sys/Cargo.toml \ libnbd-sys/build.rs \ $(NULL) diff --git a/rust/src/async_handle.rs b/rust/src/async_handle.rs new file mode 100644 index 0000000..4223b80 --- /dev/null +++ b/rust/src/async_handle.rs @@ -0,0 +1,268 @@ +// nbd client library in userspace +// Copyright Tage Johansson +// +// This library is free software; you can redistribute it and/or +// modify it under the terms of the GNU Lesser General Public +// License as published by the Free Software Foundation; either +// version 2 of the License, or (at your option) any later version. +// +// This library is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +// Lesser General Public License for more details. +// +// You should have received a copy of the GNU Lesser General Public +// License along with this library; if not, write to the Free Software +// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA + +// This module implements an asynchronous handle working on top of the +// [Tokio](https://tokio.rs) runtime. When the handle is created, +// a "polling task" is spawned on the Tokio runtime. The purpose of that +// "polling task" is to call `aio_notify_*` when appropriate. It shares a +// reference to the handle as well as some other things with the handle in the +// [HandleData] struct. The "polling task" is sleeping when no command is in +// flight, but wakes up as soon as any command is issued. +// +// The commands are implemented as +// [`async fn`s](https://doc.rust-lang.org/std/keyword.async.html) +// in async_bindings.rs. When a new command is issued, it registers a +// completion predicate with [Handle::add_command]. That predicate takes a +// reference to the handle and should return [true] iff the command is complete. +// Whenever some work is performed in the polling task, the completion +// predicates for all pending commands are called. + +#![allow(unused_imports)] // XXX: remove this +use crate::sys; +use crate::Handle; +use crate::{Error, FatalErrorKind, Result}; +use crate::{AIO_DIRECTION_BOTH, AIO_DIRECTION_READ, AIO_DIRECTION_WRITE}; +use epoll::Events; +use std::sync::Arc; +use std::sync::Mutex; +use tokio::io::{unix::AsyncFd, Interest, Ready as IoReady}; +use tokio::sync::{broadcast, Notify}; +use tokio::task; + +/// A custom result type with a shared [crate::Error] as default error type. +pub type SharedResult<T, E = Arc<Error>> = Result<T, E>; + +/// An NBD handle using Rust's `async` functionality on top of the +/// [Tokio](https://docs.rs/tokio/) runtime. +pub struct AsyncHandle { + /// Data shared both by this struct and the polling task. + pub(crate) data: Arc<HandleData>, + + /// A task which soely purpose is to poll the NBD handle. + polling_task: tokio::task::AbortHandle, +} + +pub(crate) struct HandleData { + /// The underliing handle. + pub handle: Handle, + + /// A list of all pending commands. + /// + /// For every pending command (commands in flight), a predicate will be + /// stored in this list. Whenever some progress is made on the file + /// descriptor, the predicate is called with a reference to the handle + /// and a reference to the result of that call to `aio_notify_*`. + /// Iff the predicate returns [true], the command is considered completed + /// and removed from this list. + /// + /// If The polling task dies for some reason, this [SharedResult] will be + /// set to some error. + pub pending_commands: Mutex< + SharedResult< + Vec< + Box< + dyn FnMut(&Handle, &SharedResult<()>) -> bool + + Send + + Sync + + 'static, + >, + >, + >, + >, + + /// A notifier used by commands to notify the polling task when a new + /// asynchronous command is issued. + pub new_command: Notify, +} + +impl AsyncHandle { + pub fn new() -> Result<Self> { + let handle_data = Arc::new(HandleData { + handle: Handle::new()?, + pending_commands: Mutex::new(Ok(Vec::new())), + new_command: Notify::new(), + }); + + let handle_data_2 = handle_data.clone(); + let polling_task = task::spawn(async move { + // The polling task should never finish without an error. If the + // handle is dropped, the task is aborted so it'll not return in + // that case either. + let Err(err) = polling_task(&handle_data_2).await else { + unreachable!() + }; + let err = Arc::new(Error::Fatal(err)); + // Call the completion predicates for all pending commands with the + // error. + let mut pending_cmds = + handle_data_2.pending_commands.lock().unwrap(); + let res = Err(err); + for f in pending_cmds.as_mut().unwrap().iter_mut() { + f(&handle_data_2.handle, &res); + } + *pending_cmds = Err(res.unwrap_err()); + }) + .abort_handle(); + Ok(Self { + data: handle_data, + polling_task, + }) + } + + /// Get the underliing C pointer to the handle. + pub(crate) fn raw_handle(&self) -> *mut sys::nbd_handle { + self.data.handle.raw_handle() + } + + /// Call this method when a new command is issued. As argument is passed a + /// predicate which should return [true] iff the command is completed. + pub(crate) fn add_command( + &self, + mut completion_predicate: impl FnMut(&Handle, &SharedResult<()>) -> bool + + Send + + Sync + + 'static, + ) -> SharedResult<()> { + if !completion_predicate(&self.data.handle, &Ok(())) { + let mut pending_cmds_lock = + self.data.pending_commands.lock().unwrap(); + pending_cmds_lock + .as_mut() + .map_err(|e| e.clone())? + .push(Box::new(completion_predicate)); + self.data.new_command.notify_one(); + } + Ok(()) + } +} + +impl Drop for AsyncHandle { + fn drop(&mut self) { + self.polling_task.abort(); + } +} + +/// Get the read/write direction that the handle wants on the file descriptor. +fn get_fd_interest(handle: &Handle) -> Option<Interest> { + match handle.aio_get_direction() { + 0 => None, + AIO_DIRECTION_READ => Some(Interest::READABLE), + AIO_DIRECTION_WRITE => Some(Interest::WRITABLE), + AIO_DIRECTION_BOTH => Some(Interest::READABLE | Interest::WRITABLE), + _ => unreachable!(), + } +} + +/// A task that will run as long as the handle is alive. It will poll the +/// file descriptor when new data is availlable. +async fn polling_task(handle_data: &HandleData) -> Result<(), FatalErrorKind> { + let HandleData { + handle, + pending_commands, + new_command, + } = handle_data; + let fd = handle.aio_get_fd().map_err(Error::to_fatal)?; + // XXX: Might the file descriptor ever be changed? + let tokio_fd = AsyncFd::new(fd)?; + let epfd = epoll::create(false)?; + epoll::ctl( + epfd, + epoll::ControlOptions::EPOLL_CTL_ADD, + fd, + epoll::Event::new(Events::EPOLLIN | Events::EPOLLOUT, 42), + )?; + + // The following loop does approximately the following things: + // + // 1. Determine what Libnbd wants to do next on the file descriptor, + // (read/write/both/none), and store that in [fd_interest]. + // 2. Wait for either: + // a) That interest to be available on the file descriptor in which case: + // I. Call the correct `aio_notify_*` method. + // II. Execute step 1. + // III. Send the result of the call to `aio_notify_*` on + // [result_channel] to notify pending commands that some progress + // has been made. + // IV. Resume execution from step 2. + // b) A notification was received on [new_command] signaling that a new + // command was registered and that the intrest on the file descriptor + // might has changed. Resume execution from step 1. + loop { + let Some(fd_interest) = get_fd_interest(handle) else { + // The handle does not wait for any data of the file descriptor, + // so we wait until some command is issued. + new_command.notified().await; + continue; + }; + + if pending_commands + .lock() + .unwrap() + .as_ref() + .unwrap() + .is_empty() + { + // No command is pending so there is no point to do anything. + new_command.notified().await; + continue; + } + + // Wait for the requested interest to be available on the fd. + let mut ready_guard = tokio_fd.ready(fd_interest).await?; + let readyness = ready_guard.ready(); + let res = if readyness.is_readable() && fd_interest.is_readable() { + handle.aio_notify_read() + } else if readyness.is_writable() && fd_interest.is_writable() { + handle.aio_notify_write() + } else { + continue; + }; + let res = match res { + Ok(()) => Ok(()), + Err(e @ Error::Recoverable(_)) => Err(Arc::new(e)), + Err(Error::Fatal(e)) => return Err(e), + }; + + // Call the completion predicates of all pending commands. + let mut pending_cmds_lock = pending_commands.lock().unwrap(); + let pending_cmds = pending_cmds_lock.as_mut().unwrap(); + let mut i = 0; + while i < pending_cmds.len() { + if (pending_cmds[i])(handle, &res) { + let _ = pending_cmds.swap_remove(i); + } else { + i += 1; + } + } + drop(pending_cmds_lock); + + // Use epoll to check the current read/write availability on the fd. + // This is needed because Tokio does only support edge-triggered + // notifications but Libnbd requires level-triggered notifications. + let mut revent = epoll::Event { data: 0, events: 0 }; + // Setting timeout to 0 means that it will return immediately. + epoll::wait(epfd, 0, std::slice::from_mut(&mut revent))?; + let revents = Events::from_bits(revent.events).unwrap(); + if !revents.contains(Events::EPOLLIN) { + ready_guard.clear_ready_matching(IoReady::READABLE); + } + if !revents.contains(Events::EPOLLOUT) { + ready_guard.clear_ready_matching(IoReady::WRITABLE); + } + ready_guard.retain_ready(); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index a6f3131..56316b4 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -17,11 +17,19 @@ #![deny(warnings)] +#[cfg(feature = "tokio")] +mod async_bindings; +#[cfg(feature = "tokio")] +mod async_handle; mod bindings; mod error; mod handle; pub mod types; mod utils; +#[cfg(feature = "tokio")] +pub use async_bindings::*; +#[cfg(feature = "tokio")] +pub use async_handle::{AsyncHandle, SharedResult}; pub use bindings::*; pub use error::{Error, ErrorKind, FatalErrorKind, Result}; pub use handle::Handle; -- 2.41.0