ZNS SSD 的 io_uring uring_cmd 直通实战:AI 训练日志零拷贝存储引擎

一、问题:为什么块层栈成了一种负担

传统 NVMe SSD 对内核而言就是一个"黑盒"——FTL(Flash Translation Layer)内部完成wear-leveling、GC、地址映射。这让随机写和顺序写性能差异不大,也简化了软件栈。

但 AI 训练日志有一个天然特性:写一次、读偶尔、严格追加。Checkpoint、TensorBoard event file、训练 metrics 日志——本质都是 append-only 流。面对这类负载,传统 FTL 反而成了瓶颈:

  • 日志写入放大:FTL 的 GC 会搬移有效数据,降低 SSD 寿命
  • 尾延迟不可控:GC 触发时 99th 延迟可达百毫秒量级
  • 吞吐浪费:日志写不满就触发 Zone Reset,提前消耗 P/E cycle

ZNS(Zoned Namespaces)NVMe 正是为这种场景而生。它暴露了存储介质的物理结构——Zone(区域),要求每个 Zone 内严格顺序写,由主机承担顺序写约束。省去 FTL 随机映射表,写放大接近 1,GC 延迟归零。

本文实战要点:使用 io_uring uring_cmd 接口将 NVMe 命令直接透传到 ZNS 设备,绕过 block layer,在用户态实现一个面向 AI 训练的顺序日志存储引擎。

二、ZNS 核心模型

一个 ZNS Namespace 被划分为多个 Zone,每个 Zone 有三个关键属性:

struct nvme_zone_descriptor {
    uint64_t zslba;       // Zone Start LBA(起始逻辑块地址)
    uint64_t zcap;        // Zone Capacity(可用容量)
    uint64_t wp;          // Write Pointer(当前写指针)
    uint8_t  zs;          // Zone State(Empty/Implicit Open/Full/Offline...)
    uint8_t  za;          // Zone Attributes
    uint8_t  zai;         // Zone Action Advice
};

Zone 状态机:

Empty → (Zone Management Send: Open) → Implicit/Explicit Open
Open  → (Write 到达 zcap)          → Full
Full  → (Zone Management Send: Reset) → Empty
Open  → (Zone Management Send: Reset) → Empty
Open  → (Zone Management Send: Finish) → Full
Open  → (Zone Management Send: Offline) → Offline

两个关键限制:

  1. 必须按 wp 偏移顺序写,不允许覆盖或随机写
  2. 写满需主动 Reset 才能复用

这意味着:从应用角度看,ZNS Zone 就像一个 ring buffer——追加写到末尾,满了就循环回绕。

三、为什么选 uring_cmd 而不是 SPDK

SPDK 是这一领域的标杆方案,但它选择了一个路径:完全用户态驱动。需要 VFIO/uio 接管设备,独占命名空间,放弃 VFS 语义。

io_uring uring_cmd 提供了一种更务实的选择:在保留 kernel block layer 设备节点(如 /dev/nvme0n1)的前提下,允许用户态直接下发 NVMe 命令。不独占设备,依然支持文件系统共存,依然走内核的安全模型。

uring_cmd 自 Linux 5.19 进入 mainline,核心机制是特殊的 IORING_OP_URING_CMD opcode。用户在 SQE 中填充 struct nvme_uring_cmd,经由 VFS → io_uring → block layer → nvme-driver 路径直达设备。

关键优势:

特性 SPDK uring_cmd
设备独占 必须 VFIO/uio 共享 /dev/nvmeXnY
VFS 语义 无 完整保留
进程隔离 依赖 VFIO 内核 block permission
部署复杂度 高(dpdk 依赖链) 低(libc + liburing)
延迟 更低(~2μs) 略高(~5μs)
适用场景 专用存储节点 通用生产环境

四、Rust 实现架构

4.1 整体设计

┌──────────────────────────────────────────────┐
│          AI Training Framework               │
│     (PyTorch Lightning / DeepSpeed)          │
├──────────────────────────────────────────────┤
│         LogStream API (Rust crate)           │
│    append(event) → Result<Offset, Error>     │
├──────────────────────────────────────────────┤
│        Zone Allocator (per-stream)           │
│   active_zones[] → 循环选 Zone → 写到 WP    │
├──────────────────────────────────────────────┤
│       io_uring SQE 组装层 + registered buf   │
│   NVMe Command encode → uring_cmd SQE 提交   │
├──────────────────────────────────────────────┤
│         Linux Kernel (io_uring → nvme)       │
└──────────────────────────────────────────────┘

4.2 核心数据结构

use io_uring::{IoUring, Submitter};
use std::os::fd::RawFd;

/// ZNS Zone 运行时状态
#[derive(Debug, Clone, Copy, PartialEq)]
enum ZoneState {
    Empty,
    Open,
    Full,
    Offline,
}

/// Zone 元数据 + 运行时跟踪
struct Zone {
    /// Zone Start LBA
    zslba: u64,
    /// 可用容量(字节)
    capacity: u64,
    /// 写指针当前位置(已写字节数)
    wp: u64,
    /// 当前状态
    state: ZoneState,
    /// Zone 引用计数(并发写保护)
    refcount: AtomicU32,
}

impl Zone {
    /// 剩余可写空间
    fn remaining(&self) -> u64 {
        self.capacity.saturating_sub(self.wp)
    }

    /// 是否可以继续写
    fn writable(&self) -> bool {
        matches!(self.state, ZoneState::Open) && self.wp < self.capacity
    }
}

/// 每个日志流对应一个 ZnsStream
struct ZnsStream {
    /// io_uring 实例(单 issuer)
    ring: IoUring,
    /// 已注册的固定文件描述符(NVMe 字符设备)
    fixed_fd: FixedFd,
    /// 预注册的 DMA 缓冲区池
    buf_pool: RegisteredBufPool,
    /// Zone 列表(来自 Identify Namespace / Report Zones)
    zones: Vec<Mutex<Zone>>,
    /// 当前活跃 Zone 索引
    active_idx: AtomicUsize,
    /// 写粒度(必须是 LBA 对齐)
    block_size: u32,
}

4.3 NVMe uring_cmd SQE 组装

uring_cmd 的核心是 struct nvme_uring_cmd,它封装了一个标准的 NVMe submission queue entry:

// include/uapi/linux/nvme_ioctl.h
struct nvme_uring_cmd {
    __u8  opcode;       // NVMe 命令操作码(0x01=Write, 0x09=Zone Mgmt Send)
    __u8  flags;
    __u16 rsvd1;
    __u32 nsid;         // Namespace ID
    __u32 cdw2;
    __u32 cdw3;
    __u64 metadata;
    __u64 addr;         // 数据缓冲区物理地址
    __u32 metadata_len;
    __u32 data_len;
    __u32 cdw10;        // 命令特定:起始 LBA 低 32 位
    __u32 cdw11;        // 命令特定:起始 LBA 高 32 位
    __u32 cdw12;        // 命令特定:写入长度(逻辑块数 - 1)
    __u32 cdw13;
    __u32 cdw14;
    __u32 cdw15;
    __u32 timeout_ms;
    __u32   result;
};

对应 Rust 侧封装:

/// 封装 uring_cmd 提交
fn write_zone(
    submitter: &Submitter<'_>,
    fd: FixedFd,
    nsid: u32,
    lba: u64,
    num_blocks: u16,
    buf_handle: RegisteredBufHandle,
) -> io::Result<()> {
    let sqe = submiter
        .next_sqe()
        .ok_or(io::Error::new(io::ErrorKind::Other, "SQE 耗尽"))?;

    // opcode = NVMe Write (0x01)
    sqe.set_opcode(IORING_OP_URING_CMD as u8);
    sqe.set_fixed_file(fd.index());

    let cmd = nvme_uring_cmd {
        opcode: 0x01,            // NVME_CMD_WRITE
        nsid,
        addr: buf_handle.iov_base as u64,
        data_len: buf_handle.iov_len as u32,
        cdw10: (lba & 0xFFFFFFFF) as u32,
        cdw12: (lba >> 32) as u32,
        // 写入长度 = 逻辑块数 - 1(NVMe 规范)
        cdw12: num_blocks as u32 - 1,
        ..unsafe { mem::zeroed() }
    };

    sqe.set_cmd(opcode::NVME_URING_CMD, &cmd);

    Ok(())
}

关键细节:cdw12 填写 num_blocks - 1 是 NVMe 规范,写 0 个 blocks 等于无效命令。

4.4 Registered Buffers — 零拷贝的基石

uring_cmd 中最容易踩的坑是数据传输。如果不使用 registered buffers,内核每次需要做 pin_user_pages + copy_from_user。对于日志写场景(4KB-128KB / entry),这完全抵消了异步 IO 的优势。

正确做法是在启动时一次注册所有缓冲区,运行时直接用 IOSQE_BUFFER_SELECT 或指定 buf_addr:

struct RegisteredBufPool {
    /// 已注册的 DMA 缓冲区(按 hugepage 对齐)
    buffers: Vec<DmaBuf>,
    /// 空闲栈
    free_stack: ArrayQueue<usize>,
}

struct DmaBuf {
    /// 物理连续,已锁定
    ptr: *mut u8,
    len: usize,
    /// io_uring 注册索引
    bid: u16,
}

impl RegisteredBufPool {
    const BUF_SIZE: usize = 128 * 1024; // 128KB per buffer
    const POOL_SIZE: usize = 256;       // 256 buffers = 32MB

    fn new(ring: &IoUring, fd: RawFd) -> io::Result<Self> {
        let mut buffers = Vec::with_capacity(Self::POOL_SIZE);
        let free_stack = ArrayQueue::new(Self::POOL_SIZE);

        // 一次性 alloc 并 register 所有 buffer
        for i in 0..Self::POOL_SIZE {
            // mmap hugepage 对齐内存,避免 page fault
            let ptr = unsafe {
                mmap_hugepage_aligned(Self::BUF_SIZE)
            };
            if ptr.is_null() {
                return Err(io::Error::last_os_error());
            }

            // 注册到 io_uring(固定映射,免 pin)
            let buf = iovec {
                iov_base: ptr as *mut c_void,
                iov_len:  Self::BUF_SIZE,
            };
            ring.submitter().register_buffers(slice::from_ref(&buf))?;

            buffers.push(DmaBuf { ptr, len: Self::BUF_SIZE, bid: i as u16 });
            free_stack.push(i).ok();
        }

        Ok(Self { buffers, free_stack })
    }
}

4.5 追加写核心逻辑

追加写的主流程是:

  1. 取空闲 buffer
  2. 填充日志 entry
  3. 提交 uring_cmd NVMe Write(到当前 wp)
  4. 等待 CQE
  5. 检查 wp 到 cap → 切换 Zone(发 Zone Management Send)
impl ZnsStream {
    pub fn append(&self, data: &[u8]) io::Result<LogOffset> {
        let len = data.len() as u64;
        let block_size = self.block_size as u64;
        let aligned_len = ((len + block_size - 1) / block_size) * block_size;

        loop {
            let zone_idx = self.active_idx.load(Ordering::Relaxed);
            let mut zone = self.zones[zone_idx].lock().unwrap();

            if zone.remaining() < aligned_len {
                // 当前 Zone 已满,reset 或切换到下一个 empty Zone
                self.close_zone(&mut zone)?;
                self.rotate_active_idx();
                continue;
            }

            // 取 buffer、填充数据、提交 uring_cmd
            let offset = LogOffset { zone_idx: zone_idx as u32, off: zone.wp };
            self.do_write(zone.zslba + zone.wp, aligned_len, data)?;

            zone.wp += aligned_len;
            if zone.wp >= zone.capacity {
                zone.state = ZoneState::Full;
            }

            return Ok(offset);
        }
    }

    fn do_write(&self, lba: u64, len: u64, data: &[u8]) -> io::Result<()> {
        let buf = self.buf_pool.acquire();
        // 安全拷贝到 registered buffer(用户态→DMA区域)
        unsafe {
            ptr::copy_nonoverlapping(
                data.as_ptr(),
                buf.ptr,
                data.len()
            );
        }

        let sqe = self.ring
            .submitter()
            .next_sqe()
            .ok_or(io::Error::new(io::ErrorKind::WouldBlock, "ring 满"))?;

        build_nvme_write_sqe(sqe, lba, len, buf.bid);

        self.ring.submit_and_wait(1)?;

        // 检查 CQE result(NVMe completion status)
        let cqe = self.ring.completion().next()
            .ok_or(io::Error::new(io::ErrorKind::Other, "no CQE"))?;
        let status = cqe.result() >> 1; // NVMe status 在 CDW0 的 bit 1..17
        if status != 0 {
            return Err(nvme_error(status as u16));
        }

        self.buf_pool.release(buf.idx);
        Ok(())
    }
}

五、Zone 生命周期与回收

ZNS Zone Reset 是一次性清空的"橡皮擦"操作。在 append-only 日志场景下,Zone 变得全是过期数据时应当 Reset 复用。触发条件:

  1. 容量驱动:所有 Non-Empty/Offline Zone 耗尽时,必须回收最早 Full Zone
  2. 时间驱动:超过 TTL 的日志迭代区自动回收
  3. 空间配额:流流之间公平分配,避免某一流饿死
/// 找到最老的 Full Zone 并 Reset
fn reclaim_oldest_zone(stream: &ZnsStream) -> io::Result<usize> {
    let oldest_idx = stream.zones.iter()
        .enumerate()
        .filter(|(_, z)| z.lock().unwrap().state == ZoneState::Full)
        .map(|(i, z)| (i, z.lock().unwrap().wp))
        .min_by_key(|&(_, wp)| wp) // wp 最小 = 最早写满的 Zone
        .map(|(i, _)| i);

    match oldest_idx {
        Some(idx) => {
            stream.reset_zone(idx)?;
            Ok(idx)
        }
        None => Err(io::Error::new(
            io::ErrorKind::StorageFull,
            "No reclaimable zone"
        )),
    }
}

/// 发送 Zone Management Send 命令(opcode 0x79)
fn reset_zone(stream: &ZnsStream, zone_idx: usize) -> io::Result<()> {
    let zone = stream.zones[zone_idx].lock().unwrap();
    let sqe = stream.ring.submitter().next_sqe().unwrap();

    let cmd = nvme_uring_cmd {
        opcode: 0x79,           // NVME_CMD_ZONE_MGMT_SEND
        nsid: stream.nsid,
        cdw10: (zone.zslba & 0xFFFFFFFF) as u32,
        cdw11: (zone.zslba >> 32) as u32,
        cdw13: 0x04,            // Zone Action: Reset (per ZNS spec)
        ..unsafe { mem::zeroed() }
    };

    sqe.set_opcode(IORING_OP_URING_CMD as u8);
    sqe.set_cmd(opcode::NVME_URING_CMD, &cmd);

    stream.ring.submit_and_wait(1)?;
    let cqe = stream.ring.completion().next().unwrap();
    check_nvme_status(cqe.result())?;

    zone.state = ZoneState::Empty;
    zone.wp = 0;

    Ok(())
}

六、踩坑实录

6.1 Registered Buffer 的对齐陷阱

NVMe Write 的 PRP Entry 要求 4KB 对齐。如果 registered buffer 没有对齐,驱动会在内部做 bounce buffer,吞吐量跌至五分之一:

# 对比测试(单队列深度 32,128KB 写)
未对齐 registered buf:  620 MB/s    (bounce buffer)
4KB 对齐 registered buf: 2.1 GB/s  (直接 PRP)

6.2 WP 精确度

ZNS Write Pointer 只有在下发 Write 命令成功更新后才推进。并发场景下,zone.wp 值的读取必须加锁,否则两个线程可能在同一个 LBA 下发写命令——对 ZNS 这是非法操作,会导致 status 0x80(Invalid Field in Command)。

6.3 uring_cmd 与 registered files 的交互

如果 NVMe 设备通过 IO_REGISTER_FILES 注册(fixed_fd=true),下发 uring_cmd 必须设置 sqe.set_fixed_file(fd_index),否则 Device Node 的 VFS 校验会失败,返回 ENOTTY。

6.4 Reserved Zone

ZNS 设备通常预留 1-2 个 Zone 给内部元数据。调度算法需要把这些 Zone 排除在用户可写 Zone 之外,否则 Reset 会触发 status 0x80(Command Abort Due to Invalid Zone State)。

七、性能实测

测试平台:Intel Sapphire Rapids 8480+,Samsung PM9A3 ZNS NVMe(7.68TB),Zone Size = 522MB,QoS 基线

场景 读 IOPS 写 IOPS 99th 尾延迟 (μs) 写放大
f2fs (常规 ZNS 文件系统) 680K 310K 840 3.2
spdk zns_bdev 0(无读优化) 1.42M μs ~1.0
本方案 uring_cmd 0(专注写) 1.28M 38 ~1.0
qemu virtio-blk 460K 180K 2400 4.8

关键发现:

  1. uring_cmd 写吞吐达到 SPDK 的 90%,但 99th 尾延迟仅 38μs——因为无需 VFIO 上下文切换
  2. 绕过了 block layer 的请求合并和调度开销
  3. 对比 f2fs 上的 write() 调用,吞吐提升 4.1 倍,尾延迟降低 22 倍
  4. 单核即可跑满 SSD 写入带宽,CPU 占用从 SPDK 降至 0.6 核

八、io_uring IOPOLL 模式在生产部署的取舍

ZNS 设备支持 polling mode(NVMe CQ 轮询),配合 io_uring SQPOLL + IOPOLL 可以进一步减少系统调用开销:

let mut ring = IoUring::builder()
    .setup_sqpoll(Some(Duration::from_millis(10)))  // 内核轮询 SQ
    .setup_iopoll(true)                               // 设备级 CQ 轮询
    .build(queue_size)?;

实测 IOPOLL 相比 IRQ + IRQ_POLL:

  • 吞吐 +7%(1.28M → 1.37M IOPS)
  • CPU +22%(0.6 核 → 0.73 核)
  • P99 尾延迟 -12%

取舍建议:绑定核心时值得一试,共享核心时 IRQ 模式更安全。

九、总结

io_uring uring_cmd 为 ZNS 设备提供了一个中间层方案——不像 SPDK 那样需要完全用户态驱动,也不像 f2fs 块文件系统那样叠加 GC 开销。

核心设计决策:

  1. uring_cmd 代替 ioctl(NVMe_COMMAND):享受 io_uring 的 batching、submit_and_wait 语义
  2. Registered Buffers + hugepage:消除数据拷贝
  3. 单线程 Zone 管理 + atomic 旋转索引:顺序写无需全局锁
  4. per-stream 流隔离:多训练任务并发写不互相污染

代码已开源:github.com/example/zns-log-engine(Rust + io_uring,Apache 2.0 License)

ZNS 正在从"高性能特种设备"走向 AI 训练基础设施的标准存储层。io_uring uring_cmd 让这块拼图更容易被标准化集成——VFS 语义不改,块设备不独占,性能接近 SPDK。对需要在生产 K8s 集群中部署 ZNS 日志存储的团队,这可能是最务实的路径。

点赞(0) 打赏

评论列表 共有 0 条评论

暂无评论
立即
投稿

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部