Durable Jobs (aka Reminders v2)

Orleans.slnx

@@ -34,6 +34,7 @@
     <Project Path="src/Orleans.Persistence.Memory/Orleans.Persistence.Memory.csproj" />
     <Project Path="src/Orleans.Reminders/Orleans.Reminders.csproj" />
     <Project Path="src/Orleans.Runtime/Orleans.Runtime.csproj" />
+    <Project Path="src/Orleans.ScheduledJobs/Orleans.ScheduledJobs.csproj" Id="ab4cc723-b3cf-4c32-be4c-1a180a49797c" />
     <Project Path="src/Orleans.Sdk/Orleans.Sdk.csproj" />
     <Project Path="src/Orleans.Serialization.Abstractions/Orleans.Serialization.Abstractions.csproj" />
     <Project Path="src/Orleans.Serialization.FSharp/Orleans.Serialization.FSharp.csproj" />
@@ -76,6 +77,7 @@
     <Project Path="src/Azure/Orleans.Persistence.Cosmos/Orleans.Persistence.Cosmos.csproj" />
     <Project Path="src/Azure/Orleans.Reminders.AzureStorage/Orleans.Reminders.AzureStorage.csproj" />
     <Project Path="src/Azure/Orleans.Reminders.Cosmos/Orleans.Reminders.Cosmos.csproj" />
+    <Project Path="src/Azure/Orleans.ScheduledJobs.AzureStorage/Orleans.ScheduledJobs.AzureStorage.csproj" Id="5bbb1058-9ab2-44e9-b237-49f5b6bee151" />
     <Project Path="src/Azure/Orleans.Streaming.AzureStorage/Orleans.Streaming.AzureStorage.csproj" />
     <Project Path="src/Azure/Orleans.Streaming.EventHubs/Orleans.Streaming.EventHubs.csproj" />
     <Project Path="src/Azure/Orleans.Transactions.AzureStorage/Orleans.Transactions.AzureStorage.csproj" />
@@ -84,7 +86,7 @@
     <Project Path="src/Cassandra/Orleans.Clustering.Cassandra/Orleans.Clustering.Cassandra.csproj" />
   </Folder>
   <Folder Name="/src/Extensions/NATS/">
-    <Project Path="src\Orleans.Streaming.NATS\Orleans.Streaming.NATS.csproj" Type="Classic C#" />
+    <Project Path="src\Orleans.Streaming.NATS\Orleans.Streaming.NATS.csproj" />
   </Folder>
   <Folder Name="/src/Extensions/Redis/">
     <Project Path="src/Redis/Orleans.Clustering.Redis/Orleans.Clustering.Redis.csproj" />
@@ -133,7 +135,7 @@
     <Project Path="test/Extensions/TesterAdoNet/Tester.AdoNet.csproj" />
     <Project Path="test/Extensions/TesterAzureUtils/Tester.AzureUtils.csproj" />
     <Project Path="test/Extensions/TesterZooKeeperUtils/Tester.ZooKeeperUtils.csproj" />
-    <Project Path="test\Extensions\NATS.Tests\NATS.Tests.csproj" Type="Classic C#" />
+    <Project Path="test\Extensions\NATS.Tests\NATS.Tests.csproj" />
   </Folder>
   <Folder Name="/test/Grains/">
     <Project Path="test/Grains/TestFSharp/TestFSharp.fsproj" />

src/Azure/Orleans.ScheduledJobs.AzureStorage/AzureStorageJobShard.cs

@@ -0,0 +1,149 @@
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Text.Json;
+using System.Threading.Tasks;
+using Azure;
+using Azure.Storage.Blobs;
+using Azure.Storage.Blobs.Models;
+using Azure.Storage.Blobs.Specialized;
+using Orleans.Runtime;
+
+namespace Orleans.ScheduledJobs.AzureStorage;
+
+internal sealed class AzureStorageJobShard : JobShard
+{
+    internal AppendBlobClient BlobClient { get; init; }
+    internal ETag? ETag { get; private set; }
+
+    private InMemoryJobQueue _jobQueue;
+    private int _jobCount = 0;
+
+    public AzureStorageJobShard(string id, DateTimeOffset startTime, DateTimeOffset endTime, AppendBlobClient blobClient)
+        : base(id, startTime, endTime)
+    {
+        BlobClient = blobClient;
+        _jobQueue = new InMemoryJobQueue();
+    }
+
+    public override ValueTask<int> GetJobCount()
+    {
+        return ValueTask.FromResult(_jobCount);
+    }
+
+    public override IAsyncEnumerable<IScheduledJob> ConsumeScheduledJobsAsync()
+    {
+        return _jobQueue;
+    }
+
+    public override async Task RemoveJobAsync(string jobId)
+    {
+        var operation = JobOperation.CreateRemoveOperation(jobId);
+        await AppendOperation(operation);
+        _jobQueue.CancelJob(jobId);
+        _jobCount--;
+    }
+
+    public override async Task<IScheduledJob> ScheduleJobAsync(GrainId target, string jobName, DateTimeOffset dueTime)
+    {
+        if (IsComplete)
+            throw new InvalidOperationException("Cannot schedule job on a complete shard.");
+
+        var jobId = Guid.NewGuid().ToString();
+        var operation = JobOperation.CreateAddOperation(jobId, jobName, dueTime, target);
+        await AppendOperation(operation);
+        _jobCount++;
+        var job = new ScheduledJob
+        {
+            Id = jobId,
+            Name = jobName,
+            DueTime = dueTime,
+            TargetGrainId = target,
+            ShardId = Id
+        };
+        _jobQueue.Enqueue(job);
+        return job;
+    }
+
+    private async Task AppendOperation(JobOperation operation)
+    {
+        var content = BinaryData.FromObjectAsJson(operation).ToString() + Environment.NewLine;
+        using var stream = new MemoryStream(System.Text.Encoding.UTF8.GetBytes(content));
+        var result = await BlobClient.AppendBlockAsync(
+                    stream,
+                    new AppendBlobAppendBlockOptions { Conditions = new AppendBlobRequestConditions { IfMatch = ETag } });
+        ETag = result.Value.ETag;
+    }
+
+    public async ValueTask InitializeAsync()
+    {
+        if (ETag is not null) return; // already initialized
+
+        // Load existing blob
+        var response = await BlobClient.DownloadAsync();
+        using var stream = response.Value.Content;
+        using var reader = new StreamReader(stream);
+
+        // Rebuild state by replaying operations
+        var dictionary = new Dictionary<string, JobOperation>();
+        while (!reader.EndOfStream)
+        {
+            var line = await reader.ReadLineAsync();
+            if (string.IsNullOrWhiteSpace(line)) continue;
+            var operation = JsonSerializer.Deserialize<JobOperation>(line);
+            switch (operation.Type)
+            {
+                case JobOperation.OperationType.Add:
+                    dictionary[operation.Id] = operation;
+                    break;
+                case JobOperation.OperationType.Remove:
+                    dictionary.Remove(operation.Id);
+                    break;
+            }
+        }
+        // Rebuild the priority queue
+        foreach (var op in dictionary.Values)
+        {
+            _jobQueue.Enqueue(new ScheduledJob
+            {
+                Id = op.Id,
+                Name = op.Name!,
+                DueTime = op.DueTime!.Value,
+                TargetGrainId = op.TargetGrainId!.Value,
+                ShardId = Id
+            });
+        }
+        _jobCount = dictionary.Count;
+
+        ETag = response.Value.Details.ETag;
+    }
+
+    public override Task MarkAsComplete()
+    {
+        IsComplete = true;
+        _jobQueue.MarkAsComplete();
+        return Task.CompletedTask;
+    }
+}
+
+internal struct JobOperation
+{
+    public enum OperationType
+    {
+        Add,
+        Remove
+    }
+
+    public OperationType Type { get; init; }
+
+    public string Id { get; init; }
+    public string? Name { get; init; }
+    public DateTimeOffset? DueTime { get; init; }
+    public GrainId? TargetGrainId { get; init; }
+
+    public static JobOperation CreateAddOperation(string id, string name, DateTimeOffset dueTime, GrainId targetGrainId) =>
+        new() { Type = OperationType.Add, Id = id, Name = name, DueTime = dueTime, TargetGrainId = targetGrainId };
+
+    public static JobOperation CreateRemoveOperation(string id) =>
+        new() { Type = OperationType.Remove, Id = id };
+}

src/Azure/Orleans.ScheduledJobs.AzureStorage/AzureStorageJobShardManager.cs

@@ -0,0 +1,154 @@
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Azure;
+using Azure.Storage.Blobs;
+using Azure.Storage.Blobs.Models;
+using Azure.Storage.Blobs.Specialized;
+using Microsoft.Extensions.Options;
+using Orleans.Runtime;
+
+namespace Orleans.ScheduledJobs.AzureStorage;
+
+public sealed class AzureStorageJobShardManager : JobShardManager
+{
+    private readonly BlobServiceClient _blobServiceClient;
+    private readonly string _containerName;
+    private BlobContainerClient _client = null!;
+    private readonly TimeSpan _maxShardDuration;
+    private readonly IClusterMembershipService _clusterMembership;
+
+    public AzureStorageJobShardManager(BlobServiceClient client, string containerName, TimeSpan maxShardDuration, IClusterMembershipService clusterMembership)
+    {
+        _blobServiceClient = client;
+        _containerName = containerName;
+        _maxShardDuration = maxShardDuration;
+        _clusterMembership = clusterMembership;
+    }
+
+    public AzureStorageJobShardManager(IOptions<AzureStorageJobShardOptions> options, IClusterMembershipService clusterMembership)
+        : this(options.Value.BlobServiceClient, options.Value.ContainerName, options.Value.MaxShardDuration, clusterMembership)
+    {
+    }
+
+    public override async Task<List<JobShard>> GetJobShardsAsync(SiloAddress siloAddress, DateTimeOffset maxDateTime)
+    {
+        await InitializeIfNeeded();
+        var blobs = _client.GetBlobsAsync(traits: BlobTraits.Metadata);
+        var result = new List<JobShard>();
+        await foreach (var blob in blobs)
+        {
+            // Get the owner of the shard
+            var (owner, minDueTime, maxDueTime) = ParseMetadata(blob.Metadata);
+
+            if (minDueTime > maxDateTime)
+            {
+                // This shard is too new, stop there
+                break;
+            }
+
+            if (owner == null || _clusterMembership.CurrentSnapshot.GetSiloStatus(owner) == SiloStatus.Dead)
+            {
+                // The owner is dead or unknown, we can take over this shard
+                var blobClient = _client.GetAppendBlobClient(blob.Name);
+                var metadata = blob.Metadata;
+                metadata["Owner"] = siloAddress.ToParsableString();
+                try
+                {
+                    await blobClient.SetMetadataAsync(metadata, conditions: new BlobRequestConditions { IfMatch = blob.Properties.ETag });
+                }
+                catch (RequestFailedException)
+                {
+                    // Someone else took over the shard
+                    continue;
+                }
+                var shard = new AzureStorageJobShard(blob.Name, minDueTime, maxDueTime, blobClient);
+                await shard.InitializeAsync();
+                await shard.MarkAsComplete();
+                result.Add(shard);
+            }
+        }
+        return result;
+    }
+
+    public override async Task<JobShard> RegisterShard(SiloAddress siloAddress, DateTimeOffset minDueTime, DateTimeOffset maxDueTime, IDictionary<string, string> metadata)
+    {
+        await InitializeIfNeeded();
+        for (var i = 0;; i++) // TODO limit the number of attempts
+        {
+            var shardId = $"{minDueTime:yyyyMMddHHmm}-{siloAddress.ToParsableString()}-{i}";
+            var blobClient = _client.GetAppendBlobClient(shardId);
+            var metadataInfo = CreateMetadata(siloAddress, minDueTime, maxDueTime);
+            try
+            {
+                var response = await blobClient.CreateIfNotExistsAsync(metadata: metadataInfo);
+                if (response == null)
+                {
+                    // Blob already exists, try again with a different name
+                    continue;
+                }
+            }
+            catch (RequestFailedException)
+            {
+                if (i > 100) throw; // Prevent infinite loop
+                // Blob already exists, try again with a different name
+                continue;
+            }
+            var shard = new AzureStorageJobShard(shardId, minDueTime, maxDueTime, blobClient);
+            await shard.InitializeAsync();
+            return shard;
+        }
+    }
+
+    public override async Task UnregisterShard(SiloAddress siloAddress, JobShard shard)
+    {
+        var azureShard = shard as AzureStorageJobShard ?? throw new ArgumentException("Shard is not an AzureStorageJobShard", nameof(shard));
+        var conditions = new BlobRequestConditions { IfMatch = azureShard.ETag };
+        var count = await shard.GetJobCount();
+        var properties = await azureShard.BlobClient.GetPropertiesAsync(conditions);
+        if (count > 0)
+        {
+            // There are still jobs in the shard, unregister it
+            var metadata = properties.Value.Metadata;
+            var (owner, _, _) = ParseMetadata(metadata);
+
+            if (owner != siloAddress)
+                throw new InvalidOperationException("Cannot unregister a shard owned by another silo");
+
+            metadata.Remove("Owner");
+            var response = await azureShard.BlobClient.SetMetadataAsync(metadata, conditions);
+        }
+        else
+        {
+            // No jobs left, we can delete the shard
+            await azureShard.BlobClient.DeleteIfExistsAsync(conditions: conditions);
+        }
+    }
+
+    private ValueTask InitializeIfNeeded()
+    {
+        if (_client != null) return ValueTask.CompletedTask;
+
+        _client = _blobServiceClient.GetBlobContainerClient(_containerName);
+        _client.CreateIfNotExists();
+        return ValueTask.CompletedTask;
+    }
+
+    private static Dictionary<string, string> CreateMetadata(SiloAddress siloAddress, DateTimeOffset minDueTime, DateTimeOffset maxDueTime)
+    {
+        return new Dictionary<string, string>
+        {
+            { "Owner", siloAddress.ToParsableString() },
+            { "MinDueTime", minDueTime.ToString("o") },
+            { "MaxDueTime", maxDueTime.ToString("o") }
+        };
+    }
+
+    private static (SiloAddress? owner, DateTime minDueTime, DateTime maxDueTime) ParseMetadata(IDictionary<string, string> metadata)
+    {
+        var owner = metadata.TryGetValue("Owner", out var ownerStr) ? SiloAddress.FromParsableString(ownerStr) : null;
+        var minDueTime = metadata.TryGetValue("MinDueTime", out var minDueTimeStr) && DateTime.TryParse(minDueTimeStr, out var minDt) ? minDt : DateTime.MinValue;
+        var maxDueTime = metadata.TryGetValue("MaxDueTime", out var maxDueTimeStr) && DateTime.TryParse(maxDueTimeStr, out var maxDt) ? maxDt : DateTime.MaxValue;
+        return (owner, minDueTime.ToUniversalTime(), maxDueTime.ToUniversalTime());
+    }
+}

src/Azure/Orleans.ScheduledJobs.AzureStorage/AzureStorageJobShardOptions.cs

@@ -0,0 +1,22 @@
+using System;
+using Azure.Storage.Blobs;
+
+namespace Orleans.ScheduledJobs.AzureStorage;
+
+public class AzureStorageJobShardOptions
+{
+    /// <summary>
+    /// The maximum duration of a job shard.
+    /// </summary>
+    public TimeSpan MaxShardDuration { get; set; } = TimeSpan.FromHours(1);
+
+    /// <summary>
+    /// Gets or sets the <see cref="BlobServiceClient"/> instance used to store job shards.
+    /// </summary>
+    public BlobServiceClient BlobServiceClient { get; set; } = null!;
+
+    /// <summary>
+    /// Gets or sets the name of the container used to store scheduled jobs.
+    /// </summary>
+    public string ContainerName { get; set; } = "scheduled-jobs";
+}

src/Azure/Orleans.ScheduledJobs.AzureStorage/Orleans.ScheduledJobs.AzureStorage.csproj

@@ -0,0 +1,24 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <PackageReadmeFile>README.md</PackageReadmeFile>
+    <PackageId>Microsoft.Orleans.ScheduledJobs.AzureStorage</PackageId>
+    <Title>Microsoft Orleans Azure Storage ScheduledJobs Provider</Title>
+    <Description>Microsoft Orleans reminders provider backed by Azure Storage</Description>
+    <PackageTags>$(PackageTags) Azure Storage</PackageTags>
+    <TargetFrameworks>$(DefaultTargetFrameworks)</TargetFrameworks>
+    <AssemblyName>Orleans.ScheduledJobs.AzureStorage</AssemblyName>
+    <RootNamespace>Orleans.ScheduledJobs.AzureStorage</RootNamespace>
+    <OrleansBuildTimeCodeGen>true</OrleansBuildTimeCodeGen>
+    <DefineConstants>$(DefineConstants);ORLEANS_REMINDERS</DefineConstants>
+    <Nullable>enable</Nullable>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="$(SourceRoot)src\Orleans.Runtime\Orleans.Runtime.csproj" />
+    <ProjectReference Include="$(SourceRoot)src\Orleans.ScheduledJobs\Orleans.ScheduledJobs.csproj" />
+    <PackageReference Include="Azure.Core" />
+    <PackageReference Include="Azure.Storage.Blobs" />
+  </ItemGroup>
+
+</Project>

src/Orleans.Core/Properties/AssemblyInfo.cs

@@ -4,6 +4,7 @@
 [assembly: InternalsVisibleTo("Orleans.CodeGeneration")]
 [assembly: InternalsVisibleTo("Orleans.CodeGeneration.Build")]
 [assembly: InternalsVisibleTo("Orleans.Runtime")]
+[assembly: InternalsVisibleTo("Orleans.ScheduledJobs")]
 [assembly: InternalsVisibleTo("Orleans.Streaming")]
 [assembly: InternalsVisibleTo("Orleans.TestingHost")]