gongxiangp2/lib/disk-space.ts

/**
 * R38修复：磁盘空间检查工具
 *
 * 用途：在大文件操作前检查磁盘剩余空间，防止因磁盘写满导致服务崩溃
 *
 * 支持平台：
 * - Linux: 使用 `df` 命令
 * - Windows (WSL): 使用 `df` 命令
 * - Windows (原生): 使用 PowerShell 的 Get-PSDrive
 */

import { exec } from 'child_process'
import { promisify } from 'util'

const execAsync = promisify(exec)

interface DiskSpaceInfo {
  /** 磁盘总容量（字节） */
  total: number
  /** 已使用空间（字节） */
  used: number
  /** 可用空间（字节） */
  available: number
  /** 使用百分比 (0-100) */
  usePercent: number
  /** 挂载点/盘符 */
  mountPoint: string
}

/**
 * 检查磁盘空间（Linux/WSL 使用 df 命令）
 */
async function checkDiskSpaceLinux(path: string): Promise<DiskSpaceInfo | null> {
  try {
    // df -B1: 以字节为单位显示
    // --output: 指定输出字段
    const { stdout } = await execAsync(`df -B1 --output=size,used,avail,pcent,target "${path}"`, {
      timeout: 5000, // 5秒超时
    })

    const lines = stdout.trim().split('\n')
    if (lines.length < 2) {
      console.warn('[DiskSpace] Failed to parse df output')
      return null
    }

    // 跳过标题行，解析数据行
    const data = lines[1].trim().split(/\s+/)
    if (data.length < 5) {
      console.warn('[DiskSpace] Invalid df output format')
      return null
    }

    const total = parseInt(data[0], 10)
    const used = parseInt(data[1], 10)
    const available = parseInt(data[2], 10)
    const usePercentStr = data[3].replace('%', '')
    const usePercent = parseInt(usePercentStr, 10)
    const mountPoint = data[4]

    return {
      total,
      used,
      available,
      usePercent,
      mountPoint,
    }
  } catch (error) {
    console.error('[DiskSpace] Failed to check disk space:', error)
    return null
  }
}

/**
 * 检查磁盘空间（Windows 原生使用 PowerShell）
 */
async function checkDiskSpaceWindows(path: string): Promise<DiskSpaceInfo | null> {
  try {
    // 获取盘符
    const driveLetter = path.charAt(0).toUpperCase()
    const drive = driveLetter + ':'

    // 使用 PowerShell Get-PSDrive 获取磁盘信息
    const { stdout } = await execAsync(
      `powershell -Command "Get-PSDrive -Name '${driveLetter}' | Select-Object Used,Free | ConvertTo-Json"`,
      { timeout: 5000 }
    )

    const result = JSON.parse(stdout)
    if (!result) {
      return null
    }

    const used = result.Used || 0
    const free = result.Free || 0
    const total = used + free
    const usePercent = total > 0 ? Math.round((used / total) * 100) : 0

    return {
      total,
      used,
      available: free,
      usePercent,
      mountPoint: drive,
    }
  } catch (error) {
    console.error('[DiskSpace] Failed to check disk space (Windows):', error)
    return null
  }
}

/**
 * 检查磁盘空间（自动检测平台）
 * @param path 要检查的路径（默认为临时目录）
 * @returns 磁盘空间信息，失败返回 null
 */
export async function checkDiskSpace(path: string = require('os').tmpdir()): Promise<DiskSpaceInfo | null> {
  const platform = process.platform

  if (platform === 'win32' && !process.env.WSL_DISTRO_NAME) {
    // Windows 原生环境（非 WSL）
    return checkDiskSpaceWindows(path)
  } else {
    // Linux、macOS、WSL
    return checkDiskSpaceLinux(path)
  }
}

/**
 * 检查是否有足够磁盘空间
 * @param requiredBytes 需要的字节数
 * @param path 检查的路径
 * @param minFreePercent 最小剩余空间百分比（默认 10%）
 * @returns 是否有足够空间
 */
export async function hasEnoughDiskSpace(
  requiredBytes: number,
  path?: string,
  minFreePercent: number = 10
): Promise<boolean> {
  const space = await checkDiskSpace(path)
  if (!space) {
    // 无法检查磁盘空间时，保守处理：假设有足够空间
    console.warn('[DiskSpace] Unable to check disk space, assuming sufficient space')
    return true
  }

  // 检查绝对值
  if (space.available < requiredBytes) {
    console.warn(
      `[DiskSpace] Insufficient disk space: need ${formatBytes(requiredBytes)}, available ${formatBytes(space.available)}`
    )
    return false
  }

  // 检查百分比（确保不会把磁盘填满）
  if (space.usePercent > (100 - minFreePercent)) {
    console.warn(
      `[DiskSpace] Disk usage too high: ${space.usePercent}% (max allowed: ${100 - minFreePercent}%)`
    )
    return false
  }

  return true
}

/**
 * R38修复：检查磁盘空间并在不足时抛出错误
 * @param requiredBytes 需要的字节数
 * @param operation 操作名称（用于错误消息）
 * @param path 检查的路径
 * @throws Error 当磁盘空间不足时
 */
export async function ensureEnoughDiskSpace(
  requiredBytes: number,
  operation: string = '操作',
  path?: string
): Promise<void> {
  const space = await checkDiskSpace(path)
  if (!space) {
    console.warn(`[DiskSpace] Unable to check disk space before ${operation}, proceeding anyway`)
    return
  }

  // 计算最小保留空间（磁盘总容量的 10%，或至少 5GB）
  const minReservedBytes = Math.max(space.total * 0.1, 5 * 1024 * 1024 * 1024)
  const requiredWithReserve = requiredBytes + minReservedBytes

  if (space.available < requiredWithReserve) {
    throw new Error(
      `磁盘空间不足，无法${operation}。` +
        `需要 ${formatBytes(requiredBytes)}，` +
        `但只剩余 ${formatBytes(space.available)} ` +
        `(使用率 ${space.usePercent}%)。` +
        `请清理磁盘空间后再试。`
    )
  }

  console.log(
    `[DiskSpace] Disk space check passed for ${operation}: ` +
    `need ${formatBytes(requiredBytes)}, available ${formatBytes(space.available)} (${space.usePercent}% used)`
  )
}

/**
 * 格式化字节数为可读格式
 */
export function formatBytes(bytes: number): string {
  const units = ['B', 'KB', 'MB', 'GB', 'TB']
  let size = bytes
  let unitIndex = 0

  while (size >= 1024 && unitIndex < units.length - 1) {
    size /= 1024
    unitIndex++
  }

  return `${size.toFixed(2)} ${units[unitIndex]}`
}

/**
 * 获取磁盘使用率警告级别
 * @returns 'ok' (正常), 'warning' (>70%), 'critical' (>90%)
 */
export async function getDiskUsageLevel(path?: string): Promise<'ok' | 'warning' | 'critical'> {
  const space = await checkDiskSpace(path)
  if (!space) {
    return 'ok'
  }

  if (space.usePercent >= 90) {
    return 'critical'
  } else if (space.usePercent >= 70) {
    return 'warning'
  }
  return 'ok'
}