waf-platform/EdgeNode/internal/utils/fs/ssd.go

// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .

package fsutils

import (
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
)

// StorageType 存储类型
type StorageType int

const (
	StorageTypeUnknown StorageType = 0
	StorageTypeSSD     StorageType = 1
	StorageTypeHDD     StorageType = 2
)

var (
	ssdCache     = map[string]StorageType{} // path => StorageType
	ssdCacheLock sync.RWMutex
)

// IsSSD 检测指定路径是否为 SSD
// 方法：
// 1. 检查 /sys/block/*/queue/rotational（Linux）
// 2. 使用 lsblk 命令（Linux）
// 3. 检查设备名称（nvme, ssd 等关键词）
func IsSSD(path string) (bool, error) {
	storageType, err := DetectStorageType(path)
	if err != nil {
		return false, err
	}
	return storageType == StorageTypeSSD, nil
}

// DetectStorageType 检测存储类型
func DetectStorageType(path string) (StorageType, error) {
	// 检查缓存
	ssdCacheLock.RLock()
	if cachedType, ok := ssdCache[path]; ok {
		ssdCacheLock.RUnlock()
		return cachedType, nil
	}
	ssdCacheLock.RUnlock()

	// 获取绝对路径
	absPath, err := filepath.Abs(path)
	if err != nil {
		return StorageTypeUnknown, err
	}

	// 获取设备路径
	devicePath, err := getDevicePath(absPath)
	if err != nil {
		return StorageTypeUnknown, err
	}

	// 方法 1：检查 /sys/block/*/queue/rotational（Linux）
	if isSSD, err := checkRotational(devicePath); err == nil {
		storageType := StorageTypeHDD
		if isSSD {
			storageType = StorageTypeSSD
		}
		ssdCacheLock.Lock()
		ssdCache[path] = storageType
		ssdCacheLock.Unlock()
		return storageType, nil
	}

	// 方法 2：使用 lsblk 命令（Linux）
	if isSSD, err := checkWithLsblk(devicePath); err == nil {
		storageType := StorageTypeHDD
		if isSSD {
			storageType = StorageTypeSSD
		}
		ssdCacheLock.Lock()
		ssdCache[path] = storageType
		ssdCacheLock.Unlock()
		return storageType, nil
	}

	// 方法 3：检查设备名称（nvme, ssd 等关键词）
	if isSSD := checkDeviceName(devicePath); isSSD {
		ssdCacheLock.Lock()
		ssdCache[path] = StorageTypeSSD
		ssdCacheLock.Unlock()
		return StorageTypeSSD, nil
	}

	// 默认返回未知
	ssdCacheLock.Lock()
	ssdCache[path] = StorageTypeUnknown
	ssdCacheLock.Unlock()
	return StorageTypeUnknown, nil
}

// getDevicePath 获取设备路径
func getDevicePath(path string) (string, error) {
	// 从路径提取设备名（简化实现）
	// 实际应该通过 stat 获取设备信息
	// 这里使用路径的第一级目录作为设备标识
	parts := strings.Split(strings.Trim(path, "/"), "/")
	if len(parts) > 0 {
		return parts[0], nil
	}

	return "", os.ErrNotExist
}

// checkRotational 检查 /sys/block/*/queue/rotational
func checkRotational(devicePath string) (bool, error) {
	// 提取设备名（如 sda, nvme0n1）
	deviceName := extractDeviceName(devicePath)
	if len(deviceName) == 0 {
		return false, os.ErrNotExist
	}

	// 检查 /sys/block/{device}/queue/rotational
	rotationalPath := "/sys/block/" + deviceName + "/queue/rotational"
	data, err := os.ReadFile(rotationalPath)
	if err != nil {
		return false, err
	}

	// 0 表示 SSD，1 表示 HDD
	value := strings.TrimSpace(string(data))
	return value == "0", nil
}

// checkWithLsblk 使用 lsblk 命令检查
func checkWithLsblk(devicePath string) (bool, error) {
	cmd := exec.Command("lsblk", "-d", "-o", "name,rota", "-n")
	output, err := cmd.Output()
	if err != nil {
		return false, err
	}

	lines := strings.Split(string(output), "\n")
	deviceName := extractDeviceName(devicePath)
	for _, line := range lines {
		fields := strings.Fields(line)
		if len(fields) >= 2 && fields[0] == deviceName {
			// rota=0 表示 SSD，rota=1 表示 HDD
			return fields[1] == "0", nil
		}
	}

	return false, os.ErrNotExist
}

// checkDeviceName 检查设备名称
func checkDeviceName(devicePath string) bool {
	deviceName := strings.ToLower(extractDeviceName(devicePath))

	// 检查常见 SSD 关键词
	ssdKeywords := []string{"nvme", "ssd", "flash"}
	for _, keyword := range ssdKeywords {
		if strings.Contains(deviceName, keyword) {
			return true
		}
	}

	return false
}

// extractDeviceName 从路径提取设备名
func extractDeviceName(path string) string {
	// 简化实现：从路径中提取可能的设备名
	// 实际应该通过更可靠的方法获取
	parts := strings.Split(path, "/")
	for _, part := range parts {
		if len(part) > 0 && (strings.HasPrefix(part, "sd") ||
			strings.HasPrefix(part, "nvme") ||
			strings.HasPrefix(part, "hd") ||
			strings.HasPrefix(part, "vd")) {
			return part
		}
	}
	return ""
}