From beda8b202dc338a0be203d65da7b2863b004dd83 Mon Sep 17 00:00:00 2001
From: Andrei Ionescu <t-andionescu@microsoft.com>
Date: Mon, 6 Jan 2025 12:03:22 +0000
Subject: [PATCH] Use extensions to set health check result

---
 src/Abstractions/Activities/ActivityExtensions.cs         | 8 ++++++++
 .../Microsoft.Omex.Extensions.Abstractions.csproj         | 1 +
 .../Composables/ObservableHealthCheck.cs                  | 5 +++--
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/Abstractions/Activities/ActivityExtensions.cs b/src/Abstractions/Activities/ActivityExtensions.cs
index 24f33c7a7..26624232f 100644
--- a/src/Abstractions/Activities/ActivityExtensions.cs
+++ b/src/Abstractions/Activities/ActivityExtensions.cs
@@ -4,6 +4,7 @@
 using System;
 using System.Diagnostics;
 using System.Globalization;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
 using Microsoft.Omex.Extensions.Abstractions.Activities.Processing;
 
 namespace Microsoft.Omex.Extensions.Abstractions.Activities
@@ -72,6 +73,13 @@ public static Activity MarkAsHealthCheck(this Activity activity) =>
 		public static Activity SetResult(this Activity activity, ActivityResult result) =>
 			activity.SetTag(ActivityTagKeys.Result, ActivityResultStrings.ResultToString(result));
 
+		/// <summary>
+		/// Set health check result
+		/// </summary>
+		/// <remarks>This property won't be transferred to child activity or via web requests</remarks>
+		public static Activity SetHealthCheckResult(this Activity activity, HealthStatus result) =>
+			activity.SetTag(ActivityTagKeys.HealthCheckResult, result.ToString());
+
 		/// <summary>
 		/// Set activity result to Success
 		/// </summary>
diff --git a/src/Abstractions/Microsoft.Omex.Extensions.Abstractions.csproj b/src/Abstractions/Microsoft.Omex.Extensions.Abstractions.csproj
index fbf0876d6..33a3cb2ee 100644
--- a/src/Abstractions/Microsoft.Omex.Extensions.Abstractions.csproj
+++ b/src/Abstractions/Microsoft.Omex.Extensions.Abstractions.csproj
@@ -14,6 +14,7 @@
     <Content Include="BuildTransitive\*" Pack="true" PackagePath="buildTransitive" />
   </ItemGroup>
   <ItemGroup>
+    <PackageReference Include="Microsoft.Extensions.Diagnostics.HealthChecks" />
     <PackageReference Include="Microsoft.Extensions.Hosting.Abstractions" />
     <PackageReference Include="Microsoft.Extensions.Logging.Abstractions" />
     <PackageReference Include="System.Diagnostics.DiagnosticSource" />
diff --git a/src/Diagnostics.HealthChecks/Composables/ObservableHealthCheck.cs b/src/Diagnostics.HealthChecks/Composables/ObservableHealthCheck.cs
index 56a15ecd1..eb5970623 100644
--- a/src/Diagnostics.HealthChecks/Composables/ObservableHealthCheck.cs
+++ b/src/Diagnostics.HealthChecks/Composables/ObservableHealthCheck.cs
@@ -1,4 +1,5 @@
-// Copyright (C) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
 
 using System;
 using System.Diagnostics;
@@ -66,7 +67,7 @@ public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context
 				? result.Status
 				: context.Registration.FailureStatus;
 
-			activity?.SetTag(ActivityTagKeys.HealthCheckResult, healthCheckStatus.ToString());
+			activity?.SetHealthCheckResult(healthCheckStatus);
 
 			return new HealthCheckResult(
 				healthCheckStatus,