spinlock.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cpu_io.h>
 
#include <atomic>
#include <concepts>
#include <cstddef>
#include <limits>
 
#include "expected.hpp"
#include "kernel_log.hpp"
#include "kstd_cstdio"
 
class SpinLock {
 public:
  const char* name{"unnamed"};
 
  [[nodiscard]] __always_inline auto Lock() -> Expected<void> {
    auto intr_enable = cpu_io::GetInterruptStatus();
    cpu_io::DisableInterrupt();
 
    // 先尝试获取锁
    while (locked_.test_and_set(std::memory_order_acquire)) {
      // 在等待时检查是否是当前核心持有锁（递归锁检测）
      if (core_id_.load(std::memory_order_acquire) ==
          cpu_io::GetCurrentCoreId()) {
        // 递归锁定，恢复中断状态并返回失败
        if (intr_enable) {
          cpu_io::EnableInterrupt();
        }
        // klog::Err("spinlock {}: {} recursive lock detected.",
        //           cpu_io::GetCurrentCoreId(), name);
        return std::unexpected(Error{ErrorCode::kSpinLockRecursiveLock});
      }
      cpu_io::Pause();
    }
 
    // 获取锁成功后立即设置 core_id_
    core_id_.store(cpu_io::GetCurrentCoreId(), std::memory_order_release);
    saved_intr_enable_ = intr_enable;
    return {};
  }
 
  [[nodiscard]] __always_inline auto UnLock() -> Expected<void> {
    if (!IsLockedByCurrentCore()) {
      // klog::Err("spinlock {}: {} unlock by non-owner detected.",
      //           cpu_io::GetCurrentCoreId(), name);
      return std::unexpected(Error{ErrorCode::kSpinLockNotOwned});
    }
 
    // 先重置 core_id_，再释放锁
    core_id_.store(std::numeric_limits<size_t>::max(),
                   std::memory_order_release);
    locked_.clear(std::memory_order_release);
 
    if (saved_intr_enable_) {
      cpu_io::EnableInterrupt();
    }
    return {};
  }
 
 
  explicit SpinLock(const char* _name) : name(_name) {}
 
  SpinLock() = default;
  SpinLock(const SpinLock&) = delete;
  SpinLock(SpinLock&&) = default;
  auto operator=(const SpinLock&) -> SpinLock& = delete;
  auto operator=(SpinLock&&) -> SpinLock& = default;
  ~SpinLock() = default;
 
 protected:
  std::atomic_flag locked_{ATOMIC_FLAG_INIT};
  std::atomic<size_t> core_id_{std::numeric_limits<size_t>::max()};
  bool saved_intr_enable_{false};
 
  __always_inline auto IsLockedByCurrentCore() -> bool {
    return locked_.test(std::memory_order_acquire) &&
           (core_id_.load(std::memory_order_acquire) ==
            cpu_io::GetCurrentCoreId());
  }
};
 
template <typename Mutex>
  requires requires(Mutex& m) {
    { m.Lock() } -> std::same_as<Expected<void>>;
    { m.UnLock() } -> std::same_as<Expected<void>>;
  }
class LockGuard {
 public:
  using mutex_type = Mutex;
 
 
  explicit LockGuard(mutex_type& mutex) : mutex_(mutex) {
    mutex_.Lock().or_else([&](auto&& err) {
      char core_buf[4] = {};
      auto core_id = cpu_io::GetCurrentCoreId();
      size_t pos = 0;
      if (core_id == 0) {
        core_buf[pos++] = '0';
      } else {
        char tmp[4] = {};
        size_t tmp_pos = 0;
        while (core_id > 0 && tmp_pos < sizeof(tmp)) {
          tmp[tmp_pos++] = static_cast<char>('0' + (core_id % 10));
          core_id /= 10;
        }
        while (tmp_pos > 0) {
          core_buf[pos++] = tmp[--tmp_pos];
        }
      }
      core_buf[pos] = '\0';
      klog::RawPut("PANIC: LockGuard failed to acquire lock '");
      klog::RawPut(mutex_.name);
      klog::RawPut("' on core ");
      klog::RawPut(core_buf);
      klog::RawPut(": ");
      klog::RawPut(err.message());
      klog::RawPut("\n");
      while (true) {
        cpu_io::Pause();
      }
      return Expected<void>{};
    });
  }
 
  ~LockGuard() {
    mutex_.UnLock().or_else([&](auto&& err) {
      char core_buf[4] = {};
      auto core_id = cpu_io::GetCurrentCoreId();
      size_t pos = 0;
      if (core_id == 0) {
        core_buf[pos++] = '0';
      } else {
        char tmp[4] = {};
        size_t tmp_pos = 0;
        while (core_id > 0 && tmp_pos < sizeof(tmp)) {
          tmp[tmp_pos++] = static_cast<char>('0' + (core_id % 10));
          core_id /= 10;
        }
        while (tmp_pos > 0) {
          core_buf[pos++] = tmp[--tmp_pos];
        }
      }
      core_buf[pos] = '\0';
      klog::RawPut("PANIC: LockGuard failed to release lock '");
      klog::RawPut(mutex_.name);
      klog::RawPut("' on core ");
      klog::RawPut(core_buf);
      klog::RawPut(": ");
      klog::RawPut(err.message());
      klog::RawPut("\n");
      while (true) {
        cpu_io::Pause();
      }
      return Expected<void>{};
    });
  }
 
  LockGuard() = delete;
  LockGuard(const LockGuard&) = delete;
  LockGuard(LockGuard&&) = delete;
  auto operator=(const LockGuard&) -> LockGuard& = delete;
  auto operator=(LockGuard&&) -> LockGuard& = delete;
 
 private:
  mutex_type& mutex_;
};