LenovoLegionToolkit

1
using System;
2
using System.Collections.Generic;
3
using System.Diagnostics;
4
using System.Linq;
5
using System.Threading;
6
using System.Threading.Tasks;
7
using LenovoLegionToolkit.Lib.Extensions;
8
using LenovoLegionToolkit.Lib.Resources;
9
using LenovoLegionToolkit.Lib.System;
10
using LenovoLegionToolkit.Lib.System.Management;
11
using LenovoLegionToolkit.Lib.Utils;
12
using NeoSmart.AsyncLock;
13

14
namespace LenovoLegionToolkit.Lib.Controllers;
15

16
public class GPUController
17
{
18
    private readonly AsyncLock _lock = new();
19

20
    private Task? _refreshTask;
21
    private CancellationTokenSource? _refreshCancellationTokenSource;
22

23
    private GPUState _state = GPUState.Unknown;
24
    private List<Process> _processes = new();
25
    private string? _gpuInstanceId;
26
    private string? _performanceState;
27

28
    public GPUState LastKnownState => _state;
29
    public event EventHandler<GPUStatus>? Refreshed;
30

31
    public bool IsSupported()
32
    {
33
        try
34
        {
35
            NVAPI.Initialize();
36
            return NVAPI.GetGPU() is not null;
37
        }
38
        catch
39
        {
40
            return false;
41
        }
42
        finally
43
        {
44
            try
45
            {
46
                NVAPI.Unload();
47
            }
48
            catch { /* Ignored. */ }
49
        }
50
    }
51

52
    public async Task<GPUStatus> RefreshNowAsync()
53
    {
54
        using (await _lock.LockAsync().ConfigureAwait(false))
55
        {
56
            await RefreshLoopAsync(0, 0, CancellationToken.None).ConfigureAwait(false);
57
            return new GPUStatus(_state, _performanceState, _processes);
58
        }
59
    }
60

61
    public async Task StartAsync(int delay = 1_000, int interval = 5_000)
62
    {
63
        await StopAsync(true).ConfigureAwait(false);
64

65
        if (Log.Instance.IsTraceEnabled)
66
            Log.Instance.Trace($"Starting... [delay={delay}, interval={interval}]");
67

68
        _refreshCancellationTokenSource = new CancellationTokenSource();
69
        var token = _refreshCancellationTokenSource.Token;
70
        _refreshTask = Task.Run(() => RefreshLoopAsync(delay, interval, token), token);
71
    }
72

73
    public async Task StopAsync(bool waitForFinish = false)
74
    {
75
        if (Log.Instance.IsTraceEnabled)
76
            Log.Instance.Trace($"Stopping... [refreshTask.isNull={_refreshTask is null}, _refreshCancellationTokenSource.IsCancellationRequested={_refreshCancellationTokenSource?.IsCancellationRequested}]");
77

78
        _refreshCancellationTokenSource?.Cancel();
79

80
        if (waitForFinish)
81
        {
82
            if (Log.Instance.IsTraceEnabled)
83
                Log.Instance.Trace($"Waiting to finish...");
84

85
            if (_refreshTask is not null)
86
            {
87
                try
88
                {
89
                    await _refreshTask.ConfigureAwait(false);
90
                }
91
                catch (OperationCanceledException) { }
92
            }
93

94
            if (Log.Instance.IsTraceEnabled)
95
                Log.Instance.Trace($"Finished");
96
        }
97

98
        _refreshCancellationTokenSource = null;
99
        _refreshTask = null;
100

101
        if (Log.Instance.IsTraceEnabled)
102
            Log.Instance.Trace($"Stopped");
103
    }
104

105
    public async Task RestartGPUAsync()
106
    {
107
        using (await _lock.LockAsync().ConfigureAwait(false))
108
        {
109
            if (Log.Instance.IsTraceEnabled)
110
                Log.Instance.Trace($"Deactivating... [state={_state}, gpuInstanceId={_gpuInstanceId}]");
111

112
            if (_state is not GPUState.Active and not GPUState.Inactive)
113
                return;
114

115
            if (string.IsNullOrEmpty(_gpuInstanceId))
116
                return;
117

118
            await CMD.RunAsync("pnputil", $"/restart-device \"{_gpuInstanceId}\"").ConfigureAwait(false);
119

120
            if (Log.Instance.IsTraceEnabled)
121
                Log.Instance.Trace($"Deactivating... [state= {_state}, gpuInstanceId={_gpuInstanceId}]");
122
        }
123
    }
124

125
    public async Task KillGPUProcessesAsync()
126
    {
127
        using (await _lock.LockAsync().ConfigureAwait(false))
128
        {
129
            if (Log.Instance.IsTraceEnabled)
130
                Log.Instance.Trace($"Deactivating... [state= {_state}, gpuInstanceId={_gpuInstanceId}]");
131

132
            if (_state is not GPUState.Active)
133
                return;
134

135
            if (string.IsNullOrEmpty(_gpuInstanceId))
136
                return;
137

138
            foreach (var process in _processes)
139
            {
140
                try
141
                {
142
                    process.Kill(true);
143
                    await process.WaitForExitAsync().ConfigureAwait(false);
144
                }
145
                catch (Exception ex)
146
                {
147
                    if (Log.Instance.IsTraceEnabled)
148
                        Log.Instance.Trace($"Couldn't kill process. [pid={process.Id}, name={process.ProcessName}]", ex);
149
                }
150
            }
151

152
            if (Log.Instance.IsTraceEnabled)
153
                Log.Instance.Trace($"Deactivating... [state=  {_state}, gpuInstanceId={_gpuInstanceId}]");
154
        }
155
    }
156

157
    private async Task RefreshLoopAsync(int delay, int interval, CancellationToken token)
158
    {
159
        try
160
        {
161
            if (Log.Instance.IsTraceEnabled)
162
                Log.Instance.Trace($"Initializing NVAPI...");
163

164
            NVAPI.Initialize();
165

166
            if (Log.Instance.IsTraceEnabled)
167
                Log.Instance.Trace($"Initialized NVAPI");
168

169
            await Task.Delay(delay, token).ConfigureAwait(false);
170

171
            while (true)
172
            {
173
                token.ThrowIfCancellationRequested();
174

175
                using (await _lock.LockAsync(token).ConfigureAwait(false))
176
                {
177

178
                    if (Log.Instance.IsTraceEnabled)
179
                        Log.Instance.Trace($"Will refresh...");
180

181
                    await RefreshStateAsync().ConfigureAwait(false);
182

183
                    if (Log.Instance.IsTraceEnabled)
184
                        Log.Instance.Trace($"Refreshed");
185

186
                    Refreshed?.Invoke(this, new GPUStatus(_state, _performanceState, _processes));
187
                }
188

189
                if (interval > 0)
190
                    await Task.Delay(interval, token).ConfigureAwait(false);
191
                else
192
                    break;
193
            }
194
        }
195
        catch (Exception ex) when (ex is not OperationCanceledException)
196
        {
197
            if (Log.Instance.IsTraceEnabled)
198
                Log.Instance.Trace($"Exception occurred", ex);
199

200
            throw;
201
        }
202
        finally
203
        {
204
            if (Log.Instance.IsTraceEnabled)
205
                Log.Instance.Trace($"Unloading NVAPI...");
206

207
            NVAPI.Unload();
208

209
            if (Log.Instance.IsTraceEnabled)
210
                Log.Instance.Trace($"Unloaded NVAPI");
211
        }
212
    }
213

214
    private async Task RefreshStateAsync()
215
    {
216
        if (Log.Instance.IsTraceEnabled)
217
            Log.Instance.Trace($"Refresh in progress...");
218

219
        _state = GPUState.Unknown;
220
        _processes = new();
221
        _gpuInstanceId = null;
222
        _performanceState = null;
223

224
        var gpu = NVAPI.GetGPU();
225
        if (gpu is null)
226
        {
227
            _state = GPUState.NvidiaGpuNotFound;
228

229
            if (Log.Instance.IsTraceEnabled)
230
                Log.Instance.Trace($"GPU present [state={_state}, processes.Count={_processes.Count}, gpuInstanceId={_gpuInstanceId}]");
231

232
            return;
233
        }
234

235
        try
236
        {
237
            var stateId = gpu.PerformanceStatesInfo.CurrentPerformanceState.StateId.ToString().GetUntilOrEmpty("_");
238
            _performanceState = Resource.GPUController_PoweredOn;
239
            if (!string.IsNullOrWhiteSpace(stateId))
240
                _performanceState += $", {stateId}";
241
        }
242
        catch (Exception ex) when (ex.Message == "NVAPI_GPU_NOT_POWERED")
243
        {
244
            _state = GPUState.PoweredOff;
245
            _performanceState = Resource.GPUController_PoweredOff;
246

247
            if (Log.Instance.IsTraceEnabled)
248
                Log.Instance.Trace($"Powered off [state={_state}, processes.Count={_processes.Count}, gpuInstanceId={_gpuInstanceId}]");
249

250
            return;
251
        }
252
        catch (Exception ex)
253
        {
254
            if (Log.Instance.IsTraceEnabled)
255
                Log.Instance.Trace($"GPU status exception.", ex);
256

257
            _performanceState = "Unknown";
258
        }
259

260
        var pnpDeviceIdPart = NVAPI.GetGPUId(gpu);
261

262
        if (string.IsNullOrEmpty(pnpDeviceIdPart))
263
            throw new InvalidOperationException("pnpDeviceIdPart is null or empty");
264

265
        var gpuInstanceId = await WMI.Win32.PnpEntity.GetDeviceIDAsync(pnpDeviceIdPart).ConfigureAwait(false);
266
        var processNames = NVAPIExtensions.GetActiveProcesses(gpu);
267

268
        if (NVAPI.IsDisplayConnected(gpu))
269
        {
270
            _processes = processNames;
271
            _state = GPUState.MonitorConnected;
272

273
            if (Log.Instance.IsTraceEnabled)
274
                Log.Instance.Trace(
275
                    $"Monitor connected [state={_state}, processes.Count={_processes.Count}, gpuInstanceId={_gpuInstanceId}]");
276
        }
277
        else if (processNames.Any())
278
        {
279
            _processes = processNames;
280
            _state = GPUState.Active;
281
            _gpuInstanceId = gpuInstanceId;
282

283
            if (Log.Instance.IsTraceEnabled)
284
                Log.Instance.Trace($"Active [state={_state}, processes.Count={_processes.Count}, gpuInstanceId={_gpuInstanceId}, pnpDeviceIdPart={pnpDeviceIdPart}]");
285
        }
286
        else
287
        {
288
            _state = GPUState.Inactive;
289
            _gpuInstanceId = gpuInstanceId;
290

291
            if (Log.Instance.IsTraceEnabled)
292
                Log.Instance.Trace($"Inactive [state={_state}, processes.Count={_processes.Count}, gpuInstanceId={_gpuInstanceId}]");
293
        }
294
    }
295
}
296