ProcessMemory.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Text;

namespace CoreTempGPU
{
    /// <summary>One process contribution to VRAM or RAM.</summary>
    public sealed class ProcessShare
    {
        public uint Pid;
        public string Name;
        public ulong Bytes;
        public double PercentOfUsed;
    }

    /// <summary>
    /// Cumulative top contributors until ≥67% of currently used VRAM/RAM.
    /// Cached ~1.5s to avoid UI stutter.
    /// </summary>
    internal static class ProcessMemory
    {
        public const double CoverageFraction = 0.67;
        private static readonly object _lock = new object();
        private static List<ProcessShare> _vramCache;
        private static List<ProcessShare> _ramCache;
        private static DateTime _vramCacheAt = DateTime.MinValue;
        private static DateTime _ramCacheAt = DateTime.MinValue;
        private static readonly TimeSpan CacheTtl = TimeSpan.FromMilliseconds(1500);

        public static List<ProcessShare> GetVramTopContributors(IntPtr nvmlDevice, ulong usedBytes)
        {
            lock (_lock)
            {
                if (_vramCache != null && (DateTime.UtcNow - _vramCacheAt) < CacheTtl)
                    return _vramCache;
            }

            List<ProcessShare> result;
            try
            {
                var raw = new List<Tuple<uint, ulong>>();
                if (nvmlDevice != IntPtr.Zero && Nvml.Available)
                {
                    foreach (var p in Nvml.GetMergedRunningProcesses(nvmlDevice))
                    {
                        if (p.usedGpuMemory > 0)
                            raw.Add(Tuple.Create(p.pid, p.usedGpuMemory));
                    }
                }
                if (raw.Count == 0)
                    raw = QueryPdhGpuProcessMemory();

                result = BuildCumulative(raw, usedBytes);
            }
            catch
            {
                result = new List<ProcessShare>();
            }

            lock (_lock)
            {
                _vramCache = result;
                _vramCacheAt = DateTime.UtcNow;
            }
            return result;
        }

        public static List<ProcessShare> GetRamTopContributors(ulong usedBytes)
        {
            lock (_lock)
            {
                if (_ramCache != null && (DateTime.UtcNow - _ramCacheAt) < CacheTtl)
                    return _ramCache;
            }

            List<ProcessShare> result;
            try
            {
                var raw = new List<Tuple<uint, ulong>>();
                foreach (Process proc in Process.GetProcesses())
                {
                    try
                    {
                        long ws = proc.WorkingSet64;
                        if (ws <= 0) continue;
                        raw.Add(Tuple.Create((uint)proc.Id, (ulong)ws));
                    }
                    catch { }
                    finally
                    {
                        try { proc.Dispose(); } catch { }
                    }
                }
                result = BuildCumulative(raw, usedBytes);
            }
            catch
            {
                result = new List<ProcessShare>();
            }

            lock (_lock)
            {
                _ramCache = result;
                _ramCacheAt = DateTime.UtcNow;
            }
            return result;
        }

        private static List<ProcessShare> BuildCumulative(List<Tuple<uint, ulong>> raw, ulong usedBytes)
        {
            var list = new List<ProcessShare>();
            if (raw == null || raw.Count == 0 || usedBytes == 0) return list;

            raw.Sort((a, b) => b.Item2.CompareTo(a.Item2));
            ulong cum = 0;
            int n = 0;
            foreach (var t in raw)
            {
                cum += t.Item2;
                list.Add(new ProcessShare
                {
                    Pid = t.Item1,
                    Name = ResolveName(t.Item1),
                    Bytes = t.Item2,
                    PercentOfUsed = 100.0 * t.Item2 / usedBytes
                });
                n++;
                if (cum >= usedBytes * CoverageFraction) break;
                if (n >= 12) break;
            }
            return list;
        }

        private static string ResolveName(uint pid)
        {
            try
            {
                using (var p = Process.GetProcessById((int)pid))
                {
                    string n = p.ProcessName;
                    if (string.IsNullOrEmpty(n)) return "PID " + pid;
                    return n.EndsWith(".exe", StringComparison.OrdinalIgnoreCase) ? n : n + ".exe";
                }
            }
            catch
            {
                return "PID " + pid;
            }
        }

        /// <summary>PDH fallback: \\GPU Process Memory(*)\\Local Usage instance names contain pid_.</summary>
        private static List<Tuple<uint, ulong>> QueryPdhGpuProcessMemory()
        {
            var list = new List<Tuple<uint, ulong>>();
            try
            {
                // Lightweight: sample via nvidia-less path using PerformanceCounter is heavy;
                // use ExpandWildCard + one-shot query for Local Usage.
                var paths = PdhExpand(@"\GPU Process Memory(*)\Local Usage");
                if (paths == null || paths.Count == 0) return list;

                using (var q = Pdh.OpenQuery(paths))
                {
                    if (q == null) return list;
                    q.Collect();
                    System.Threading.Thread.Sleep(200);
                    q.Collect();
                    for (int i = 0; i < q.Counters.Count; i++)
                    {
                        double v;
                        if (!q.Counters[i].TryRead(out v) || v <= 0) continue;
                        uint pid = ParsePidFromInstance(q.Counters[i].Path);
                        if (pid == 0) continue;
                        list.Add(Tuple.Create(pid, (ulong)v));
                    }
                }
            }
            catch { }
            return list;
        }

        private static List<string> PdhExpand(string wild)
        {
            var result = new List<string>();
            try
            {
                int sz = 0;
                int rc = PdhExpandWildCardPathW(null, wild, null, ref sz, 0);
                if (sz <= 0) return result;
                var buf = new char[sz];
                rc = PdhExpandWildCardPathW(null, wild, buf, ref sz, 0);
                if (rc != 0) return result;
                var cur = new StringBuilder();
                for (int i = 0; i < buf.Length; i++)
                {
                    if (buf[i] == '\0')
                    {
                        if (cur.Length > 0) { result.Add(cur.ToString()); cur.Length = 0; }
                        else break;
                    }
                    else cur.Append(buf[i]);
                }
            }
            catch { }
            return result;
        }

        [DllImport("pdh.dll", CharSet = CharSet.Unicode)]
        private static extern int PdhExpandWildCardPathW(string dataSource, string wildCardPath, char[] expandedPathList, ref int pathListLength, uint flags);

        private static uint ParsePidFromInstance(string path)
        {
            // ...\GPU Process Memory(pid_1234_luid_...)\Local Usage
            if (string.IsNullOrEmpty(path)) return 0;
            int i = path.IndexOf("pid_", StringComparison.OrdinalIgnoreCase);
            if (i < 0) return 0;
            i += 4;
            int j = i;
            while (j < path.Length && char.IsDigit(path[j])) j++;
            uint pid;
            if (uint.TryParse(path.Substring(i, j - i), out pid)) return pid;
            return 0;
        }

        public static string FormatTooltip(string title, List<ProcessShare> shares, ulong usedBytes)
        {
            var sb = new StringBuilder();
            sb.Append(title);
            if (shares == null || shares.Count == 0)
            {
                sb.AppendLine();
                sb.Append("No process breakdown available.");
                return sb.ToString();
            }
            sb.AppendLine();
            sb.AppendLine("(≥67% of " + GpuSnapshot.FormatBytes(usedBytes) + " used)");
            int i = 1;
            foreach (var s in shares)
            {
                sb.Append('#');
                sb.Append(i++);
                sb.Append(' ');
                sb.Append(s.Name);
                sb.Append(":  ");
                sb.Append(GpuSnapshot.FormatBytes(s.Bytes));
                sb.Append(" (");
                sb.Append(s.PercentOfUsed.ToString("0.0"));
                sb.AppendLine("%)");
            }
            string text = sb.ToString();
            // ToolTip practical limit ~hundreds of chars; keep under ~900.
            if (text.Length > 900) text = text.Substring(0, 897) + "...";
            return text;
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct MemoryStatusEx
        {
            public uint Length;
            public uint MemoryLoad;
            public ulong TotalPhys;
            public ulong AvailPhys;
            public ulong TotalPageFile;
            public ulong AvailPageFile;
            public ulong TotalVirtual;
            public ulong AvailVirtual;
            public ulong AvailExtendedVirtual;
        }

        [DllImport("kernel32.dll", SetLastError = true)]
        private static extern bool GlobalMemoryStatusEx(ref MemoryStatusEx lpBuffer);

        public static bool TryGetSystemRam(out ulong total, out ulong used)
        {
            total = 0;
            used = 0;
            try
            {
                var st = new MemoryStatusEx();
                st.Length = (uint)Marshal.SizeOf(typeof(MemoryStatusEx));
                if (!GlobalMemoryStatusEx(ref st)) return false;
                total = st.TotalPhys;
                if (st.TotalPhys >= st.AvailPhys)
                    used = st.TotalPhys - st.AvailPhys;
                return true;
            }
            catch { return false; }
        }
    }
}