From 79c2d005f3e592a2c33ebb9dcf2f1fc203c2b09f Mon Sep 17 00:00:00 2001
From: Danny Guinther <danny@phave.com>
Date: Mon, 9 Dec 2024 11:16:05 -0500
Subject: [PATCH] Evolve to only make IServiceProvider accessible to
 InterceptActivity

---
 .../IServiceProviderAccessor.cs               | 16 -------
 .../ServiceProviderAccessor.cs                | 42 -------------------
 .../ServiceProviderExtensions.cs              | 13 +-----
 .../ActivityServiceProviderAccessor.cs        | 31 ++++++++++++++
 .../IWorkerInterceptorWithServiceProvider.cs  | 31 ++++++++++++++
 5 files changed, 64 insertions(+), 69 deletions(-)
 delete mode 100644 src/Temporalio.Extensions.Hosting/IServiceProviderAccessor.cs
 delete mode 100644 src/Temporalio.Extensions.Hosting/ServiceProviderAccessor.cs
 create mode 100644 src/Temporalio/Activities/ActivityServiceProviderAccessor.cs
 create mode 100644 src/Temporalio/Worker/Interceptors/IWorkerInterceptorWithServiceProvider.cs

diff --git a/src/Temporalio.Extensions.Hosting/IServiceProviderAccessor.cs b/src/Temporalio.Extensions.Hosting/IServiceProviderAccessor.cs
deleted file mode 100644
index ff0639c6..00000000
--- a/src/Temporalio.Extensions.Hosting/IServiceProviderAccessor.cs
+++ /dev/null
@@ -1,16 +0,0 @@
-using System;
-
-namespace Temporalio.Extensions.Hosting
-{
-    /// <summary>
-    /// Provides access to the current, scoped <see cref="IServiceProvider"/> if
-    /// one is available.
-    /// </summary>
-    public interface IServiceProviderAccessor
-    {
-        /// <summary>
-        /// Gets or sets the current service provider.
-        /// </summary>
-        IServiceProvider? ServiceProvider { get; set; }
-    }
-}
\ No newline at end of file
diff --git a/src/Temporalio.Extensions.Hosting/ServiceProviderAccessor.cs b/src/Temporalio.Extensions.Hosting/ServiceProviderAccessor.cs
deleted file mode 100644
index ad3f9802..00000000
--- a/src/Temporalio.Extensions.Hosting/ServiceProviderAccessor.cs
+++ /dev/null
@@ -1,42 +0,0 @@
-using System;
-using System.Threading;
-
-namespace Temporalio.Extensions.Hosting
-{
-    /// <summary>
-    /// Provides an implementation of <see cref="IServiceProvider" /> based on
-    /// the current execution context.
-    /// </summary>
-    public class ServiceProviderAccessor : IServiceProviderAccessor
-    {
-        private static readonly AsyncLocal<ServiceProviderHolder> ServiceProviderCurrent = new();
-
-        /// <inheritdoc/>
-        public IServiceProvider? ServiceProvider
-        {
-            get => ServiceProviderCurrent.Value?.ServiceProvider;
-
-            set
-            {
-                var holder = ServiceProviderCurrent.Value;
-                if (holder != null)
-                {
-                    // Clear current IServiceProvider trapped in the AsyncLocals, as its done.
-                    holder.ServiceProvider = null;
-                }
-
-                if (value != null)
-                {
-                    // Use an object indirection to hold the IServiceProvider in the AsyncLocal,
-                    // so it can be cleared in all ExecutionContexts when its cleared.
-                    ServiceProviderCurrent.Value = new ServiceProviderHolder { ServiceProvider = value };
-                }
-            }
-        }
-
-        private sealed class ServiceProviderHolder
-        {
-            public IServiceProvider? ServiceProvider { get; set; }
-        }
-    }
-}
\ No newline at end of file
diff --git a/src/Temporalio.Extensions.Hosting/ServiceProviderExtensions.cs b/src/Temporalio.Extensions.Hosting/ServiceProviderExtensions.cs
index c6c0e113..54620362 100644
--- a/src/Temporalio.Extensions.Hosting/ServiceProviderExtensions.cs
+++ b/src/Temporalio.Extensions.Hosting/ServiceProviderExtensions.cs
@@ -68,13 +68,7 @@ public static ActivityDefinition CreateTemporalActivityDefinition(
 #else
                 var scope = provider.CreateScope();
 #endif
-                IServiceProviderAccessor? serviceProviderAccessor =
-                    scope.ServiceProvider.GetService<IServiceProviderAccessor>();
-
-                if (serviceProviderAccessor is not null)
-                {
-                    serviceProviderAccessor.ServiceProvider = scope.ServiceProvider;
-                }
+                ActivityServiceProviderAccessor.AsyncLocalCurrent.Value = scope.ServiceProvider;
 
                 try
                 {
@@ -119,10 +113,7 @@ public static ActivityDefinition CreateTemporalActivityDefinition(
                 }
                 finally
                 {
-                    if (serviceProviderAccessor is not null)
-                    {
-                        serviceProviderAccessor.ServiceProvider = null;
-                    }
+                    ActivityServiceProviderAccessor.AsyncLocalCurrent.Value = null;
 #if NET6_0_OR_GREATER
                     await scope.DisposeAsync().ConfigureAwait(false);
 #else
diff --git a/src/Temporalio/Activities/ActivityServiceProviderAccessor.cs b/src/Temporalio/Activities/ActivityServiceProviderAccessor.cs
new file mode 100644
index 00000000..becdf2db
--- /dev/null
+++ b/src/Temporalio/Activities/ActivityServiceProviderAccessor.cs
@@ -0,0 +1,31 @@
+using System;
+using System.Threading;
+
+namespace Temporalio.Activities
+{
+    /// <summary>
+    /// Provides access to the scoped <see cref="IServiceProvider" /> for the
+    /// current activity execution context.
+    /// </summary>
+    public static class ActivityServiceProviderAccessor
+    {
+        /// <summary>
+        /// Gets the async local current value.
+        /// </summary>
+        public static readonly AsyncLocal<IServiceProvider?> AsyncLocalCurrent = new();
+
+        /// <summary>
+        /// Gets a value indicating whether the current code is running in an
+        /// activity with a service provider.
+        /// </summary>
+        public static bool HasCurrent => AsyncLocalCurrent.Value is not null;
+
+        /// <summary>
+        /// Gets the current activity's scoped <see cref="IServiceProvider"/>.
+        /// </summary>
+        /// <exception cref="InvalidOperationException">If no <see
+        /// cref="IServiceProvider"/>  is available.</exception>
+        public static IServiceProvider Current => AsyncLocalCurrent.Value ??
+            throw new InvalidOperationException("No current service provider");
+    }
+}
diff --git a/src/Temporalio/Worker/Interceptors/IWorkerInterceptorWithServiceProvider.cs b/src/Temporalio/Worker/Interceptors/IWorkerInterceptorWithServiceProvider.cs
new file mode 100644
index 00000000..d72bf949
--- /dev/null
+++ b/src/Temporalio/Worker/Interceptors/IWorkerInterceptorWithServiceProvider.cs
@@ -0,0 +1,31 @@
+using System;
+#if NETCOREAPP3_0_OR_GREATER
+using Temporalio.Activities;
+#endif
+
+namespace Temporalio.Worker.Interceptors
+{
+    /// <summary>
+    /// Interceptor for intercepting activities and workflows providing activity
+    /// interceptors with access to the service provider.
+    /// </summary>
+    public interface IWorkerInterceptorWithServiceProvider : IWorkerInterceptor
+    {
+        /// <summary>
+        /// Create an activity inbound interceptor to intercept calls, with
+        /// access to the service provider for the current activity.
+        /// </summary>
+        /// <param name="serviceProvider">Service provider instance scoped to the current activity.</param>
+        /// <param name="nextInterceptor">The next interceptor in the chain to call.</param>
+        /// <returns>Created interceptor.</returns>
+#if NETCOREAPP3_0_OR_GREATER
+        ActivityInboundInterceptor InterceptActivity(IServiceProvider serviceProvider, ActivityInboundInterceptor nextInterceptor) =>
+            nextInterceptor;
+
+        ActivityInboundInterceptor InterceptActivity(ActivityInboundInterceptor nextInterceptor) =>
+            InterceptActivity(ActivityServiceProviderAccessor.Current, nextInterceptor);
+#else
+        ActivityInboundInterceptor InterceptActivity(IServiceProvider serviceProvider, ActivityInboundInterceptor nextInterceptor);
+#endif
+    }
+}
\ No newline at end of file