Merged PR 712962: Add a gRPC ClusterState implementation

This PR adds a ClusterState storage implementation that is updated via gRPC using protobuf-net

Author: jbayardo

Public/Src/Cache/ContentStore/Distributed/ContentLocations/MachineLocation.cs

@@ -18,11 +18,13 @@ public readonly record struct MachineLocation
     {
         public const string GrpcUriSchemePrefix = "grpc://";
 
+        public static MachineLocation Invalid { get; } = new(string.Empty);
+
         /// <summary>
         /// Gets whether the current machine location represents valid data
         /// </summary>
         [JsonIgnore]
-        public bool IsValid => Path != null;
+        public bool IsValid => !string.IsNullOrEmpty(Path);
 
         /// <summary>
         /// Gets the path representation of the machine location

Public/Src/Cache/ContentStore/Distributed/Ephemeral/GrpcClusterStateStorageClient.cs

@@ -0,0 +1,62 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading.Tasks;
+using BuildXL.Cache.ContentStore.Distributed.MetadataService;
+using BuildXL.Cache.ContentStore.Distributed.NuCache;
+using BuildXL.Cache.ContentStore.Distributed.NuCache.ClusterStateManagement;
+using BuildXL.Cache.ContentStore.Interfaces.Results;
+using BuildXL.Cache.ContentStore.Interfaces.Time;
+using BuildXL.Cache.ContentStore.Tracing;
+using BuildXL.Cache.ContentStore.Tracing.Internal;
+using BuildXL.Cache.ContentStore.Utils;
+using BuildXL.Cache.Host.Configuration;
+
+#nullable enable
+
+namespace BuildXL.Cache.ContentStore.Distributed.Ephemeral;
+
+public class GrpcClusterStateStorageClient : GrpcCodeFirstClient<IGrpcClusterStateStorage>, IClusterStateStorage
+{
+    public record Configuration(TimeSpan OperationTimeout, RetryPolicyConfiguration RetryPolicy);
+
+    protected override Tracer Tracer { get; } = new(nameof(GrpcClusterStateStorageClient));
+
+    public GrpcClusterStateStorageClient(Configuration configuration, IClientAccessor<IGrpcClusterStateStorage> accessor)
+        : base(accessor, CreateRetryPolicy(configuration.RetryPolicy), SystemClock.Instance, configuration.OperationTimeout)
+    {
+    }
+
+    private static IRetryPolicy CreateRetryPolicy(RetryPolicyConfiguration configurationRetryPolicy)
+    {
+        return configurationRetryPolicy.AsRetryPolicy(_ => true,
+            // We use an absurdly high retry count because the actual operation timeout is controlled through
+            // PerformOperationAsync in ExecuteAsync.
+            1_000_000);
+    }
+
+    public Task<Result<IClusterStateStorage.HeartbeatOutput>> HeartbeatAsync(OperationContext context, IClusterStateStorage.HeartbeatInput request)
+    {
+        return ExecuteAsync(
+            context,
+            async (context, options, service) => Result.Success(await service.HeartbeatAsync(request, options)),
+            extraEndMessage: _ => string.Empty);
+    }
+
+    public Task<Result<ClusterStateMachine>> ReadStateAsync(OperationContext context)
+    {
+        return ExecuteAsync(
+            context,
+            async (context, options, service) => Result.Success(await service.ReadStateAsync(options)),
+            extraEndMessage: _ => string.Empty);
+    }
+
+    public Task<Result<IClusterStateStorage.RegisterMachineOutput>> RegisterMachinesAsync(OperationContext context, IClusterStateStorage.RegisterMachineInput request)
+    {
+        return ExecuteAsync(
+            context,
+            async (context, options, service) => Result.Success(await service.RegisterMachinesAsync(request, options)),
+            extraEndMessage: _ => string.Empty);
+    }
+}

Public/Src/Cache/ContentStore/Distributed/Ephemeral/GrpcClusterStateStorageService.cs

@@ -0,0 +1,69 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Threading.Tasks;
+using BuildXL.Cache.ContentStore.Distributed.MetadataService;
+using BuildXL.Cache.ContentStore.Distributed.NuCache;
+using BuildXL.Cache.ContentStore.Distributed.NuCache.ClusterStateManagement;
+using BuildXL.Cache.ContentStore.Interfaces.Logging;
+using BuildXL.Cache.ContentStore.Interfaces.Results;
+using BuildXL.Cache.ContentStore.Interfaces.Tracing;
+using BuildXL.Cache.ContentStore.Logging;
+using BuildXL.Cache.ContentStore.Tracing;
+using BuildXL.Cache.ContentStore.Tracing.Internal;
+using BuildXL.Cache.ContentStore.Utils;
+using ProtoBuf.Grpc;
+
+#nullable enable
+
+namespace BuildXL.Cache.ContentStore.Distributed.Ephemeral;
+
+public class GrpcClusterStateStorageService : StartupShutdownComponentBase, IGrpcClusterStateStorage
+{
+    protected override Tracer Tracer { get; } = new(nameof(GrpcClusterStateStorageService));
+
+    private readonly IClusterStateStorage _clusterStateStorage;
+
+    private ILogger _logger = NullLogger.Instance;
+
+    public GrpcClusterStateStorageService(IClusterStateStorage clusterStateStorage)
+    {
+        _clusterStateStorage = clusterStateStorage;
+        LinkLifetime(_clusterStateStorage);
+    }
+
+    protected override Task<BoolResult> StartupComponentAsync(OperationContext context)
+    {
+        _logger = context.TracingContext.Logger;
+        return base.StartupComponentAsync(context);
+    }
+
+    public Task<IClusterStateStorage.RegisterMachineOutput> RegisterMachinesAsync(IClusterStateStorage.RegisterMachineInput request, CallContext callContext = default)
+    {
+        var operationContext = CreateOperationContext(callContext);
+        return _clusterStateStorage.RegisterMachinesAsync(operationContext, request).ThrowIfFailureAsync();
+    }
+
+    public Task<IClusterStateStorage.HeartbeatOutput> HeartbeatAsync(IClusterStateStorage.HeartbeatInput request, CallContext callContext = default)
+    {
+        var operationContext = CreateOperationContext(callContext);
+        return _clusterStateStorage.HeartbeatAsync(operationContext, request).ThrowIfFailureAsync();
+    }
+
+    public Task<ClusterStateMachine> ReadStateAsync(CallContext callContext = default)
+    {
+        var operationContext = CreateOperationContext(callContext);
+        return _clusterStateStorage.ReadStateAsync(operationContext).ThrowIfFailureAsync();
+    }
+
+    private OperationContext CreateOperationContext(CallContext callContext)
+    {
+        var contextId = MetadataServiceSerializer.TryGetContextId(callContext.RequestHeaders);
+        var tracingContext = contextId != null
+            ? new Context(contextId, _logger)
+            : new Context(_logger);
+
+        var operationContext = new OperationContext(tracingContext, callContext.CancellationToken);
+        return operationContext;
+    }
+}

Public/Src/Cache/ContentStore/Distributed/Ephemeral/IGrpcClusterStateStorage.cs

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System.Threading.Tasks;
+using BuildXL.Cache.ContentStore.Distributed.NuCache;
+using BuildXL.Cache.ContentStore.Distributed.NuCache.ClusterStateManagement;
+using ProtoBuf.Grpc;
+using ProtoBuf.Grpc.Configuration;
+
+#nullable enable
+
+namespace BuildXL.Cache.ContentStore.Distributed.Ephemeral;
+
+/// <summary>
+/// This interface serves as a code-first service declaration for protobuf-net gRPC. Clients may talk to a service
+/// implementing this interface, which behind it servers a <see cref="IClusterStateStorage"/>.
+/// </summary>
+[Service("Cache.ClusterState")]
+public interface IGrpcClusterStateStorage
+{
+    Task<IClusterStateStorage.RegisterMachineOutput> RegisterMachinesAsync(IClusterStateStorage.RegisterMachineInput request, CallContext callContext = default);
+
+    Task<IClusterStateStorage.HeartbeatOutput> HeartbeatAsync(IClusterStateStorage.HeartbeatInput request, CallContext callContext = default);
+
+    Task<ClusterStateMachine> ReadStateAsync(CallContext callContext = default);
+}

Public/Src/Cache/ContentStore/Distributed/MetadataService/ClientGlobalCacheStore.cs

@@ -1,10 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Runtime.CompilerServices;
 using System.Threading.Tasks;
 using BuildXL.Cache.ContentStore.Distributed.NuCache;
 using BuildXL.Cache.ContentStore.Hashing;
@@ -15,8 +13,6 @@
 using BuildXL.Cache.ContentStore.Utils;
 using BuildXL.Cache.MemoizationStore.Interfaces.Results;
 using BuildXL.Cache.MemoizationStore.Interfaces.Sessions;
-using BuildXL.Utilities.Core.Tracing;
-using Grpc.Core;
 
 #nullable enable
 
@@ -25,104 +21,31 @@ namespace BuildXL.Cache.ContentStore.Distributed.MetadataService
     /// <summary>
     /// A global content metadata store client which routes requests to a remote machine.
     /// </summary>
-    public class ClientGlobalCacheStore : StartupShutdownComponentBase, IGlobalCacheStore
+    public class ClientGlobalCacheStore : GrpcCodeFirstClient<IGlobalCacheService>, IGlobalCacheStore
     {
         /// <inheritdoc />
         public override bool AllowMultipleStartupAndShutdowns => true;
 
         /// <inheritdoc />
-        protected override Tracer Tracer { get; } = new Tracer(nameof(ClientGlobalCacheStore));
-
-        private readonly IClientAccessor<IGlobalCacheService> _serviceClientFactory;
-
-        private readonly ClientContentMetadataStoreConfiguration _configuration;
-
-        private readonly IClock _clock;
-
-        private readonly IRetryPolicy _retryPolicy;
+        protected override Tracer Tracer { get; } = new(nameof(ClientGlobalCacheStore));
 
         public ClientGlobalCacheStore(
             IClientAccessor<IGlobalCacheService> metadataServiceClientFactory,
             ClientContentMetadataStoreConfiguration configuration)
+            : base(metadataServiceClientFactory, CreateRetryPolicy(configuration), SystemClock.Instance, configuration.OperationTimeout)
         {
-            _serviceClientFactory = metadataServiceClientFactory;
-            _configuration = configuration;
-            _clock = SystemClock.Instance;
+        }
 
-            _retryPolicy = RetryPolicyFactory.GetExponentialPolicy(
+        private static IRetryPolicy CreateRetryPolicy(ClientContentMetadataStoreConfiguration configuration)
+        {
+            return RetryPolicyFactory.GetExponentialPolicy(
                 _ => true,
                 // We use an absurdly high retry count because the actual operation timeout is controlled through
                 // PerformOperationAsync in ExecuteAsync.
                 1_000_000,
-                _configuration.RetryMinimumWaitTime,
-                _configuration.RetryMaximumWaitTime,
-                _configuration.RetryDelta);
-
-            LinkLifetime(_serviceClientFactory);
-        }
-
-        private async Task<TResult> ExecuteAsync<TResult>(
-            OperationContext originalContext,
-            Func<OperationContext, CallOptions, IGlobalCacheService, Task<TResult>> executeAsync,
-            Func<TResult, string?> extraEndMessage,
-            string? extraStartMessage = null,
-            [CallerMemberName] string caller = null!)
-            where TResult : ResultBase
-        {
-            var attempt = -1;
-            using var contextWithShutdown = TrackShutdown(originalContext);
-            var context = contextWithShutdown.Context;
-            var callerAttempt = $"{caller}_Attempt";
-
-            return await context.PerformOperationWithTimeoutAsync(
-                Tracer,
-                context =>
-                {
-                    var callOptions = new CallOptions(
-                        headers: new Metadata()
-                        {
-                            MetadataServiceSerializer.CreateContextIdHeaderEntry(context.TracingContext.TraceId)
-                        },
-                        deadline: _clock.UtcNow + _configuration.OperationTimeout,
-                        cancellationToken: context.Token);
-
-                    return _retryPolicy.ExecuteAsync(async () =>
-                    {
-                        await Task.Yield();
-
-                        attempt++;
-
-                        var stopwatch = StopwatchSlim.Start();
-                        var clientCreationTime = TimeSpan.Zero;
-
-                        var result = await context.PerformOperationAsync(Tracer, () =>
-                            {
-                                return _serviceClientFactory.UseAsync(context, service =>
-                                {
-                                    clientCreationTime = stopwatch.Elapsed;
-
-                                    return executeAsync(context, callOptions, service);
-                                });
-                            },
-                            extraStartMessage: extraStartMessage,
-                            extraEndMessage: r => $"Attempt=[{attempt}] ClientCreationTimeMs=[{clientCreationTime.TotalMilliseconds}] {extraEndMessage(r)}",
-                            caller: callerAttempt,
-                            traceErrorsOnly: true);
-
-                        await Task.Yield();
-
-                        // Because we capture exceptions inside the PerformOperation, we need to make sure that they
-                        // get propagated for the retry policy to kick in.
-                        result.RethrowIfFailure();
-
-                        return result;
-                    }, context.Token);
-                },
-                caller: caller,
-                traceErrorsOnly: true,
-                extraStartMessage: extraStartMessage,
-                extraEndMessage: r => $"Attempts=[{attempt + 1}] {extraEndMessage(r)}",
-                timeout: _configuration.OperationTimeout);
+                configuration.RetryMinimumWaitTime,
+                configuration.RetryMaximumWaitTime,
+                configuration.RetryDelta);
         }
 
         public Task<Result<IReadOnlyList<ContentLocationEntry>>> GetBulkAsync(OperationContext context, IReadOnlyList<ShortHash> contentHashes)
@@ -243,15 +166,5 @@ public Task<Result<SerializedMetadataEntry>> GetContentHashListAsync(OperationCo
             // TODO: What to log here?
             extraEndMessage: r => r.GetValueOrDefault()?.ToString());
         }
-
-        public Task<Result<MachineMapping>> RegisterMachineAsync(OperationContext context, MachineLocation machineLocation)
-        {
-            throw new NotImplementedException($"Attempt to use {nameof(ClientGlobalCacheStore)} for machine registration is unsupported");
-        }
-
-        public Task<BoolResult> ForceRegisterMachineAsync(OperationContext context, MachineMapping mapping)
-        {
-            throw new NotImplementedException($"Attempt to use {nameof(ClientGlobalCacheStore)} for machine registration is unsupported");
-        }
     }
 }