virtio_blk.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cpu_io.h>
 
#include <array>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
 
#include "expected.hpp"
#include "io_buffer.hpp"
#include "kernel_log.hpp"
#include "virtio/defs.h"
#include "virtio/device/blk/virtio_blk_defs.h"
#include "virtio/device/device_initializer.hpp"
#include "virtio/transport/mmio.hpp"
#include "virtio/virt_queue/split.hpp"
 
namespace virtio::blk {
 
template <typename TransportT = MmioTransport,
          typename VirtqueueT = SplitVirtqueue>
class VirtioBlk {
 public:
  using UserData = void*;
 
  static constexpr uint16_t kMaxInflight = 64;
 
  static constexpr size_t kMaxSgElements = 18;
 
  [[nodiscard]] static constexpr auto GetRequiredVqMemSize(uint16_t queue_count,
                                                           uint32_t queue_size)
      -> std::pair<size_t, size_t> {
    // 始终按 event_idx=true 分配，因为特性协商在分配之后
    size_t per_queue =
        VirtqueueT::CalcSize(static_cast<uint16_t>(queue_size), true);
    return {per_queue * queue_count, 4096};
  }
 
  [[nodiscard]] static constexpr auto CalcDmaSize(uint16_t queue_size = 128)
      -> size_t {
    // 始终按 event_idx=true 分配，确保空间充足
    return VirtqueueT::CalcSize(queue_size, true);
  }
 
  [[nodiscard]] static constexpr auto GetRequiredSlotMemSize()
      -> std::pair<size_t, size_t> {
    return {sizeof(RequestSlot) * kMaxInflight, alignof(RequestSlot)};
  }
 
  [[nodiscard]] static auto Create(
      uint64_t mmio_base, const DmaRegion& vq_dma, const DmaRegion& slot_dma,
      VirtToPhysFunc virt_to_phys = IdentityVirtToPhys,
      uint16_t queue_count = 1, uint32_t queue_size = 128,
      uint64_t driver_features = 0) -> Expected<VirtioBlk> {
    if (queue_count == 0) {
      return std::unexpected(Error{ErrorCode::kInvalidArgument});
    }
    if (queue_count > 1) {
    }
 
    // 1. 创建传输层
    TransportT transport(mmio_base);
    if (!transport.IsValid()) {
      return std::unexpected(Error{ErrorCode::kTransportNotInitialized});
    }
 
    // 2. 设备初始化序列
    DeviceInitializer<TransportT> initializer(transport);
 
    uint64_t wanted_features =
        static_cast<uint64_t>(ReservedFeature::kVersion1) |
        static_cast<uint64_t>(ReservedFeature::kEventIdx) | driver_features;
    auto negotiated_result = initializer.Init(wanted_features);
    if (!negotiated_result) {
      return std::unexpected(negotiated_result.error());
    }
    uint64_t negotiated = *negotiated_result;
 
    if ((negotiated & static_cast<uint64_t>(ReservedFeature::kVersion1)) == 0) {
      return std::unexpected(Error{ErrorCode::kFeatureNegotiationFailed});
    }
 
    // 根据协商结果决定是否启用 Event Index
    bool event_idx =
        (negotiated & static_cast<uint64_t>(ReservedFeature::kEventIdx)) != 0;
    if (event_idx) {
    }
    // 3. 创建 Virtqueue
    VirtqueueT vq(vq_dma, static_cast<uint16_t>(queue_size), event_idx);
    if (!vq.IsValid()) {
      return std::unexpected(Error{ErrorCode::kInvalidArgument});
    }
 
    // 4. 配置队列
    const uint32_t queue_idx = 0;
    auto setup_result = initializer.SetupQueue(
        queue_idx, vq.DescPhys(), vq.AvailPhys(), vq.UsedPhys(), vq.Size());
    if (!setup_result) {
      return std::unexpected(setup_result.error());
    }
 
    // 5. 激活设备
    auto activate_result = initializer.Activate();
    if (!activate_result) {
      return std::unexpected(activate_result.error());
    }
 
    return VirtioBlk(std::move(transport), std::move(vq), negotiated, slot_dma,
                     virt_to_phys);
  }
 
  // ======== 异步 IO 接口 (Enqueue/Kick/HandleInterrupt) ========
 
  [[nodiscard]] auto EnqueueRead(uint16_t queue_index, uint64_t sector,
                                 const IoVec* buffers, size_t buffer_count,
                                 UserData token = nullptr) -> Expected<void> {
    return DoEnqueue(ReqType::kIn, queue_index, sector, buffers, buffer_count,
                     token);
  }
 
  [[nodiscard]] auto EnqueueWrite(uint16_t queue_index, uint64_t sector,
                                  const IoVec* buffers, size_t buffer_count,
                                  UserData token = nullptr) -> Expected<void> {
    return DoEnqueue(ReqType::kOut, queue_index, sector, buffers, buffer_count,
                     token);
  }
 
  auto Kick(uint16_t queue_index) -> void {
    if (queue_index != 0) {
      return;
    }
    // 写屏障：确保 Available Ring 更新对设备可见
    cpu_io::Wmb();
 
    if (vq_.EventIdxEnabled()) {
      auto* avail_event_ptr = vq_.UsedAvailEvent();
      if (avail_event_ptr != nullptr) {
        uint16_t avail_event = *avail_event_ptr;
        uint16_t new_idx = vq_.AvailIdx();
        if (VringNeedEvent(avail_event, new_idx, old_avail_idx_)) {
          transport_.NotifyQueue(queue_index);
        } else {
          stats_.kicks_elided++;
        }
        old_avail_idx_ = new_idx;
      } else {
        transport_.NotifyQueue(queue_index);
      }
    } else {
      transport_.NotifyQueue(queue_index);
    }
  }
 
  template <typename CompletionCallback>
  auto HandleInterrupt(CompletionCallback&& on_complete) -> void {
    // 确认中断
    uint32_t isr_status = transport_.GetInterruptStatus();
    if (isr_status != 0) {
      transport_.AckInterrupt(isr_status);
    }
    stats_.interrupts_handled++;
 
    ProcessCompletions(static_cast<CompletionCallback&&>(on_complete));
    UpdateUsedEvent();
  }
 
  auto HandleInterrupt() -> void {
    uint32_t status = transport_.GetInterruptStatus();
    if (status != 0) {
      transport_.AckInterrupt(status);
    }
    stats_.interrupts_handled++;
    request_completed_ = true;
    cpu_io::Wmb();
  }
 
  // ======== 同步便捷方法 ========
 
  [[nodiscard]] auto Read(uint64_t sector, uint8_t* data) -> Expected<void> {
    if (data == nullptr) {
      return std::unexpected(Error{ErrorCode::kInvalidArgument});
    }
    IoVec data_iov{virt_to_phys_(reinterpret_cast<uintptr_t>(data)),
                   kSectorSize};
    return SubmitSyncRequest(ReqType::kIn, sector, &data_iov, 1);
  }
 
  [[nodiscard]] auto Write(uint64_t sector, const uint8_t* data)
      -> Expected<void> {
    if (data == nullptr) {
      return std::unexpected(Error{ErrorCode::kInvalidArgument});
    }
    IoVec data_iov{
        virt_to_phys_(reinterpret_cast<uintptr_t>(const_cast<uint8_t*>(data))),
        kSectorSize};
    return SubmitSyncRequest(ReqType::kOut, sector, &data_iov, 1);
  }
 
  // ======== 配置与监控 ========
 
  [[nodiscard]] auto ReadConfig() const -> BlkConfig {
    BlkConfig config{};
 
    config.capacity = transport_.ReadConfigU64(
        static_cast<uint32_t>(BlkConfigOffset::kCapacity));
    config.size_max = transport_.ReadConfigU32(
        static_cast<uint32_t>(BlkConfigOffset::kSizeMax));
    config.seg_max = transport_.ReadConfigU32(
        static_cast<uint32_t>(BlkConfigOffset::kSegMax));
 
    if ((negotiated_features_ &
         static_cast<uint64_t>(BlkFeatureBit::kGeometry)) != 0) {
      config.geometry.cylinders = transport_.ReadConfigU16(
          static_cast<uint32_t>(BlkConfigOffset::kGeometryCylinders));
      config.geometry.heads = transport_.ReadConfigU8(
          static_cast<uint32_t>(BlkConfigOffset::kGeometryHeads));
      config.geometry.sectors = transport_.ReadConfigU8(
          static_cast<uint32_t>(BlkConfigOffset::kGeometrySectors));
    }
 
    if ((negotiated_features_ &
         static_cast<uint64_t>(BlkFeatureBit::kBlkSize)) != 0) {
      config.blk_size = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kBlkSize));
    }
 
    if ((negotiated_features_ &
         static_cast<uint64_t>(BlkFeatureBit::kTopology)) != 0) {
      config.topology.physical_block_exp = transport_.ReadConfigU8(
          static_cast<uint32_t>(BlkConfigOffset::kTopologyPhysBlockExp));
      config.topology.alignment_offset = transport_.ReadConfigU8(
          static_cast<uint32_t>(BlkConfigOffset::kTopologyAlignOffset));
      config.topology.min_io_size = transport_.ReadConfigU16(
          static_cast<uint32_t>(BlkConfigOffset::kTopologyMinIoSize));
      config.topology.opt_io_size = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kTopologyOptIoSize));
    }
 
    if ((negotiated_features_ &
         static_cast<uint64_t>(BlkFeatureBit::kConfigWce)) != 0) {
      config.writeback = transport_.ReadConfigU8(
          static_cast<uint32_t>(BlkConfigOffset::kWriteback));
    }
 
    if ((negotiated_features_ &
         static_cast<uint64_t>(BlkFeatureBit::kDiscard)) != 0) {
      config.max_discard_sectors = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kMaxDiscardSectors));
      config.max_discard_seg = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kMaxDiscardSeg));
      config.discard_sector_alignment = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kDiscardSectorAlignment));
    }
 
    if ((negotiated_features_ &
         static_cast<uint64_t>(BlkFeatureBit::kWriteZeroes)) != 0) {
      config.max_write_zeroes_sectors = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kMaxWriteZeroesSectors));
      config.max_write_zeroes_seg = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kMaxWriteZeroesSeg));
      config.write_zeroes_may_unmap = transport_.ReadConfigU8(
          static_cast<uint32_t>(BlkConfigOffset::kWriteZeroesMayUnmap));
    }
 
    if ((negotiated_features_ &
         static_cast<uint64_t>(BlkFeatureBit::kSecureErase)) != 0) {
      config.max_secure_erase_sectors = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kMaxSecureEraseSectors));
      config.max_secure_erase_seg = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kMaxSecureEraseSeg));
      config.secure_erase_sector_alignment = transport_.ReadConfigU32(
          static_cast<uint32_t>(BlkConfigOffset::kSecureEraseSectorAlignment));
    }
 
    if ((negotiated_features_ & static_cast<uint64_t>(BlkFeatureBit::kMq)) !=
        0) {
      config.num_queues = transport_.ReadConfigU16(
          static_cast<uint32_t>(BlkConfigOffset::kNumQueues));
    }
 
    return config;
  }
 
  [[nodiscard]] auto GetCapacity() const -> uint64_t {
    return transport_.ReadConfigU64(
        static_cast<uint32_t>(BlkConfigOffset::kCapacity));
  }
 
  [[nodiscard]] auto GetNegotiatedFeatures() const -> uint64_t {
    return negotiated_features_;
  }
 
  [[nodiscard]] auto GetStats() const -> VirtioStats { return stats_; }
 
  VirtioBlk() = delete;
  VirtioBlk(VirtioBlk&& other) noexcept
      : transport_(std::move(other.transport_)),
        vq_(std::move(other.vq_)),
        negotiated_features_(other.negotiated_features_),
        slot_dma_(other.slot_dma_),
        slots_(other.slots_),
        virt_to_phys_(other.virt_to_phys_),
        stats_(other.stats_),
        slot_bitmap_(other.slot_bitmap_),
        old_avail_idx_(other.old_avail_idx_),
        request_completed_(other.request_completed_) {
    other.slots_ = nullptr;
    other.slot_bitmap_ = 0;
  }
  auto operator=(VirtioBlk&& other) noexcept -> VirtioBlk& {
    if (this != &other) {
      transport_ = std::move(other.transport_);
      vq_ = std::move(other.vq_);
      negotiated_features_ = other.negotiated_features_;
      slot_dma_ = other.slot_dma_;
      slots_ = other.slots_;
      virt_to_phys_ = other.virt_to_phys_;
      stats_ = other.stats_;
      slot_bitmap_ = other.slot_bitmap_;
      old_avail_idx_ = other.old_avail_idx_;
      request_completed_ = other.request_completed_;
      other.slots_ = nullptr;
      other.slot_bitmap_ = 0;
    }
    return *this;
  }
  VirtioBlk(const VirtioBlk&) = delete;
  auto operator=(const VirtioBlk&) -> VirtioBlk& = delete;
  ~VirtioBlk() = default;
 
 private:
  struct RequestSlot {
    alignas(16) BlkReqHeader header;
    alignas(4) volatile uint8_t status;
    UserData token;
    uint16_t desc_head;
  };
 
  VirtioBlk(TransportT transport, VirtqueueT vq, uint64_t features,
            const DmaRegion& slot_dma, VirtToPhysFunc v2p)
      : transport_(std::move(transport)),
        vq_(std::move(vq)),
        negotiated_features_(features),
        slot_dma_(slot_dma),
        slots_(reinterpret_cast<RequestSlot*>(slot_dma.Data())),
        virt_to_phys_(v2p),
        stats_{},
        slot_bitmap_(0),
        old_avail_idx_(0),
        request_completed_(false) {}
 
  [[nodiscard]] auto DoEnqueue(ReqType type, uint16_t queue_index,
                               uint64_t sector, const IoVec* buffers,
                               size_t buffer_count, UserData token)
      -> Expected<void> {
    if (queue_index != 0) {
      return std::unexpected(Error{ErrorCode::kInvalidArgument});
    }
 
    if (buffer_count + 2 > kMaxSgElements) {
      return std::unexpected(Error{ErrorCode::kInvalidArgument});
    }
 
    auto slot_result = AllocRequestSlot();
    if (!slot_result) {
      stats_.queue_full_errors++;
      return std::unexpected(slot_result.error());
    }
    uint16_t slot_idx = *slot_result;
    auto& slot = slots_[slot_idx];
 
    slot.header.type = static_cast<uint32_t>(type);
    slot.header.reserved = 0;
    slot.header.sector = sector;
    slot.status = 0xFF;  // sentinel：设备完成后会覆写
    slot.token = token;
 
    std::array<IoVec, kMaxSgElements> readable_iovs{};
    std::array<IoVec, kMaxSgElements> writable_iovs{};
    size_t readable_count = 0;
    size_t writable_count = 0;
 
    auto slot_base_phys =
        slot_dma_.phys + static_cast<size_t>(slot_idx) * sizeof(RequestSlot);
    readable_iovs[readable_count++] = {
        slot_base_phys + offsetof(RequestSlot, header), sizeof(BlkReqHeader)};
 
    if (type == ReqType::kIn) {
      for (size_t i = 0; i < buffer_count; ++i) {
        writable_iovs[writable_count++] = buffers[i];
      }
    } else {
      for (size_t i = 0; i < buffer_count; ++i) {
        readable_iovs[readable_count++] = buffers[i];
      }
    }
 
    // 状态字节始终为 device-writable
    writable_iovs[writable_count++] = {
        slot_base_phys + offsetof(RequestSlot, status), sizeof(uint8_t)};
 
    cpu_io::Wmb();
 
    auto chain_result = vq_.SubmitChain(readable_iovs.data(), readable_count,
                                        writable_iovs.data(), writable_count);
    if (!chain_result) {
      FreeRequestSlot(slot_idx);
      stats_.queue_full_errors++;
      return std::unexpected(chain_result.error());
    }
 
    slot.desc_head = *chain_result;
 
    return {};
  }
 
  template <typename CompletionCallback>
  auto ProcessCompletions(CompletionCallback&& on_complete) -> void {
    cpu_io::Rmb();
 
    while (vq_.HasUsed()) {
      auto elem_result = vq_.PopUsed();
      if (!elem_result) {
        break;
      }
 
      auto elem = *elem_result;
      auto head = static_cast<uint16_t>(elem.id);
 
      uint16_t slot_idx = FindSlotByDescHead(head);
      if (slot_idx < kMaxInflight) {
        auto& slot = slots_[slot_idx];
 
        cpu_io::Rmb();
 
        ErrorCode ec = MapBlkStatus(slot.status);
        on_complete(slot.token, ec);
        stats_.bytes_transferred += elem.len;
        FreeRequestSlot(slot_idx);
      }
 
      (void)vq_.FreeChain(head);
    }
  }
 
  [[nodiscard]] auto AllocRequestSlot() -> Expected<uint16_t> {
    uint64_t free_bits = ~slot_bitmap_;
    if (free_bits == 0) {
      return std::unexpected(Error{ErrorCode::kNoFreeDescriptors});
    }
    auto idx = static_cast<uint16_t>(__builtin_ctzll(free_bits));
    if (idx >= kMaxInflight) {
      return std::unexpected(Error{ErrorCode::kNoFreeDescriptors});
    }
    slot_bitmap_ |= (uint64_t{1} << idx);
    return idx;
  }
 
  auto FreeRequestSlot(uint16_t idx) -> void {
    if (idx < kMaxInflight) {
      slot_bitmap_ &= ~(uint64_t{1} << idx);
    }
  }
 
  [[nodiscard]] auto FindSlotByDescHead(uint16_t desc_head) const -> uint16_t {
    uint64_t used = slot_bitmap_;
    while (used != 0) {
      auto i = static_cast<uint16_t>(__builtin_ctzll(used));
      if (slots_[i].desc_head == desc_head) {
        return i;
      }
      used &= used - 1;  // clear lowest set bit
    }
    return kMaxInflight;
  }
 
  [[nodiscard]] static auto MapBlkStatus(uint8_t status) -> ErrorCode {
    switch (status) {
      case static_cast<uint8_t>(BlkStatus::kOk):
        return ErrorCode::kSuccess;
      case static_cast<uint8_t>(BlkStatus::kIoErr):
        return ErrorCode::kIoError;
      case static_cast<uint8_t>(BlkStatus::kUnsupp):
        return ErrorCode::kNotSupported;
      default:
        return ErrorCode::kDeviceError;
    }
  }
 
  [[nodiscard]] static auto VringNeedEvent(uint16_t event_idx, uint16_t new_idx,
                                           uint16_t old_idx) -> bool {
    return static_cast<uint16_t>(new_idx - event_idx - 1) <
           static_cast<uint16_t>(new_idx - old_idx);
  }
 
  auto UpdateUsedEvent() -> void {
    if (vq_.EventIdxEnabled()) {
      auto* used_event_ptr = vq_.AvailUsedEvent();
      if (used_event_ptr != nullptr) {
        *used_event_ptr = vq_.LastUsedIdx();
        cpu_io::Wmb();
      }
    }
  }
 
  [[nodiscard]] auto SubmitSyncRequest(ReqType type, uint64_t sector,
                                       const IoVec* buffers,
                                       size_t buffer_count) -> Expected<void> {
    auto enq = DoEnqueue(type, 0, sector, buffers, buffer_count, nullptr);
    if (!enq) {
      return std::unexpected(enq.error());
    }
 
    Kick(0);
 
    uint32_t spin_limit = 100000000U;
 
    for (uint32_t i = 0; i < spin_limit; ++i) {
      cpu_io::Rmb();
      if (vq_.HasUsed()) {
        break;
      }
    }
 
    if (!vq_.HasUsed()) {
      klog::Warn("Sync request timeout: sector={}", sector);
      return std::unexpected(Error{ErrorCode::kTimeout});
    }
 
    ErrorCode result = ErrorCode::kSuccess;
    bool done = false;
    ProcessCompletions([&done, &result](UserData, ErrorCode status) {
      done = true;
      result = status;
    });
    UpdateUsedEvent();
 
    if (!done) {
      return std::unexpected(Error{ErrorCode::kTimeout});
    }
    if (result != ErrorCode::kSuccess) {
      return std::unexpected(Error{result});
    }
    return {};
  }
  TransportT transport_;
  VirtqueueT vq_;
  uint64_t negotiated_features_{0};
  DmaRegion slot_dma_{};
  RequestSlot* slots_{nullptr};
  VirtToPhysFunc virt_to_phys_{IdentityVirtToPhys};
  VirtioStats stats_;
  uint64_t slot_bitmap_{};
  uint16_t old_avail_idx_{0};
  volatile bool request_completed_{false};
};
 
}  // namespace virtio::blk