forge/pkg/syscheck/disk.go

package syscheck

import (
	"context"
	"crypto/rand"
	"fmt"
	"io"
	"os"
	"strconv"
	"time"

	"golang.org/x/sys/unix"
)

// GetDiskSpace returns available disk space in bytes for the root partition
func GetDiskSpace(ctx context.Context) (int64, error) {
	var stat unix.Statfs_t
	if err := unix.Statfs("/", &stat); err != nil {
		return 0, fmt.Errorf("failed to get disk space: %w", err)
	}

	// Available space = Available blocks * Block size
	availableSpace := int64(stat.Bavail) * int64(stat.Bsize)
	return availableSpace, nil
}

// GetDiskSpeed measures disk read/write speed by writing and reading a 1GB test file
// Returns read speed and write speed in bytes per second
func GetDiskSpeed(ctx context.Context) (int64, int64, error) {
	const (
		testSize   = 1024 * 1024 * 1024 // 1GB
		bufferSize = 1024 * 1024        // 1MB buffer
	)

	tmpFile := "/tmp/diskspeed_test_" + strconv.FormatInt(time.Now().UnixNano(), 10)
	defer func() {
		// Clean up test file
		_ = os.Remove(tmpFile)
	}()

	// Test write speed
	writeSpeed, err := measureWriteSpeed(ctx, tmpFile, testSize, bufferSize)
	if err != nil {
		return 0, 0, fmt.Errorf("failed to measure write speed: %w", err)
	}

	// Test read speed
	readSpeed, err := measureReadSpeed(ctx, tmpFile, bufferSize)
	if err != nil {
		return 0, 0, fmt.Errorf("failed to measure read speed: %w", err)
	}

	return readSpeed, writeSpeed, nil
}

// measureWriteSpeed writes test data to a file and measures the speed
func measureWriteSpeed(ctx context.Context, filename string, totalSize, bufferSize int64) (int64, error) {
	// Create file
	file, err := os.Create(filename)
	if err != nil {
		return 0, fmt.Errorf("failed to create test file: %w", err)
	}
	defer file.Close()

	// Prepare buffer with random data
	buffer := make([]byte, bufferSize)
	if _, err := rand.Read(buffer); err != nil {
		return 0, fmt.Errorf("failed to generate random data: %w", err)
	}

	// Start timing
	startTime := time.Now()
	var written int64

	// Write data in chunks
	for written < totalSize {
		// Check context cancellation
		select {
		case <-ctx.Done():
			return 0, ctx.Err()
		default:
		}

		// Write buffer
		n, err := file.Write(buffer)
		if err != nil {
			return 0, fmt.Errorf("failed to write data: %w", err)
		}
		written += int64(n)
	}

	// Sync to ensure data is written to disk
	if err := file.Sync(); err != nil {
		return 0, fmt.Errorf("failed to sync file: %w", err)
	}

	// Calculate speed
	duration := time.Since(startTime)
	speed := int64(float64(written) / duration.Seconds())

	return speed, nil
}

// measureReadSpeed reads the test file and measures the speed
func measureReadSpeed(ctx context.Context, filename string, bufferSize int64) (int64, error) {
	// Open file
	file, err := os.Open(filename)
	if err != nil {
		return 0, fmt.Errorf("failed to open test file: %w", err)
	}
	defer file.Close()

	// Get file size
	fileInfo, err := file.Stat()
	if err != nil {
		return 0, fmt.Errorf("failed to stat file: %w", err)
	}
	totalSize := fileInfo.Size()

	// Prepare buffer
	buffer := make([]byte, bufferSize)

	// Start timing
	startTime := time.Now()
	var totalRead int64

	// Read data in chunks
	for {
		// Check context cancellation
		select {
		case <-ctx.Done():
			return 0, ctx.Err()
		default:
		}

		// Read buffer
		n, err := file.Read(buffer)
		if err == io.EOF {
			break
		}
		if err != nil {
			return 0, fmt.Errorf("failed to read data: %w", err)
		}
		totalRead += int64(n)
	}

	// Calculate speed
	duration := time.Since(startTime)
	speed := int64(float64(totalRead) / duration.Seconds())

	// Verify we read the entire file
	if totalRead != totalSize {
		return 0, fmt.Errorf("read size mismatch: expected %d, got %d", totalSize, totalRead)
	}

	return speed, nil
}