Add support for entities to OrchestrationServiceClientShim package

src/Client/OrchestrationServiceClientShim/Client.OrchestrationServiceClientShim.csproj

@@ -8,13 +8,6 @@ The client is responsible for interacting with orchestrations from outside the w
     <EnableStyleCop>true</EnableStyleCop>
   </PropertyGroup>
 
-  <PropertyGroup>
-    <!-- We are still working on this package for entities preview. -->
-    <VersionPrefix>1.0.5</VersionPrefix>
-    <VersionSuffix></VersionSuffix> <!-- Need this here to set it back to empty. -->
-  </PropertyGroup>
-
-
   <ItemGroup>
     <ProjectReference Include="../Core/Client.csproj" />
   </ItemGroup>

src/Client/OrchestrationServiceClientShim/DependencyInjection/DurableTaskClientBuilderExtensions.cs

@@ -2,8 +2,10 @@
 // Licensed under the MIT License.
 
 using DurableTask.Core;
+using DurableTask.Core.Entities;
 using Microsoft.DurableTask.Client.OrchestrationServiceClientShim;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Options;
 
 namespace Microsoft.DurableTask.Client;
 
@@ -58,17 +60,51 @@ public static IDurableTaskClientBuilder UseOrchestrationService(
         builder.Services.AddOptions<ShimDurableTaskClientOptions>(builder.Name)
             .PostConfigure<IServiceProvider>((opt, sp) =>
             {
-                if (opt.Client is not null)
-                {
-                    return;
-                }
-
-                // Try to resolve client from service container.
-                opt.Client = sp.GetService<IOrchestrationServiceClient>()
-                    ?? sp.GetService<IOrchestrationService>() as IOrchestrationServiceClient;
+                ConfigureClient(sp, opt);
+                ConfigureEntities(builder.Name, sp, opt);
             })
-            .Validate(x => x.Client is not null, "ShimDurableTaskClientOptions.Client must not be null.");
+            .Validate(x => x.Client is not null, "ShimDurableTaskClientOptions.Client must not be null.")
+            .Validate(
+                x => !x.EnableEntitySupport || x.Entities.Queries is not null,
+                "ShimDurableTaskClientOptions.Entities.Queries must not be null when entity support is enabled.");
 
         return builder.UseBuildTarget<ShimDurableTaskClient, ShimDurableTaskClientOptions>();
     }
+
+    static void ConfigureClient(IServiceProvider services, ShimDurableTaskClientOptions options)
+    {
+        if (options.Client is not null)
+        {
+            return;
+        }
+
+        // Try to resolve client from service container.
+        options.Client = services.GetService<IOrchestrationServiceClient>()
+            ?? services.GetService<IOrchestrationService>() as IOrchestrationServiceClient;
+    }
+
+    static void ConfigureEntities(string name, IServiceProvider services, ShimDurableTaskClientOptions options)
+    {
+        if (options.Entities.Queries is null)
+        {
+            options.Entities.Queries = services.GetService<EntityBackendQueries>()
+                ?? GetEntityService(services, options)?.EntityBackendQueries;
+        }
+
+        if (options.Entities.MaxSignalDelayTime is null)
+        {
+            EntityBackendProperties? properties = services.GetService<IOptionsMonitor<EntityBackendProperties>>()?.Get(name)
+                ?? GetEntityService(services, options)?.EntityBackendProperties;
+            options.Entities.MaxSignalDelayTime = properties?.MaximumSignalDelayTime;
+        }
+    }
+
+    static IEntityOrchestrationService? GetEntityService(
+        IServiceProvider services, ShimDurableTaskClientOptions options)
+    {
+        return services.GetService<IEntityOrchestrationService>()
+            ?? services.GetService<IOrchestrationService>() as IEntityOrchestrationService
+            ?? services.GetService<IOrchestrationServiceClient>() as IEntityOrchestrationService
+            ?? options.Client as IEntityOrchestrationService;
+    }
 }

src/Client/OrchestrationServiceClientShim/ShimDurableEntityClient.cs

@@ -0,0 +1,177 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using DurableTask.Core;
+using DurableTask.Core.Entities;
+using Microsoft.DurableTask.Client.Entities;
+using Microsoft.DurableTask.Entities;
+
+namespace Microsoft.DurableTask.Client.OrchestrationServiceClientShim;
+
+/// <summary>
+/// A shim client for interacting with entities backend via <see cref="IOrchestrationServiceClient"/>.
+/// </summary>
+class ShimDurableEntityClient : DurableEntityClient
+{
+    readonly ShimDurableTaskClientOptions options;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ShimDurableEntityClient"/> class.
+    /// </summary>
+    /// <param name="name">The name of this client.</param>
+    /// <param name="options">The client options..</param>
+    public ShimDurableEntityClient(string name, ShimDurableTaskClientOptions options)
+        : base(name)
+    {
+        this.options = Check.NotNull(options);
+    }
+
+    EntityBackendQueries Queries => this.options.Entities.Queries!;
+
+    DataConverter Converter => this.options.DataConverter;
+
+    /// <inheritdoc/>
+    public override async Task<CleanEntityStorageResult> CleanEntityStorageAsync(
+        CleanEntityStorageRequest? request = null,
+        bool continueUntilComplete = true,
+        CancellationToken cancellation = default)
+    {
+        CleanEntityStorageRequest r = request ?? CleanEntityStorageRequest.Default;
+        EntityBackendQueries.CleanEntityStorageResult result = await this.Queries.CleanEntityStorageAsync(
+            new EntityBackendQueries.CleanEntityStorageRequest()
+            {
+                RemoveEmptyEntities = r.RemoveEmptyEntities,
+                ReleaseOrphanedLocks = r.ReleaseOrphanedLocks,
+                ContinuationToken = r.ContinuationToken,
+            },
+            cancellation);
+
+        return new()
+        {
+            EmptyEntitiesRemoved = result.EmptyEntitiesRemoved,
+            OrphanedLocksReleased = result.OrphanedLocksReleased,
+            ContinuationToken = result.ContinuationToken,
+        };
+    }
+
+    /// <inheritdoc/>
+    public override AsyncPageable<EntityMetadata> GetAllEntitiesAsync(EntityQuery? filter = null)
+        => this.GetAllEntitiesAsync(this.Convert, filter);
+
+    /// <inheritdoc/>
+    public override AsyncPageable<EntityMetadata<T>> GetAllEntitiesAsync<T>(EntityQuery? filter = null)
+        => this.GetAllEntitiesAsync(this.Convert<T>, filter);
+
+    /// <inheritdoc/>
+    public override async Task<EntityMetadata?> GetEntityAsync(
+        EntityInstanceId id, bool includeState = true, CancellationToken cancellation = default)
+        => this.Convert(await this.Queries.GetEntityAsync(
+            new EntityId(id.Name, id.Key), includeState, false, cancellation));
+
+    /// <inheritdoc/>
+    public override async Task<EntityMetadata<T>?> GetEntityAsync<T>(
+        EntityInstanceId id, bool includeState = true, CancellationToken cancellation = default)
+        => this.Convert<T>(await this.Queries.GetEntityAsync(
+            new EntityId(id.Name, id.Key), includeState, false, cancellation));
+
+    /// <inheritdoc/>
+    public override async Task SignalEntityAsync(
+        EntityInstanceId id,
+        string operationName,
+        object? input = null,
+        SignalEntityOptions? options = null,
+        CancellationToken cancellation = default)
+    {
+        Check.NotNullOrEmpty(id.Name);
+        Check.NotNull(id.Key);
+
+        DateTimeOffset? scheduledTime = options?.SignalTime;
+        string? serializedInput = this.Converter.Serialize(input);
+
+        EntityMessageEvent eventToSend = ClientEntityHelpers.EmitOperationSignal(
+            new OrchestrationInstance() { InstanceId = id.ToString() },
+            Guid.NewGuid(),
+            operationName,
+            serializedInput,
+            EntityMessageEvent.GetCappedScheduledTime(
+                DateTime.UtcNow,
+                this.options.Entities.MaxSignalDelayTimeOrDefault,
+                scheduledTime?.UtcDateTime));
+
+        await this.options.Client!.SendTaskOrchestrationMessageAsync(eventToSend.AsTaskMessage());
+    }
+
+    AsyncPageable<TMetadata> GetAllEntitiesAsync<TMetadata>(
+        Func<EntityBackendQueries.EntityMetadata, TMetadata> select,
+        EntityQuery? filter)
+        where TMetadata : notnull
+    {
+        bool includeState = filter?.IncludeState ?? true;
+        bool includeTransient = filter?.IncludeTransient ?? false;
+        string startsWith = filter?.InstanceIdStartsWith ?? string.Empty;
+        DateTime? lastModifiedFrom = filter?.LastModifiedFrom?.UtcDateTime;
+        DateTime? lastModifiedTo = filter?.LastModifiedTo?.UtcDateTime;
+
+        return Pageable.Create(async (continuation, size, cancellation) =>
+        {
+            size ??= filter?.PageSize;
+            EntityBackendQueries.EntityQueryResult result = await this.Queries.QueryEntitiesAsync(
+                new EntityBackendQueries.EntityQuery()
+                {
+                    InstanceIdStartsWith = startsWith,
+                    LastModifiedFrom = lastModifiedFrom,
+                    LastModifiedTo = lastModifiedTo,
+                    IncludeTransient = includeTransient,
+                    IncludeState = includeState,
+                    ContinuationToken = continuation,
+                    PageSize = size,
+                },
+                cancellation);
+
+            return new Page<TMetadata>(result.Results.Select(select).ToList(), result.ContinuationToken);
+        });
+    }
+
+    EntityMetadata<T> Convert<T>(EntityBackendQueries.EntityMetadata metadata)
+    {
+        return new(
+            new EntityInstanceId(metadata.EntityId.Name, metadata.EntityId.Key),
+            this.Converter.Deserialize<T>(metadata.SerializedState))
+            {
+                LastModifiedTime = metadata.LastModifiedTime,
+                BacklogQueueSize = metadata.BacklogQueueSize,
+                LockedBy = metadata.LockedBy,
+            };
+    }
+
+    EntityMetadata<T>? Convert<T>(EntityBackendQueries.EntityMetadata? metadata)
+    {
+        if (metadata is null)
+        {
+            return null;
+        }
+
+        return this.Convert<T>(metadata.Value);
+    }
+
+    EntityMetadata Convert(EntityBackendQueries.EntityMetadata metadata)
+    {
+        SerializedData? data = metadata.SerializedState is null ? null : new(metadata.SerializedState, this.Converter);
+        return new(new EntityInstanceId(metadata.EntityId.Name, metadata.EntityId.Key), data)
+        {
+            LastModifiedTime = metadata.LastModifiedTime,
+            BacklogQueueSize = metadata.BacklogQueueSize,
+            LockedBy = metadata.LockedBy,
+        };
+    }
+
+    EntityMetadata? Convert(EntityBackendQueries.EntityMetadata? metadata)
+    {
+        if (metadata is null)
+        {
+            return null;
+        }
+
+        return this.Convert(metadata.Value);
+    }
+}

src/Client/OrchestrationServiceClientShim/ShimDurableTaskClient.cs

@@ -3,8 +3,10 @@
 
 using System.Diagnostics.CodeAnalysis;
 using DurableTask.Core;
+using DurableTask.Core.Entities;
 using DurableTask.Core.History;
 using DurableTask.Core.Query;
+using Microsoft.DurableTask.Client.Entities;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 using Core = DurableTask.Core;
@@ -18,6 +20,7 @@ namespace Microsoft.DurableTask.Client.OrchestrationServiceClientShim;
 class ShimDurableTaskClient : DurableTaskClient
 {
     readonly ShimDurableTaskClientOptions options;
+    ShimDurableEntityClient? entities;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="ShimDurableTaskClient"/> class.
@@ -42,6 +45,31 @@ public ShimDurableTaskClient(string name, ShimDurableTaskClientOptions options)
         this.options = Check.NotNull(options);
     }
 
+    /// <inheritdoc/>
+    public override DurableEntityClient Entities
+    {
+        get
+        {
+            if (!this.options.EnableEntitySupport)
+            {
+                throw new InvalidOperationException("Entity support is not enabled.");
+            }
+
+            if (this.entities is null)
+            {
+                if (this.options.Entities.Queries is null)
+                {
+                    throw new NotSupportedException(
+                        "The configured IOrchestrationServiceClient does not support entities.");
+                }
+
+                this.entities = new(this.Name, this.options);
+            }
+
+            return this.entities;
+        }
+    }
+
     DataConverter DataConverter => this.options.DataConverter;
 
     IOrchestrationServiceClient Client => this.options.Client!;

src/Client/OrchestrationServiceClientShim/ShimDurableTaskClientOptions.cs

@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 
 using DurableTask.Core;
+using DurableTask.Core.Entities;
+using Microsoft.DurableTask.Client.Entities;
 
 namespace Microsoft.DurableTask.Client.OrchestrationServiceClientShim;
 
@@ -15,4 +17,33 @@ public sealed class ShimDurableTaskClientOptions : DurableTaskClientOptions
     /// If not manually set, this will be resolved from the <see cref="IServiceProvider" />, if available.
     /// </summary>
     public IOrchestrationServiceClient? Client { get; set; }
+
+    /// <summary>
+    /// Gets the <see cref="ShimDurableTaskEntityOptions"/> to configure entity support.
+    /// </summary>
+    public ShimDurableTaskEntityOptions Entities { get; } = new();
+}
+
+/// <summary>
+/// Options for entities.
+/// </summary>
+public class ShimDurableTaskEntityOptions
+{
+    /// <summary>
+    /// Gets or sets the <see cref="EntityBackendQueries"/> to use in the <see cref="DurableEntityClient" />.
+    /// If not set manually, this will attempt to be resolved automatically by looking for
+    /// <see cref="IEntityOrchestrationService"/> in the <see cref="IServiceProvider"/>.
+    /// </summary>
+    public EntityBackendQueries? Queries { get; set; }
+
+    /// <summary>
+    /// Gets or sets the maximum time span to use for signal delay. If not set manually, will attempt to be resolved
+    /// through the service container. This will finally default to 3 days if it cannot be any other means.
+    /// </summary>
+    public TimeSpan? MaxSignalDelayTime { get; set; }
+
+    /// <summary>
+    /// Gets the max signal delay time.
+    /// </summary>
+    internal TimeSpan MaxSignalDelayTimeOrDefault => this.MaxSignalDelayTime ?? TimeSpan.FromDays(3);
 }