From 0d6cf6682cd37349c8e18d8bbf671757f8a4c3ac Mon Sep 17 00:00:00 2001
From: Joe Schmitt <joschmit@microsoft.com>
Date: Mon, 1 Jul 2024 13:15:04 -0700
Subject: [PATCH] Nullability fixes

---
 .../Controllers/DiagController.cs                           | 4 +---
 .../Controllers/MetricsController.cs                        | 4 +---
 .../TaskCompletionSourceExtensions.cs                       | 6 ++----
 3 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/DiagController.cs b/src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/DiagController.cs
index 60bea0b1635..f4097f4509e 100644
--- a/src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/DiagController.cs
+++ b/src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/DiagController.cs
@@ -35,9 +35,7 @@ public partial class DiagController : DiagnosticsControllerBase
     {
         private const TraceProfile DefaultTraceProfiles = TraceProfile.Cpu | TraceProfile.Http | TraceProfile.Metrics | TraceProfile.GcCollect;
 
-#nullable disable
         private readonly IOptions<DiagnosticPortOptions> _diagnosticPortOptions;
-#nullable restore
         private readonly IOptions<CallStacksOptions> _callStacksOptions;
         private readonly IOptions<ParameterCapturingOptions> _parameterCapturingOptions;
         private readonly IOptionsMonitor<GlobalCounterOptions> _counterOptions;
@@ -54,7 +52,7 @@ public partial class DiagController : DiagnosticsControllerBase
         public DiagController(IServiceProvider serviceProvider, ILogger<DiagController> logger)
             : base(serviceProvider.GetRequiredService<IDiagnosticServices>(), serviceProvider.GetRequiredService<IEgressOperationStore>(), logger)
         {
-            _diagnosticPortOptions = serviceProvider.GetService<IOptions<DiagnosticPortOptions>>();
+            _diagnosticPortOptions = serviceProvider.GetRequiredService<IOptions<DiagnosticPortOptions>>();
             _callStacksOptions = serviceProvider.GetRequiredService<IOptions<CallStacksOptions>>();
             _parameterCapturingOptions = serviceProvider.GetRequiredService<IOptions<ParameterCapturingOptions>>();
             _counterOptions = serviceProvider.GetRequiredService<IOptionsMonitor<GlobalCounterOptions>>();
diff --git a/src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/MetricsController.cs b/src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/MetricsController.cs
index e4f8538b921..1b14cb49b6c 100644
--- a/src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/MetricsController.cs
+++ b/src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/MetricsController.cs
@@ -21,9 +21,7 @@ public class MetricsController : ControllerBase
         private const string ArtifactType_Metrics = "metrics";
 
         private readonly ILogger<MetricsController> _logger;
-#nullable disable
         private readonly MetricsStoreService _metricsStore;
-#nullable restore
         private readonly MetricsOptions _metricsOptions;
 
         public MetricsController(ILogger<MetricsController> logger,
@@ -31,7 +29,7 @@ public MetricsController(ILogger<MetricsController> logger,
             IOptions<MetricsOptions> metricsOptions)
         {
             _logger = logger;
-            _metricsStore = serviceProvider.GetService<MetricsStoreService>();
+            _metricsStore = serviceProvider.GetRequiredService<MetricsStoreService>();
             _metricsOptions = metricsOptions.Value;
         }
 
diff --git a/src/Microsoft.Diagnostics.Monitoring.WebApi/TaskCompletionSourceExtensions.cs b/src/Microsoft.Diagnostics.Monitoring.WebApi/TaskCompletionSourceExtensions.cs
index 321789485ef..f92df62209b 100644
--- a/src/Microsoft.Diagnostics.Monitoring.WebApi/TaskCompletionSourceExtensions.cs
+++ b/src/Microsoft.Diagnostics.Monitoring.WebApi/TaskCompletionSourceExtensions.cs
@@ -1,8 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-#nullable disable
-
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -22,7 +20,7 @@ internal static class TaskCompletionSourceExtensions
         /// </remarks>
         public static async Task<T> WithCancellation<T>(this TaskCompletionSource<T> source, CancellationToken token)
         {
-            using (token.Register(source => ((TaskCompletionSource<T>)source).TrySetCanceled(token), source))
+            using (token.Register(() => source.TrySetCanceled(token)))
             {
                 return await source.Task.ConfigureAwait(false);
             }
@@ -30,7 +28,7 @@ public static async Task<T> WithCancellation<T>(this TaskCompletionSource<T> sou
 
         public static async Task WithCancellation(this TaskCompletionSource source, CancellationToken token)
         {
-            using (token.Register(source => ((TaskCompletionSource)source).TrySetCanceled(token), source))
+            using (token.Register(() => source.TrySetCanceled(token)))
             {
                 await source.Task.ConfigureAwait(false);
             }