BQP/srcs/RobotNet10/Commons/RobotNet10.Realtime/CpuAffinity.cs

using System;
using System.Collections.Generic;
using System.Linq;

namespace RobotNet10.Realtime;

/// <summary>
/// Provides CPU affinity management for Linux real-time applications.
/// </summary>
public class CpuAffinity : IDisposable
{
    private LinuxNative.cpu_set_t _cpuSet;
    private bool _disposed = false;

    /// <summary>
    /// Initializes a new instance of CpuAffinity.
    /// </summary>
    public CpuAffinity()
    {
        LinuxNative.CPU_ZERO(ref _cpuSet);
    }

    /// <summary>
    /// Adds a CPU to the affinity set.
    /// </summary>
    /// <param name="cpu">CPU number (0-based)</param>
    public void AddCpu(int cpu)
    {
        if (cpu < 0 || cpu >= LinuxNative.CPU_SETSIZE)
        {
            throw new ArgumentOutOfRangeException(nameof(cpu), 
                $"CPU number must be between 0 and {LinuxNative.CPU_SETSIZE - 1}");
        }
        LinuxNative.CPU_SET(cpu, ref _cpuSet);
    }

    /// <summary>
    /// Removes a CPU from the affinity set.
    /// </summary>
    /// <param name="cpu">CPU number (0-based)</param>
    public void RemoveCpu(int cpu)
    {
        if (cpu < 0 || cpu >= LinuxNative.CPU_SETSIZE)
        {
            throw new ArgumentOutOfRangeException(nameof(cpu), 
                $"CPU number must be between 0 and {LinuxNative.CPU_SETSIZE - 1}");
        }
        LinuxNative.CPU_CLR(cpu, ref _cpuSet);
    }

    /// <summary>
    /// Checks if a CPU is in the affinity set.
    /// </summary>
    /// <param name="cpu">CPU number (0-based)</param>
    /// <returns>True if CPU is in the set, false otherwise</returns>
    public bool IsCpuSet(int cpu)
    {
        if (cpu < 0 || cpu >= LinuxNative.CPU_SETSIZE)
        {
            return false;
        }
        return LinuxNative.CPU_ISSET(cpu, ref _cpuSet) != 0;
    }

    /// <summary>
    /// Gets the number of CPUs in the affinity set.
    /// </summary>
    public int CpuCount => LinuxNative.CPU_COUNT(ref _cpuSet);

    /// <summary>
    /// Gets all CPUs in the affinity set.
    /// </summary>
    /// <returns>List of CPU numbers</returns>
    public List<int> GetCpus()
    {
        var cpus = new List<int>();
        for (int i = 0; i < LinuxNative.CPU_SETSIZE; i++)
        {
            if (IsCpuSet(i))
            {
                cpus.Add(i);
            }
        }
        return cpus;
    }

    /// <summary>
    /// Sets CPU affinity for the current thread.
    /// </summary>
    /// <param name="cpus">List of CPU numbers to set</param>
    public static void SetAffinity(params int[] cpus)
    {
        var affinity = new CpuAffinity();
        foreach (var cpu in cpus)
        {
            affinity.AddCpu(cpu);
        }
        affinity.Apply();
    }

    /// <summary>
    /// Applies the CPU affinity to the current thread.
    /// </summary>
    public void Apply()
    {
        if (_disposed)
        {
            throw new ObjectDisposedException(nameof(CpuAffinity));
        }

        IntPtr cpusetsize = new IntPtr(LinuxNative.CPU_SETSIZE / 8); // Size in bytes
        int result = LinuxNative.sched_setaffinity(0, cpusetsize, ref _cpuSet);
        if (result != 0)
        {
            LinuxNative.ThrowLastError("sched_setaffinity");
        }
    }

    /// <summary>
    /// Gets the current CPU affinity for the current thread.
    /// </summary>
    /// <returns>CpuAffinity object representing current affinity</returns>
    public static CpuAffinity GetAffinity()
    {
        var cpuSet = new LinuxNative.cpu_set_t();
        IntPtr cpusetsize = new IntPtr(LinuxNative.CPU_SETSIZE / 8);
        int result = LinuxNative.sched_getaffinity(0, cpusetsize, ref cpuSet);
        if (result != 0)
        {
            LinuxNative.ThrowLastError("sched_getaffinity");
        }

        var affinity = new CpuAffinity();
        affinity._cpuSet = cpuSet;
        return affinity;
    }

    /// <summary>
    /// Clears all CPUs from the affinity set.
    /// </summary>
    public void Clear()
    {
        LinuxNative.CPU_ZERO(ref _cpuSet);
    }

    public void Dispose()
    {
        if (!_disposed)
        {
            _disposed = true;
        }
    }
}