Clean transient entities and preserve state-included metadata

src/Client/Core/Entities/EntityMetadata.cs

@@ -11,12 +11,8 @@ namespace Microsoft.DurableTask.Client.Entities;
 /// Represents entity metadata.
 /// </summary>
 /// <typeparam name="TState">The type of state held by the metadata.</typeparam>
-/// <remarks>
-/// Initializes a new instance of the <see cref="EntityMetadata{TState}"/> class.
-/// </remarks>
-/// <param name="id">The ID of the entity.</param>
 [JsonConverter(typeof(EntityMetadataConverter))]
-public class EntityMetadata<TState>(EntityInstanceId id)
+public class EntityMetadata<TState>
 {
     readonly TState? state;
 
@@ -25,17 +21,39 @@ public class EntityMetadata<TState>(EntityInstanceId id)
     /// </summary>
     /// <param name="id">The ID of the entity.</param>
     /// <param name="state">The state of the entity.</param>
-    public EntityMetadata(EntityInstanceId id, TState? state)
-        : this(id)
+    /// <param name="includesState">
+    /// A value indicating whether the entity state was included in the response, even when the state is <see langword="null"/>.
+    /// </param>
+    public EntityMetadata(EntityInstanceId id, TState? state, bool includesState)
     {
-        this.IncludesState = state is not null;
+        this.Id = Check.NotDefault(id);
+        this.IncludesState = includesState;
         this.state = state;
     }
 
+    /// <summary>
+    /// Initializes a new instance of the <see cref="EntityMetadata{TState}"/> class.
+    /// </summary>
+    /// <param name="id">The ID of the entity.</param>
+    /// <param name="state">The state of the entity.</param>
+    public EntityMetadata(EntityInstanceId id, TState? state)
+        : this(id, state, state is not null)
+    {
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="EntityMetadata{TState}"/> class.
+    /// </summary>
+    /// <param name="id">The ID of the entity.</param>
+    public EntityMetadata(EntityInstanceId id)
+        : this(id, default, false)
+    {
+    }
+
     /// <summary>
     /// Gets the ID for this entity.
     /// </summary>
-    public EntityInstanceId Id { get; } = Check.NotDefault(id);
+    public EntityInstanceId Id { get; }
 
     /// <summary>
     /// Gets the time the entity was last modified.
@@ -60,7 +78,7 @@ public EntityMetadata(EntityInstanceId id, TState? state)
     /// </remarks>
     [MemberNotNullWhen(true, nameof(State))]
     [MemberNotNullWhen(true, nameof(state))]
-    public bool IncludesState { get; } = false;
+    public bool IncludesState { get; }
 
     /// <summary>
     /// Gets the state for this entity.
@@ -90,13 +108,29 @@ public TState State
 /// <summary>
 /// Represents the metadata for a durable entity instance.
 /// </summary>
-/// <remarks>
-/// Initializes a new instance of the <see cref="EntityMetadata"/> class.
-/// </remarks>
-/// <param name="id">The ID of the entity.</param>
-/// <param name="state">The state of this entity.</param>
 [JsonConverter(typeof(EntityMetadataConverter))]
-public sealed class EntityMetadata(EntityInstanceId id, SerializedData? state = null)
-    : EntityMetadata<SerializedData>(id, state)
+public sealed class EntityMetadata : EntityMetadata<SerializedData>
 {
+    /// <summary>
+    /// Initializes a new instance of the <see cref="EntityMetadata"/> class.
+    /// </summary>
+    /// <param name="id">The ID of the entity.</param>
+    /// <param name="state">The state of this entity.</param>
+    public EntityMetadata(EntityInstanceId id, SerializedData? state = null)
+        : base(id, state)
+    {
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="EntityMetadata"/> class.
+    /// </summary>
+    /// <param name="id">The ID of the entity.</param>
+    /// <param name="state">The state of this entity.</param>
+    /// <param name="includesState">
+    /// A value indicating whether the entity state was included in the response, even when the state is <see langword="null"/>.
+    /// </param>
+    public EntityMetadata(EntityInstanceId id, SerializedData? state, bool includesState)
+        : base(id, state, includesState)
+    {
+    }
 }

src/Client/Grpc/GrpcDurableEntityClient.cs

@@ -1,7 +1,9 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Collections.Generic;
 using System.Diagnostics;
+using System.Linq;
 using Microsoft.DurableTask.Client.Entities;
 using Microsoft.DurableTask.Entities;
 using Microsoft.Extensions.Logging;
@@ -126,6 +128,11 @@ public override async Task<CleanEntityStorageResult> CleanEntityStorageAsync(
             }
             while (continueUntilComplete && continuationToken != null);
 
+            if (continueUntilComplete && req.RemoveEmptyEntities)
+            {
+                emptyEntitiesRemoved += await this.PurgeTransientEntitiesAsync(cancellation);
+            }
+
             return new CleanEntityStorageResult
             {
                 ContinuationToken = continuationToken,
@@ -220,12 +227,24 @@ EntityMetadata ToEntityMetadata(P.EntityMetadata metadata, bool includeState)
     {
         var coreEntityId = DTCore.Entities.EntityId.FromString(metadata.InstanceId);
         EntityInstanceId entityId = new(coreEntityId.Name, coreEntityId.Key);
-        bool hasState = metadata.SerializedState != null;
+        bool hasState = !string.IsNullOrEmpty(metadata.SerializedState);
+
+        DateTimeOffset lastModified = metadata.LastModifiedTime?.ToDateTimeOffset() ?? default;
+
+        if (includeState && hasState)
+        {
+            SerializedData data = new(metadata.SerializedState!, this.dataConverter);
+            return new EntityMetadata(entityId, data)
+            {
+                LastModifiedTime = lastModified,
+                BacklogQueueSize = metadata.BacklogQueueSize,
+                LockedBy = metadata.LockedBy,
+            };
+        }
 
-        SerializedData? data = (includeState && hasState) ? new(metadata.SerializedState!, this.dataConverter) : null;
-        return new EntityMetadata(entityId, data)
+        return new EntityMetadata(entityId, null, includesState: includeState)
         {
-            LastModifiedTime = metadata.LastModifiedTime.ToDateTimeOffset(),
+            LastModifiedTime = lastModified,
             BacklogQueueSize = metadata.BacklogQueueSize,
             LockedBy = metadata.LockedBy,
         };
@@ -235,27 +254,91 @@ EntityMetadata<T> ToEntityMetadata<T>(P.EntityMetadata metadata, bool includeSta
     {
         var coreEntityId = DTCore.Entities.EntityId.FromString(metadata.InstanceId);
         EntityInstanceId entityId = new(coreEntityId.Name, coreEntityId.Key);
-        DateTimeOffset lastModified = metadata.LastModifiedTime.ToDateTimeOffset();
-        bool hasState = metadata.SerializedState != null;
+        DateTimeOffset lastModified = metadata.LastModifiedTime?.ToDateTimeOffset() ?? default;
+        bool hasState = !string.IsNullOrEmpty(metadata.SerializedState);
 
         if (includeState && hasState)
         {
-            T? data = includeState ? this.dataConverter.Deserialize<T>(metadata.SerializedState) : default;
+            T? data = this.dataConverter.Deserialize<T>(metadata.SerializedState);
             return new EntityMetadata<T>(entityId, data)
             {
                 LastModifiedTime = lastModified,
                 BacklogQueueSize = metadata.BacklogQueueSize,
                 LockedBy = metadata.LockedBy,
             };
         }
-        else
-        {
-            return new EntityMetadata<T>(entityId)
+
+        return includeState
+            ? new EntityMetadata<T>(entityId, state: default, includesState: true)
+            {
+                LastModifiedTime = lastModified,
+                BacklogQueueSize = metadata.BacklogQueueSize,
+                LockedBy = metadata.LockedBy,
+            }
+            : new EntityMetadata<T>(entityId)
             {
                 LastModifiedTime = lastModified,
                 BacklogQueueSize = metadata.BacklogQueueSize,
                 LockedBy = metadata.LockedBy,
             };
+    }
+
+    /// <summary>
+    /// Iterates transient entity metadata and purges empty entities to fully clean storage.
+    /// </summary>
+    /// <param name="cancellation">The cancellation token.</param>
+    /// <returns>The number of entities removed.</returns>
+    async Task<int> PurgeTransientEntitiesAsync(CancellationToken cancellation)
+    {
+        int removed = 0;
+        string? continuation = null;
+
+        do
+        {
+            P.QueryEntitiesResponse response = await this.sidecarClient.QueryEntitiesAsync(
+                new P.QueryEntitiesRequest
+                {
+                    Query = new P.EntityQuery
+                    {
+                        IncludeState = true,
+                        IncludeTransient = true,
+                        ContinuationToken = continuation,
+                    },
+                },
+                cancellationToken: cancellation);
+
+            List<string> transientIds = response.Entities
+                .Where(IsEmptyTransientEntity)
+                .Select(e => e.InstanceId)
+                .ToList();
+
+            if (transientIds.Count > 0)
+            {
+                P.PurgeInstancesResponse purgeResponse = await this.sidecarClient.PurgeInstancesAsync(
+                    new P.PurgeInstancesRequest
+                    {
+                        InstanceBatch = new P.InstanceBatch
+                        {
+                            InstanceIds = { transientIds },
+                        },
+                    },
+                    cancellationToken: cancellation);
+
+                removed += purgeResponse.DeletedInstanceCount;
+            }
+
+            continuation = response.ContinuationToken;
         }
+        while (continuation != null);
+
+        return removed;
     }
+
+    /// <summary>
+    /// Determines whether an entity is transient and contains no persisted state or locks.
+    /// </summary>
+    /// <param name="entity">The entity metadata.</param>
+    /// <returns><c>true</c> when the entity is empty and transient; otherwise, <c>false</c>.</returns>
+    static bool IsEmptyTransientEntity(P.EntityMetadata entity) =>
+        string.IsNullOrEmpty(entity.SerializedState) && string.IsNullOrEmpty(entity.LockedBy);
 }