rublk: support ublk_loop bpf target

Signed-off-by: Ming Lei <[email protected]>

src/bpf/bpf_aio_kfunc.h

@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: GPL-2.0
+#ifndef UBLK_BPF_AIO_INTERNAL_H
+#define UBLK_BPF_AIO_INTERNAL_H
+
+extern struct bpf_aio *bpf_aio_alloc(unsigned int op, unsigned int flags) __ksym;
+extern struct bpf_aio *bpf_aio_alloc_sleepable(unsigned int op, unsigned int flags) __ksym;
+extern void bpf_aio_release(struct bpf_aio *aio) __ksym;
+extern int bpf_aio_submit(struct bpf_aio *aio, int fd, loff_t pos,
+                unsigned bytes, unsigned io_flags) __ksym;
+#endif

src/bpf/ublk_bpf.h

@@ -3,6 +3,7 @@
 #define UBLK_BPF_GEN_H
 
 #include "ublk_bpf_kfunc.h"
+#include "bpf_aio_kfunc.h"
 
 #ifdef DEBUG
 #define BPF_DBG(...) bpf_printk(__VA_ARGS__)

src/bpf/ublk_loop.bpf.c

@@ -0,0 +1,149 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include "vmlinux.h"
+#include <linux/const.h>
+#include <linux/errno.h>
+#include <linux/falloc.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+//#define DEBUG
+#include "ublk_bpf.h"
+
+/* libbpf v1.4.5 is required for struct_ops to work */
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, 1);
+	__type(key, int);
+	__type(value, int);
+} fd_map SEC(".maps");
+
+static inline void ublk_loop_comp_and_release_aio(struct bpf_aio *aio, int ret)
+{
+	ublk_bpf_dettach_and_complete_aio(aio, ret);
+	bpf_aio_release(aio);
+}
+
+SEC("struct_ops/bpf_aio_complete_cb")
+void BPF_PROG(ublk_loop_comp_cb, struct bpf_aio *aio)
+{
+	BPF_DBG("aio result %d, back_file %s pos %llx", aio->ret,
+			aio->iocb.ki_filp->f_path.dentry->d_name.name,
+			aio->iocb.ki_pos - aio->ret);
+	ublk_loop_comp_and_release_aio(aio, aio->ret);
+}
+
+SEC(".struct_ops.link")
+struct bpf_aio_complete_ops loop_ublk_bpf_aio_ops = {
+	.id = -1,
+	.bpf_aio_complete_cb = (void *)ublk_loop_comp_cb,
+};
+
+static inline int ublk_loop_submit_backing_io(const struct ublk_bpf_io *io,
+		const struct ublksrv_io_desc *iod, int backing_fd)
+{
+	unsigned int op_flags = 0;
+	struct bpf_aio *aio;
+	int res = -EINVAL;
+	int op;
+
+	/* translate ublk opcode into backing file's */
+	switch (iod->op_flags & 0xff) {
+	case 0 /*UBLK_IO_OP_READ*/:
+		op = BPF_AIO_OP_FS_READ;
+		break;
+	case 1 /*UBLK_IO_OP_WRITE*/:
+		op = BPF_AIO_OP_FS_WRITE;
+		break;
+	case 2 /*UBLK_IO_OP_FLUSH*/:
+		op = BPF_AIO_OP_FS_FSYNC;
+		break;
+	case 3 /*UBLK_IO_OP_DISCARD*/:
+		op = BPF_AIO_OP_FS_FALLOCATE;
+		op_flags = FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE;
+		break;
+	case 4 /*UBLK_IO_OP_WRITE_SAME*/:
+		op = BPF_AIO_OP_FS_FALLOCATE;
+		op_flags = FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE;
+		break;
+	case 5 /*UBLK_IO_OP_WRITE_ZEROES*/:
+		op = BPF_AIO_OP_FS_FALLOCATE;
+		op_flags = FALLOC_FL_ZERO_RANGE | FALLOC_FL_KEEP_SIZE;
+		break;
+	default:
+		return -EINVAL;
+	}
+
+	res = -ENOMEM;
+	aio = bpf_aio_alloc(op, 0);
+	if (!aio)
+		goto fail;
+
+	/* attach aio into the specified range of this io command */
+	res = ublk_bpf_attach_and_prep_aio(io, 0, iod->nr_sectors << 9, aio);
+	if (res < 0) {
+		bpf_printk("bpf aio attaching failed %d\n", res);
+		goto fail;
+	}
+
+	/* submit this aio onto the backing file */
+	res = bpf_aio_submit(aio, backing_fd, iod->start_sector << 9,
+			iod->nr_sectors << 9, op_flags);
+	if (res < 0) {
+		bpf_printk("aio submit failed %d\n", res);
+		ublk_loop_comp_and_release_aio(aio, res);
+	}
+	return 0;
+fail:
+	return res;
+}
+
+static inline int __ublk_loop_handle_io_cmd(const struct ublk_bpf_io *io)
+{
+	const struct ublksrv_io_desc *iod;
+	int res = -EINVAL;
+	int fd_key = 0;
+	int *fd;
+
+	iod = ublk_bpf_get_iod(io);
+	if (!iod) {
+		ublk_bpf_complete_io(io, res);
+		return UBLK_BPF_IO_HANDLED;
+	}
+
+	BPF_DBG("ublk dev %u qid %u: handle io cmd tag %u %lx-%d",
+			ublk_bpf_get_dev_id(io),
+			ublk_bpf_get_queue_id(io),
+			ublk_bpf_get_io_tag(io),
+			iod->start_sector << 9,
+			iod->nr_sectors << 9);
+
+	/* retrieve backing file descriptor */
+	fd = bpf_map_lookup_elem(&fd_map, &fd_key);
+	if (!fd)
+		goto exit;
+
+	/* handle this io command by submitting IOs on backing file */
+	res = ublk_loop_submit_backing_io(io, iod, *fd);
+
+exit:
+	/* io cmd can't be completes until this reference is dropped */
+	ublk_bpf_io_dec_ref(io, res);
+
+	return UBLK_BPF_IO_HANDLED;
+}
+
+SEC("struct_ops.s/ublk_bpf_queue_io_cmd_daemon")
+int BPF_PROG(ublk_loop_handle_io_cmd, struct ublk_bpf_io *io)
+{
+	return __ublk_loop_handle_io_cmd(io);
+}
+
+SEC(".struct_ops.link")
+struct ublk_bpf_ops loop_ublk_bpf_ops = {
+	.dev_id = -1,
+	.queue_io_cmd_daemon = (void *)ublk_loop_handle_io_cmd,
+};
+
+char LICENSE[] SEC("license") = "GPL";

src/loop.rs

@@ -8,6 +8,23 @@ use std::os::unix::io::AsRawFd;
 use std::path::PathBuf;
 use std::rc::Rc;
 
+use anyhow::Result;
+use libbpf_rs::skel::OpenSkel;
+use libbpf_rs::skel::Skel;
+use libbpf_rs::skel::SkelBuilder;
+use libbpf_rs::MapCore as _;
+use libbpf_rs::MapFlags;
+use std::mem::MaybeUninit;
+
+mod loopbpf {
+    include!(concat!(
+        env!("CARGO_MANIFEST_DIR"),
+        "/src/bpf/ublk_loop.skel.rs"
+    ));
+}
+
+use loopbpf::*;
+
 #[derive(clap::Args, Debug)]
 pub struct LoopArgs {
     #[command(flatten)]
@@ -191,6 +208,18 @@ fn lo_handle_io_cmd_sync(q: &UblkQueue<'_>, tag: u16, i: &UblkIOCtx, buf_addr: *
     }
 }
 
+fn q_bpf_fn(qid: u16, dev: &UblkDev) {
+    let lo_io_handler = move |q: &UblkQueue, tag: u16, io: &UblkIOCtx| {
+        lo_handle_io_cmd_sync(q, tag, io, std::ptr::null_mut());
+    };
+
+    UblkQueue::new(qid, dev)
+        .unwrap()
+        .regiser_io_bufs(None)
+        .submit_fetch_commands(None)
+        .wait_and_handle_io(lo_io_handler);
+}
+
 fn q_fn(qid: u16, dev: &UblkDev) {
     let bufs_rc = Rc::new(dev.alloc_queue_io_bufs());
     let bufs = bufs_rc.clone();
@@ -238,7 +267,11 @@ fn q_a_fn(qid: u16, dev: &UblkDev) {
     smol::block_on(async { futures::future::join_all(f_vec).await });
 }
 
-pub(crate) fn ublk_add_loop(ctrl: UblkCtrl, opt: Option<LoopArgs>) -> Result<i32, UblkError> {
+fn __ublk_add_loop(
+    ctrl: UblkCtrl,
+    opt: Option<LoopArgs>,
+    bpf_skel: Option<UblkLoopSkel>,
+) -> Result<i32, UblkError> {
     let (file, dio, ro, aa, _shm, fg) = match opt {
         Some(ref o) => {
             let parent = o.gen_arg.get_start_dir();
@@ -290,16 +323,72 @@ pub(crate) fn ublk_add_loop(ctrl: UblkCtrl, opt: Option<LoopArgs>) -> Result<i32
         async_await: aa,
     };
 
+    let bpf;
+    if let Some(skel) = bpf_skel {
+        let key = (0_i32).to_ne_bytes();
+        let val = lo.back_file.as_raw_fd().to_ne_bytes();
+        let _ = skel.maps.fd_map.update(&key, &val, MapFlags::ANY);
+        bpf = true;
+    } else {
+        bpf = false;
+    }
+
     //todo: USER_COPY should be the default option
-    if (ctrl.dev_info().flags & (libublk::sys::UBLK_F_USER_COPY as u64)) != 0 {
+    if !bpf && (ctrl.dev_info().flags & (libublk::sys::UBLK_F_USER_COPY as u64)) != 0 {
+        log::error!("wrong user_copy flag");
         return Err(UblkError::OtherError(-libc::EINVAL));
     }
 
     ctrl.run_target(
         |dev: &mut UblkDev| lo_init_tgt(dev, &lo, opt, dio),
-        move |qid, dev: &_| if aa { q_a_fn(qid, dev) } else { q_fn(qid, dev) },
+        move |qid, dev: &_| {
+            if bpf {
+                q_bpf_fn(qid, dev)
+            } else if aa {
+                q_a_fn(qid, dev)
+            } else {
+                q_fn(qid, dev)
+            }
+        },
         move |ctrl: &UblkCtrl| crate::rublk_prep_dump_dev(_shm, fg, ctrl),
     )
     .unwrap();
     Ok(0)
 }
+
+fn __ublk_add_loop_bpf(ctrl: UblkCtrl, opt: Option<LoopArgs>) -> Result<()> {
+    let skel_builder = UblkLoopSkelBuilder::default();
+    let mut open_object = MaybeUninit::uninit();
+    let mut open_skel = skel_builder.open(&mut open_object)?;
+    let dev_id = ctrl.dev_info().dev_id as i32;
+
+    open_skel.struct_ops.loop_ublk_bpf_ops_mut().dev_id = dev_id;
+    open_skel.struct_ops.loop_ublk_bpf_aio_ops_mut().id = dev_id;
+
+    let mut skel = open_skel.load()?;
+    let _link = skel.maps.loop_ublk_bpf_ops.attach_struct_ops()?;
+    let _aio_link = skel.maps.loop_ublk_bpf_aio_ops.attach_struct_ops()?;
+
+    skel.attach()?;
+
+    __ublk_add_loop(ctrl, opt, Some(skel))?;
+
+    Ok(())
+}
+
+pub(crate) fn ublk_add_loop(ctrl: UblkCtrl, opt: Option<LoopArgs>) -> Result<i32, UblkError> {
+    let info = ctrl.dev_info();
+
+    if (info.flags & (libublk::sys::UBLK_F_BPF as u64)) != 0 {
+        let res = __ublk_add_loop_bpf(ctrl, opt);
+        match res {
+            Ok(_) => Ok(0),
+            _ => {
+                log::error!("add loop bpf failed");
+                Err(UblkError::OtherError(-libc::EINVAL))
+            }
+        }
+    } else {
+        __ublk_add_loop(ctrl, opt, None)
+    }
+}

tests/basic.rs

@@ -247,7 +247,7 @@ mod integration {
         }
     }
 
-    fn __test_ublk_add_del_loop(bs: u32, aa: bool) {
+    fn __test_ublk_add_del_loop(bs: u32, aa: bool, bpf: bool) {
         let tmp_file = tempfile::NamedTempFile::new().unwrap();
         let file_size = 32 * 1024 * 1024; // 1 MB
         let p = tmp_file.path();
@@ -263,6 +263,9 @@ mod integration {
         if aa {
             cmd_line.push("-a");
         }
+        if bpf {
+            cmd_line.push("--bpf");
+        }
         let id = run_rublk_add_dev(cmd_line);
 
         let ctrl = UblkCtrl::new_simple(id).unwrap();
@@ -275,8 +278,20 @@ mod integration {
         if !support_ublk() {
             return;
         }
-        __test_ublk_add_del_loop(4096, false);
-        __test_ublk_add_del_loop(4096, true);
+        let bs = 4096;
+        let mut aa = false;
+        let mut bpf = false;
+
+        __test_ublk_add_del_loop(bs, aa, bpf);
+
+        aa = true;
+        __test_ublk_add_del_loop(bs, aa, bpf);
+
+        bpf = true;
+        __test_ublk_add_del_loop(bs, aa, bpf);
+
+        aa = false;
+        __test_ublk_add_del_loop(bs, aa, bpf);
     }
 
     fn __test_ublk_null_read_only(cmds: &[&str], exp_ro: bool) {