Skip to content

Commit cdeec01

Browse files
TingluoHuangTingluoHuang
authored
Enable hostcontext to track auth migration. (#3776)
1 parent 2cb1f94 commit cdeec01

File tree

4 files changed

+253
-8
lines changed

4 files changed

+253
-8
lines changed

‎src/Runner.Common/AuthMigration.cs

+13
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
using System;
2+
3+
namespace GitHub.Runner.Common
4+
{
5+
public class AuthMigrationEventArgs : EventArgs
6+
{
7+
public AuthMigrationEventArgs(string trace)
8+
{
9+
Trace = trace;
10+
}
11+
public string Trace { get; private set; }
12+
}
13+
}

‎src/Runner.Common/HostContext.cs

+91
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ public interface IHostContext : IDisposable
3737
void ShutdownRunner(ShutdownReason reason);
3838
void WritePerfCounter(string counter);
3939
void LoadDefaultUserAgents();
40+
41+
bool AllowAuthMigration { get; }
42+
void EnableAuthMigration(string trace);
43+
void DeferAuthMigration(TimeSpan deferred, string trace);
44+
event EventHandler<AuthMigrationEventArgs> AuthMigrationChanged;
4045
}
4146

4247
public enum StartupType
@@ -70,12 +75,21 @@ public sealed class HostContext : EventListener, IObserver<DiagnosticListener>,
7075
private RunnerWebProxy _webProxy = new();
7176
private string _hostType = string.Empty;
7277

78+
// disable auth migration by default
79+
private readonly ManualResetEventSlim _allowAuthMigration = new ManualResetEventSlim(false);
80+
private DateTime _deferredAuthMigrationTime = DateTime.MaxValue;
81+
private readonly object _authMigrationLock = new object();
82+
private CancellationTokenSource _authMigrationAutoReenableTaskCancellationTokenSource = new();
83+
private Task _authMigrationAutoReenableTask;
84+
7385
public event EventHandler Unloading;
86+
public event EventHandler<AuthMigrationEventArgs> AuthMigrationChanged;
7487
public CancellationToken RunnerShutdownToken => _runnerShutdownTokenSource.Token;
7588
public ShutdownReason RunnerShutdownReason { get; private set; }
7689
public ISecretMasker SecretMasker => _secretMasker;
7790
public List<ProductInfoHeaderValue> UserAgents => _userAgents;
7891
public RunnerWebProxy WebProxy => _webProxy;
92+
public bool AllowAuthMigration => _allowAuthMigration.IsSet;
7993
public HostContext(string hostType, string logFile = null)
8094
{
8195
// Validate args.
@@ -207,6 +221,71 @@ public HostContext(string hostType, string logFile = null)
207221
LoadDefaultUserAgents();
208222
}
209223

224+
// marked as internal for testing
225+
internal async Task AuthMigrationAuthReenableAsync(TimeSpan refreshInterval, CancellationToken token)
226+
{
227+
try
228+
{
229+
while (!token.IsCancellationRequested)
230+
{
231+
_trace.Verbose($"Auth migration defer timer is set to expire at {_deferredAuthMigrationTime.ToString("O")}. AllowAuthMigration: {_allowAuthMigration.IsSet}.");
232+
await Task.Delay(refreshInterval, token);
233+
if (!_allowAuthMigration.IsSet && DateTime.UtcNow > _deferredAuthMigrationTime)
234+
{
235+
_trace.Info($"Auth migration defer timer expired. Allowing auth migration.");
236+
EnableAuthMigration("Auth migration defer timer expired.");
237+
}
238+
}
239+
}
240+
catch (TaskCanceledException)
241+
{
242+
// Task was cancelled, exit the loop.
243+
}
244+
catch (Exception ex)
245+
{
246+
_trace.Info("Error in auth migration reenable task.");
247+
_trace.Error(ex);
248+
}
249+
}
250+
251+
public void EnableAuthMigration(string trace)
252+
{
253+
_allowAuthMigration.Set();
254+
255+
lock (_authMigrationLock)
256+
{
257+
if (_authMigrationAutoReenableTask == null)
258+
{
259+
var refreshIntervalInMS = 60 * 1000;
260+
#if DEBUG
261+
// For L0, we will refresh faster
262+
if (!string.IsNullOrEmpty(Environment.GetEnvironmentVariable("_GITHUB_ACTION_AUTH_MIGRATION_REFRESH_INTERVAL")))
263+
{
264+
refreshIntervalInMS = int.Parse(Environment.GetEnvironmentVariable("_GITHUB_ACTION_AUTH_MIGRATION_REFRESH_INTERVAL"));
265+
}
266+
#endif
267+
_authMigrationAutoReenableTask = AuthMigrationAuthReenableAsync(TimeSpan.FromMilliseconds(refreshIntervalInMS), _authMigrationAutoReenableTaskCancellationTokenSource.Token);
268+
}
269+
}
270+
271+
_trace.Info($"Enable auth migration at {_deferredAuthMigrationTime.ToString("O")}.");
272+
AuthMigrationChanged?.Invoke(this, new AuthMigrationEventArgs(trace));
273+
}
274+
275+
public void DeferAuthMigration(TimeSpan deferred, string trace)
276+
{
277+
_allowAuthMigration.Reset();
278+
279+
// defer migration for a while
280+
lock (_authMigrationLock)
281+
{
282+
_deferredAuthMigrationTime = DateTime.UtcNow.Add(deferred);
283+
}
284+
285+
_trace.Info($"Disabled auth migration until {_deferredAuthMigrationTime.ToString("O")}.");
286+
AuthMigrationChanged?.Invoke(this, new AuthMigrationEventArgs(trace));
287+
}
288+
210289
public void LoadDefaultUserAgents()
211290
{
212291
if (string.IsNullOrEmpty(WebProxy.HttpProxyAddress) && string.IsNullOrEmpty(WebProxy.HttpsProxyAddress))
@@ -549,6 +628,18 @@ private void Dispose(bool disposing)
549628
_loadContext.Unloading -= LoadContext_Unloading;
550629
_loadContext = null;
551630
}
631+
632+
if (_authMigrationAutoReenableTask != null)
633+
{
634+
_authMigrationAutoReenableTaskCancellationTokenSource?.Cancel();
635+
}
636+
637+
if (_authMigrationAutoReenableTaskCancellationTokenSource != null)
638+
{
639+
_authMigrationAutoReenableTaskCancellationTokenSource?.Dispose();
640+
_authMigrationAutoReenableTaskCancellationTokenSource = null;
641+
}
642+
552643
_httpTraceSubscription?.Dispose();
553644
_diagListenerSubscription?.Dispose();
554645
_traceManager?.Dispose();

‎src/Test/L0/HostContextL0.cs

+129-2
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
1-
using GitHub.Runner.Common.Util;
2-
using System;
1+
using System;
32
using System.IO;
43
using System.Reflection;
54
using System.Runtime.CompilerServices;
65
using System.Text;
76
using System.Threading;
7+
using System.Threading.Tasks;
88
using Xunit;
99

1010
namespace GitHub.Runner.Common.Tests
@@ -172,6 +172,133 @@ public void SecretMaskerForProxy()
172172
}
173173
}
174174

175+
[Fact]
176+
[Trait("Level", "L0")]
177+
[Trait("Category", "Common")]
178+
public void AuthMigrationDisabledByDefault()
179+
{
180+
try
181+
{
182+
Environment.SetEnvironmentVariable("_GITHUB_ACTION_AUTH_MIGRATION_REFRESH_INTERVAL", "100");
183+
184+
// Arrange.
185+
Setup();
186+
187+
// Assert.
188+
Assert.False(_hc.AllowAuthMigration);
189+
190+
// Change migration state is error free.
191+
_hc.EnableAuthMigration("L0Test");
192+
_hc.DeferAuthMigration(TimeSpan.FromHours(1), "L0Test");
193+
}
194+
finally
195+
{
196+
Environment.SetEnvironmentVariable("_GITHUB_ACTION_AUTH_MIGRATION_REFRESH_INTERVAL", null);
197+
// Cleanup.
198+
Teardown();
199+
}
200+
}
201+
202+
[Fact]
203+
[Trait("Level", "L0")]
204+
[Trait("Category", "Common")]
205+
public async Task AuthMigrationReenableTaskNotRunningByDefault()
206+
{
207+
try
208+
{
209+
Environment.SetEnvironmentVariable("_GITHUB_ACTION_AUTH_MIGRATION_REFRESH_INTERVAL", "50");
210+
211+
// Arrange.
212+
Setup();
213+
214+
// Assert.
215+
Assert.False(_hc.AllowAuthMigration);
216+
await Task.Delay(TimeSpan.FromMilliseconds(200));
217+
}
218+
finally
219+
{
220+
Environment.SetEnvironmentVariable("_GITHUB_ACTION_AUTH_MIGRATION_REFRESH_INTERVAL", null);
221+
// Cleanup.
222+
Teardown();
223+
}
224+
225+
var logFile = Path.Combine(Path.GetDirectoryName(Assembly.GetEntryAssembly().Location), $"trace_{nameof(HostContextL0)}_{nameof(AuthMigrationReenableTaskNotRunningByDefault)}.log");
226+
var logContent = await File.ReadAllTextAsync(logFile);
227+
Assert.Contains("HostContext", logContent);
228+
Assert.DoesNotContain("Auth migration defer timer", logContent);
229+
}
230+
231+
[Fact]
232+
[Trait("Level", "L0")]
233+
[Trait("Category", "Common")]
234+
public void AuthMigrationEnableDisable()
235+
{
236+
try
237+
{
238+
// Arrange.
239+
Setup();
240+
241+
var eventFiredCount = 0;
242+
_hc.AuthMigrationChanged += (sender, e) =>
243+
{
244+
eventFiredCount++;
245+
Assert.Equal("L0Test", e.Trace);
246+
};
247+
248+
// Assert.
249+
_hc.EnableAuthMigration("L0Test");
250+
Assert.True(_hc.AllowAuthMigration);
251+
252+
_hc.DeferAuthMigration(TimeSpan.FromHours(1), "L0Test");
253+
Assert.False(_hc.AllowAuthMigration);
254+
Assert.Equal(2, eventFiredCount);
255+
}
256+
finally
257+
{
258+
// Cleanup.
259+
Teardown();
260+
}
261+
}
262+
263+
[Fact]
264+
[Trait("Level", "L0")]
265+
[Trait("Category", "Common")]
266+
public async Task AuthMigrationAutoReset()
267+
{
268+
try
269+
{
270+
Environment.SetEnvironmentVariable("_GITHUB_ACTION_AUTH_MIGRATION_REFRESH_INTERVAL", "100");
271+
272+
// Arrange.
273+
Setup();
274+
275+
var eventFiredCount = 0;
276+
_hc.AuthMigrationChanged += (sender, e) =>
277+
{
278+
eventFiredCount++;
279+
Assert.NotEmpty(e.Trace);
280+
};
281+
282+
// Assert.
283+
_hc.EnableAuthMigration("L0Test");
284+
Assert.True(_hc.AllowAuthMigration);
285+
286+
_hc.DeferAuthMigration(TimeSpan.FromMilliseconds(500), "L0Test");
287+
Assert.False(_hc.AllowAuthMigration);
288+
289+
await Task.Delay(TimeSpan.FromSeconds(1));
290+
Assert.True(_hc.AllowAuthMigration);
291+
Assert.Equal(3, eventFiredCount);
292+
}
293+
finally
294+
{
295+
Environment.SetEnvironmentVariable("_GITHUB_ACTION_AUTH_MIGRATION_REFRESH_INTERVAL", null);
296+
297+
// Cleanup.
298+
Teardown();
299+
}
300+
}
301+
175302
private void Setup([CallerMemberName] string testName = "")
176303
{
177304
_tokenSource = new CancellationTokenSource();

‎src/Test/L0/TestHostContext.cs

+20-6
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,15 @@
1-
using GitHub.Runner.Common.Util;
2-
using System;
1+
using System;
32
using System.Collections.Concurrent;
3+
using System.Collections.Generic;
44
using System.Globalization;
55
using System.IO;
6+
using System.Net.Http.Headers;
7+
using System.Reflection;
68
using System.Runtime.CompilerServices;
9+
using System.Runtime.Loader;
710
using System.Threading;
811
using System.Threading.Tasks;
9-
using System.Runtime.Loader;
10-
using System.Reflection;
11-
using System.Collections.Generic;
1212
using GitHub.DistributedTask.Logging;
13-
using System.Net.Http.Headers;
1413
using GitHub.Runner.Sdk;
1514

1615
namespace GitHub.Runner.Common.Tests
@@ -31,6 +30,7 @@ public sealed class TestHostContext : IHostContext, IDisposable
3130
private StartupType _startupType;
3231
public event EventHandler Unloading;
3332
public event EventHandler<DelayEventArgs> Delaying;
33+
public event EventHandler<AuthMigrationEventArgs> AuthMigrationChanged;
3434
public CancellationToken RunnerShutdownToken => _runnerShutdownTokenSource.Token;
3535
public ShutdownReason RunnerShutdownReason { get; private set; }
3636
public ISecretMasker SecretMasker => _secretMasker;
@@ -92,6 +92,8 @@ public StartupType StartupType
9292

9393
public RunnerWebProxy WebProxy => new();
9494

95+
public bool AllowAuthMigration { get; set; }
96+
9597
public async Task Delay(TimeSpan delay, CancellationToken token)
9698
{
9799
// Event callback
@@ -387,6 +389,18 @@ public void LoadDefaultUserAgents()
387389
{
388390
return;
389391
}
392+
393+
public void EnableAuthMigration(string trace)
394+
{
395+
AllowAuthMigration = true;
396+
AuthMigrationChanged?.Invoke(this, new AuthMigrationEventArgs(trace));
397+
}
398+
399+
public void DeferAuthMigration(TimeSpan deferred, string trace)
400+
{
401+
AllowAuthMigration = false;
402+
AuthMigrationChanged?.Invoke(this, new AuthMigrationEventArgs(trace));
403+
}
390404
}
391405

392406
public class DelayEventArgs : EventArgs

0 commit comments

Comments
 (0)