Merged PR 730455: Migrate to GRPC format in MachineLocation

Ignacio Alonso Battaglia · Ignacio Alonso Battaglia · commit b64183f7f72a · 2023-08-15T19:20:25.000Z
Add a flag to use the host in MachinoLocation _Equals_ function. This allows us to migrate to the GRPC format without breaking the match between the old and the new format.

**NOTE:** We should revert this changes after migration.
diff --git a/Public/Src/Cache/ContentStore/Distributed/ContentLocations/MachineLocation.cs b/Public/Src/Cache/ContentStore/Distributed/ContentLocations/MachineLocation.cs
@@ -16,7 +16,12 @@ namespace BuildXL.Cache.ContentStore.Distributed
     /// </summary>
     public readonly record struct MachineLocation
     {
-        public const string GrpcUriSchemePrefix = "grpc://";
+        private const string GrpcUriSchemePrefix = "grpc://";
+
+        // TODO: This is a temporary solution while we migrate to the new format
+        // When set to true the old and the new format will be considered equal
+        // Work item to remove this https://dev.azure.com/mseng/1ES/_workitems/edit/2095358
+        public static bool OnlyUseHostToCompare = false;
 
         public static MachineLocation Invalid { get; } = new(string.Empty);
 
@@ -45,12 +50,24 @@ public bool Equals(MachineLocation other)
                 return other.Path is null;
             }
 
+            if (OnlyUseHostToCompare)
+            {
+                return ExtractHost().Equals(other.ExtractHost(), StringComparison.InvariantCultureIgnoreCase);
+            }
+
             return Path.Equals(other.Path, StringComparison.InvariantCultureIgnoreCase);
         }
 
         public override int GetHashCode()
         {
-            return Path is not null ? StringComparer.InvariantCultureIgnoreCase.GetHashCode(Path) : 42;
+            if (Path is null)
+            {
+                return 42;
+            }
+            // TODO: This is a temporary solution while we migrate to the new format
+            // Same machine has same hash code
+            var host = ExtractHost();
+            return StringComparer.InvariantCultureIgnoreCase.GetHashCode(host);
         }
 
         /// <inheritdoc />
@@ -64,6 +81,12 @@ public static MachineLocation Create(string machineName, int port)
             return new MachineLocation($"{MachineLocation.GrpcUriSchemePrefix}{machineName}:{port}/");
         }
 
+        private string ExtractHost()
+        {
+            var (extractedHost, _) = ExtractHostInfo();
+            return extractedHost;
+        }
+
         public (string host, int? port) ExtractHostInfo()
         {
             if (Path.StartsWith(GrpcUriSchemePrefix))
diff --git a/Public/Src/Cache/ContentStore/Distributed/NuCache/ClusterStateManagement/ClusterStateMachine.cs b/Public/Src/Cache/ContentStore/Distributed/NuCache/ClusterStateManagement/ClusterStateMachine.cs
@@ -25,8 +25,10 @@ public static (ClusterStateMachine NextState, MachineRecord[] Result) HeartbeatM
         var priorMachineRecords = new MachineRecord[request.MachineIds.Count];
         foreach (var entry in request.MachineIds.AsIndexed())
         {
+            MachineLocation? newLocation = request.MachineLocations != null ? request.MachineLocations[entry.Index] : null;
+
             (state, priorMachineRecords[entry.Index]) =
-                state.Heartbeat(entry.Item, nowUtc, request.MachineState).ThrowIfFailure();
+                state.Heartbeat(entry.Item, nowUtc, request.MachineState, newLocation).ThrowIfFailure();
         }
 
         return (state, priorMachineRecords);
@@ -209,7 +211,7 @@ private ClusterStateMachine ForceRegisterMachineWithState(
         }
     }
 
-    public Result<(ClusterStateMachine Next, MachineRecord Previous)> Heartbeat(MachineId machineId, DateTime nowUtc, MachineState state)
+    public Result<(ClusterStateMachine Next, MachineRecord Previous)> Heartbeat(MachineId machineId, DateTime nowUtc, MachineState state, MachineLocation? machineLocation = null)
     {
         if (!machineId.Valid)
         {
@@ -222,7 +224,14 @@ private ClusterStateMachine ForceRegisterMachineWithState(
         MachineRecord? previous = null;
         foreach (var record in Records)
         {
-            if (record.Id == machineId)
+            if (record.Id == machineId && machineLocation != null)
+            {
+                // TODO: This is a hack to allow the machine location to be updated https://dev.azure.com/mseng/1ES/_workitems/edit/2095358
+                // Update the machine location in case it has changed due to migration to a different format (e.g. GRPC format)
+                records.Add(record.Heartbeat(nowUtc, state, (MachineLocation)machineLocation));
+                previous = record;
+            }
+            else if (record.Id == machineId)
             {
                 records.Add(record.Heartbeat(nowUtc, state));
                 previous = record;
diff --git a/Public/Src/Cache/ContentStore/Distributed/NuCache/ClusterStateManagement/ClusterStateManager.cs b/Public/Src/Cache/ContentStore/Distributed/NuCache/ClusterStateManagement/ClusterStateManager.cs
@@ -165,6 +165,7 @@ public Task<Result<MachineState>> HeartbeatAsync(
                 }
 
                 var localMachineIds = ClusterState.LocalMachineMappings.Select(machineMapping => machineMapping.Id).ToArray();
+                var localMachineLocations = ClusterState.LocalMachineMappings.Select(machineMapping => machineMapping.Location).ToArray();
 
                 if (_configuration.ReadOnly || localMachineIds.Length == 0)
                 {
@@ -177,7 +178,7 @@ public Task<Result<MachineState>> HeartbeatAsync(
                 }
                 else
                 {
-                    var heartbeatResponse = await _storage.HeartbeatAsync(context, new IClusterStateStorage.HeartbeatInput(localMachineIds, machineState)).ThrowIfFailureAsync();
+                    var heartbeatResponse = await _storage.HeartbeatAsync(context, new IClusterStateStorage.HeartbeatInput(localMachineIds, machineState, localMachineLocations)).ThrowIfFailureAsync();
                     Contract.Assert(heartbeatResponse.PriorRecords.Length == localMachineIds.Length, "Mismatch between number of requested heartbeats and actual heartbeats. This should never happen.");
 
                     ClusterState.Update(context, heartbeatResponse.State, configuration: _configuration.RecomputeConfiguration, nowUtc: _clock.UtcNow).ThrowIfFailure();
diff --git a/Public/Src/Cache/ContentStore/Distributed/NuCache/ClusterStateManagement/IClusterStateStorage.cs b/Public/Src/Cache/ContentStore/Distributed/NuCache/ClusterStateManagement/IClusterStateStorage.cs
@@ -39,7 +39,7 @@ public record RegisterMachineOutput(ClusterStateMachine State, MachineMapping[]
     public Task<Result<RegisterMachineOutput>> RegisterMachinesAsync(OperationContext context, RegisterMachineInput request);
 
     [ProtoContract(ImplicitFields = ImplicitFields.AllPublic)]
-    public record HeartbeatInput(IReadOnlyList<MachineId> MachineIds, MachineState MachineState)
+    public record HeartbeatInput(IReadOnlyList<MachineId> MachineIds, MachineState MachineState, IReadOnlyList<MachineLocation> MachineLocations = null)
     {
         /// <summary>
         /// This parameterless constructor exists only to allow ProtoBuf.NET initialization
diff --git a/Public/Src/Cache/ContentStore/Distributed/NuCache/ClusterStateManagement/MachineRecord.cs b/Public/Src/Cache/ContentStore/Distributed/NuCache/ClusterStateManagement/MachineRecord.cs
@@ -40,6 +40,16 @@ internal MachineRecord Heartbeat(DateTime nowUtc, MachineState nextState)
             return this with { State = nextState, LastHeartbeatTimeUtc = nowUtc };
         }
 
+        internal MachineRecord Heartbeat(DateTime nowUtc, MachineState nextState, MachineLocation machineLocation)
+        {
+            if (nextState == MachineState.Unknown)
+            {
+                return this with { LastHeartbeatTimeUtc = nowUtc, Location = machineLocation};
+            }
+
+            return this with { State = nextState, LastHeartbeatTimeUtc = nowUtc, Location = machineLocation };
+        }
+
         public bool IsOpen()
         {
             return State == MachineState.Open;
diff --git a/Public/Src/Cache/ContentStore/Distributed/NuCache/TransitioningContentLocationStore.cs b/Public/Src/Cache/ContentStore/Distributed/NuCache/TransitioningContentLocationStore.cs
@@ -75,6 +75,7 @@ protected override async Task<BoolResult> StartupCoreAsync(OperationContext cont
                 if (LocalLocationStore.ClusterState.TryResolveMachineId(LocalMachineLocation, out var localMachineId))
                 {
                     LocalMachineId = localMachineId;
+                    Tracer.Info(context, $"Resolved MachineId {LocalMachineId} for {LocalMachineLocation}");
                 }
                 else if (_configuration.DistributedContentConsumerOnly)
                 {
diff --git a/Public/Src/Cache/ContentStore/Distributed/Stores/DistributedContentStoreSettings.cs b/Public/Src/Cache/ContentStore/Distributed/Stores/DistributedContentStoreSettings.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Runtime.Serialization;
 using BuildXL.Cache.ContentStore.Distributed.NuCache.CopyScheduling;
 using BuildXL.Cache.ContentStore.Distributed.Sessions;
 using BuildXL.Cache.ContentStore.Interfaces.Distributed;
@@ -106,6 +107,11 @@ public sealed class DistributedContentStoreSettings
         /// </summary>
         public int ProactiveCopyMaxRetries { get; set; } = 0;
 
+        /// <summary>
+        /// Whether to use the value of the host inside Equals function in MachineLocation
+        /// </summary>
+        public bool UseHostInMachineLocationEquals { get; set; } = false;
+
         /// <summary>
         /// Defines pinning behavior
         /// </summary>
diff --git a/Public/Src/Cache/ContentStore/Distributed/Utilities/GrpcFileCopier.cs b/Public/Src/Cache/ContentStore/Distributed/Utilities/GrpcFileCopier.cs
@@ -206,7 +206,7 @@ public MachineLocation GetLocalMachineLocation(AbsolutePath cacheRoot)
         {
             if (_configuration.UseUniversalLocations)
             {
-                return new MachineLocation($"{MachineLocation.GrpcUriSchemePrefix}{_localMachineName}:{_configuration.GrpcPort}/");
+                return MachineLocation.Create(_localMachineName, _configuration.GrpcPort);
             }
 
             if (!cacheRoot.IsLocal)
diff --git a/Public/Src/Cache/ContentStore/DistributedTest/ContentLocation/NuCache/MachineListTests.cs b/Public/Src/Cache/ContentStore/DistributedTest/ContentLocation/NuCache/MachineListTests.cs
@@ -271,7 +271,7 @@ public MachineListFactory(ITestClock clock, int amountMachines, int designatedLo
 
                 this.MachineIds = Enumerable.Range(1, amountMachines).Select(n => (ushort)n).ToArray();
 
-                var machineMappings = MachineIds.Select(m => new MachineMapping(new MachineId(m), new MachineLocation(m.ToString()))).ToArray();
+                var machineMappings = MachineIds.Select(m => new MachineMapping(new MachineId(m), new MachineLocation($"grpc://{m}:123"))).ToArray();
                 var clusterState = new ClusterState(primaryMachineId: default, machineMappings);
                 foreach (var mapping in machineMappings)
                 {
diff --git a/Public/Src/Cache/ContentStore/DistributedTest/Stores/DistributedContentCopierTests.cs b/Public/Src/Cache/ContentStore/DistributedTest/Stores/DistributedContentCopierTests.cs
@@ -55,7 +55,7 @@ public async Task CopyFromInRingMachines()
                     new OperationContext(context),
                     hashWithLocations,
                     handleCopyAsync: tpl => Task.FromResult(new PutResult(hash, 42)),
-                    inRingMachines: new MachineLocation[] { new MachineLocation("") });
+                    inRingMachines: new MachineLocation[] { new MachineLocation(@"grpc://fun.com:123") });
 
                 result.ShouldBeSuccess();
             }
@@ -99,7 +99,7 @@ public async Task CopyFailsForWrongCopySize()
                 var hashWithLocations = new ContentHashWithSizeAndLocations(
                     hash,
                     size: 42,
-                    new MachineLocation[] {new MachineLocation("")});
+                    new MachineLocation[] {new MachineLocation(@"grpc://fun.com:123") });
 
                 mockFileCopier.CopyToAsyncResult = CopyFileResult.SuccessWithSize(41);
                 var result = await distributedCopier.TryCopyAndPutAsync(
@@ -127,7 +127,7 @@ public async Task CopyFailsForWrongHash()
                 var hashWithLocations = new ContentHashWithSizeAndLocations(
                     hash,
                     size: 42,
-                    new MachineLocation[] {new MachineLocation("")});
+                    new MachineLocation[] {new MachineLocation(@"grpc://fun.com:123") });
 
                 mockFileCopier.CopyToAsyncResult = CopyFileResult.SuccessWithSize(42);
                 var result = await distributedCopier.TryCopyAndPutAsync(
@@ -152,7 +152,7 @@ public async Task CopyRetries(int retries)
                 var (distributedCopier, mockFileCopier) = CreateMocks(FileSystem, directory.Path,TimeSpan.Zero, retries);
                 await using var _ = await distributedCopier.StartupWithAutoShutdownAsync(context);
 
-                var machineLocations = new MachineLocation[] {new MachineLocation("")};
+                var machineLocations = new MachineLocation[] {new MachineLocation(@"grpc://fun.com:123") };
 
                 var hash = ContentHash.Random();
                 var hashWithLocations = new ContentHashWithSizeAndLocations(
@@ -190,7 +190,7 @@ public async Task CopyRetriesWithRestrictions(int retries)
                     maxRetryCount: retries + 1);
                 await using var _ = await distributedCopier.StartupWithAutoShutdownAsync(context);
 
-                var empty = new MachineLocation("");
+                var empty = new MachineLocation(@"grpc://fun.com:123");
                 var machineLocations = new [] { empty, empty, empty, empty, empty };
 
                 var hash = ContentHash.Random();
@@ -245,7 +245,7 @@ public async Task CopyWithDestinationPathError(int retries)
                 var (distributedCopier, mockFileCopier) = CreateMocks(FileSystem, directory.Path, TimeSpan.FromMilliseconds((10)), retries);
                 await using var _ = await distributedCopier.StartupWithAutoShutdownAsync(context);
 
-                var machineLocations = new MachineLocation[] { new MachineLocation(""), new MachineLocation("") };
+                var machineLocations = new MachineLocation[] { new MachineLocation(@"grpc://fun.com:123"), new MachineLocation(@"grpc://fun.com:123") };
 
                 var hash = ContentHash.Random();
                 var hashWithLocations = new ContentHashWithSizeAndLocations(
@@ -306,7 +306,7 @@ public class MockFileCopier : IRemoteFileCopier
 #pragma warning restore 649
             public CopyFileResult[] CustomResults;
 
-            public MachineLocation GetLocalMachineLocation(AbsolutePath cacheRoot) => new MachineLocation("");
+            public MachineLocation GetLocalMachineLocation(AbsolutePath cacheRoot) => new MachineLocation(@"grpc://fun.com:123");
 
             /// <inheritdoc />
             public Task<CopyFileResult> CopyToAsync(OperationContext context, ContentLocation sourceLocation, Stream destinationStream, CopyOptions options)
diff --git a/Public/Src/Cache/DistributedCache.Host/Configuration/DistributedContentSettings.cs b/Public/Src/Cache/DistributedCache.Host/Configuration/DistributedContentSettings.cs
@@ -896,6 +896,9 @@ public bool UseGrpcDotNetForCopies()
         [DataMember]
         public bool UseUniversalLocations { get; set; }
 
+        [DataMember]
+        public bool UseHostInMachineLocationEquals { get; set; } = false;
+
         /// <summary>
         /// Include domain name in machine location.
         /// </summary>
diff --git a/Public/Src/Cache/DistributedCache.Host/Service/Internal/DistributedContentStoreFactory.cs b/Public/Src/Cache/DistributedCache.Host/Service/Internal/DistributedContentStoreFactory.cs
@@ -84,6 +84,9 @@ public DistributedContentStoreFactory(DistributedCacheServiceArguments arguments
             ContentLocationStoreConfiguration = CreateContentLocationStoreConfiguration();
             _distributedContentStoreSettings = CreateDistributedStoreSettings(_arguments, ContentLocationStoreConfiguration);
 
+            // TODO: This is a temporary solution while we migrate to the new format
+            MachineLocation.OnlyUseHostToCompare = _distributedContentStoreSettings.UseHostInMachineLocationEquals;
+
             // Tracing configuration before creating anything.
             if (arguments.TraceConfiguration)
             {
@@ -369,6 +372,7 @@ private static DistributedContentStoreSettings CreateDistributedStoreSettings(
             {
 
                 PinConfiguration = pinConfiguration,
+                UseHostInMachineLocationEquals = distributedSettings.UseHostInMachineLocationEquals,
                 ProactiveCopyMode = (ProactiveCopyMode)Enum.Parse(typeof(ProactiveCopyMode), distributedSettings.ProactiveCopyMode),
                 PushProactiveCopies = distributedSettings.PushProactiveCopies,
                 ProactiveCopyOnPut = distributedSettings.ProactiveCopyOnPut,

Original file line number	Diff line number	Diff line change
`@@ -25,8 +25,10 @@ public static (ClusterStateMachine NextState, MachineRecord[] Result) HeartbeatM`
`25`	`25`	`var priorMachineRecords = new MachineRecord[request.MachineIds.Count];`
`26`	`26`	`foreach (var entry in request.MachineIds.AsIndexed())`
`27`	`27`	`{`
	`28`	`+ MachineLocation? newLocation = request.MachineLocations != null ? request.MachineLocations[entry.Index] : null;`
	`29`	`+`
`28`	`30`	`(state, priorMachineRecords[entry.Index]) =`
`29`		`- state.Heartbeat(entry.Item, nowUtc, request.MachineState).ThrowIfFailure();`
	`31`	`+ state.Heartbeat(entry.Item, nowUtc, request.MachineState, newLocation).ThrowIfFailure();`
`30`	`32`	`}`
`31`	`33`
`32`	`34`	`return (state, priorMachineRecords);`
`@@ -209,7 +211,7 @@ private ClusterStateMachine ForceRegisterMachineWithState(`
`209`	`211`	`}`
`210`	`212`	`}`
`211`	`213`
`212`		`- public Result<(ClusterStateMachine Next, MachineRecord Previous)> Heartbeat(MachineId machineId, DateTime nowUtc, MachineState state)`
	`214`	`+ public Result<(ClusterStateMachine Next, MachineRecord Previous)> Heartbeat(MachineId machineId, DateTime nowUtc, MachineState state, MachineLocation? machineLocation = null)`
`213`	`215`	`{`
`214`	`216`	`if (!machineId.Valid)`
`215`	`217`	`{`
`@@ -222,7 +224,14 @@ private ClusterStateMachine ForceRegisterMachineWithState(`
`222`	`224`	`MachineRecord? previous = null;`
`223`	`225`	`foreach (var record in Records)`
`224`	`226`	`{`
`225`		`- if (record.Id == machineId)`
	`227`	`+ if (record.Id == machineId && machineLocation != null)`
	`228`	`+ {`
	`229`	`+ // TODO: This is a hack to allow the machine location to be updated https://dev.azure.com/mseng/1ES/_workitems/edit/2095358`
	`230`	`+ // Update the machine location in case it has changed due to migration to a different format (e.g. GRPC format)`
	`231`	`+ records.Add(record.Heartbeat(nowUtc, state, (MachineLocation)machineLocation));`
	`232`	`+ previous = record;`
	`233`	`+ }`
	`234`	`+ else if (record.Id == machineId)`
`226`	`235`	`{`
`227`	`236`	`records.Add(record.Heartbeat(nowUtc, state));`
`228`	`237`	`previous = record;`
Original file line number	Diff line number	Diff line change
`@@ -165,6 +165,7 @@ public Task<Result<MachineState>> HeartbeatAsync(`
`165`	`165`	`}`
`166`	`166`
`167`	`167`	`var localMachineIds = ClusterState.LocalMachineMappings.Select(machineMapping => machineMapping.Id).ToArray();`
	`168`	`+ var localMachineLocations = ClusterState.LocalMachineMappings.Select(machineMapping => machineMapping.Location).ToArray();`
`168`	`169`
`169`	`170`	`if (_configuration.ReadOnly \|\| localMachineIds.Length == 0)`
`170`	`171`	`{`
`@@ -177,7 +178,7 @@ public Task<Result<MachineState>> HeartbeatAsync(`
`177`	`178`	`}`
`178`	`179`	`else`
`179`	`180`	`{`
`180`		`- var heartbeatResponse = await _storage.HeartbeatAsync(context, new IClusterStateStorage.HeartbeatInput(localMachineIds, machineState)).ThrowIfFailureAsync();`
	`181`	`+ var heartbeatResponse = await _storage.HeartbeatAsync(context, new IClusterStateStorage.HeartbeatInput(localMachineIds, machineState, localMachineLocations)).ThrowIfFailureAsync();`
`181`	`182`	`Contract.Assert(heartbeatResponse.PriorRecords.Length == localMachineIds.Length, "Mismatch between number of requested heartbeats and actual heartbeats. This should never happen.");`
`182`	`183`
`183`	`184`	`ClusterState.Update(context, heartbeatResponse.State, configuration: _configuration.RecomputeConfiguration, nowUtc: _clock.UtcNow).ThrowIfFailure();`
Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@ public record RegisterMachineOutput(ClusterStateMachine State, MachineMapping[]`
`39`	`39`	`public Task<Result<RegisterMachineOutput>> RegisterMachinesAsync(OperationContext context, RegisterMachineInput request);`
`40`	`40`
`41`	`41`	`[ProtoContract(ImplicitFields = ImplicitFields.AllPublic)]`
`42`		`- public record HeartbeatInput(IReadOnlyList<MachineId> MachineIds, MachineState MachineState)`
	`42`	`+ public record HeartbeatInput(IReadOnlyList<MachineId> MachineIds, MachineState MachineState, IReadOnlyList<MachineLocation> MachineLocations = null)`
`43`	`43`	`{`
`44`	`44`	`/// <summary>`
`45`	`45`	`/// This parameterless constructor exists only to allow ProtoBuf.NET initialization`
Original file line number	Diff line number	Diff line change
`@@ -75,6 +75,7 @@ protected override async Task<BoolResult> StartupCoreAsync(OperationContext cont`
`75`	`75`	`if (LocalLocationStore.ClusterState.TryResolveMachineId(LocalMachineLocation, out var localMachineId))`
`76`	`76`	`{`
`77`	`77`	`LocalMachineId = localMachineId;`
	`78`	`+ Tracer.Info(context, $"Resolved MachineId {LocalMachineId} for {LocalMachineLocation}");`
`78`	`79`	`}`
`79`	`80`	`else if (_configuration.DistributedContentConsumerOnly)`
`80`	`81`	`{`
Original file line number	Diff line number	Diff line change
`@@ -206,7 +206,7 @@ public MachineLocation GetLocalMachineLocation(AbsolutePath cacheRoot)`
`206`	`206`	`{`
`207`	`207`	`if (_configuration.UseUniversalLocations)`
`208`	`208`	`{`
`209`		`- return new MachineLocation($"{MachineLocation.GrpcUriSchemePrefix}{_localMachineName}:{_configuration.GrpcPort}/");`
	`209`	`+ return MachineLocation.Create(_localMachineName, _configuration.GrpcPort);`
`210`	`210`	`}`
`211`	`211`
`212`	`212`	`if (!cacheRoot.IsLocal)`
Original file line number	Diff line number	Diff line change
`@@ -271,7 +271,7 @@ public MachineListFactory(ITestClock clock, int amountMachines, int designatedLo`
`271`	`271`
`272`	`272`	`this.MachineIds = Enumerable.Range(1, amountMachines).Select(n => (ushort)n).ToArray();`
`273`	`273`
`274`		`- var machineMappings = MachineIds.Select(m => new MachineMapping(new MachineId(m), new MachineLocation(m.ToString()))).ToArray();`
	`274`	`+ var machineMappings = MachineIds.Select(m => new MachineMapping(new MachineId(m), new MachineLocation($"grpc://{m}:123"))).ToArray();`
`275`	`275`	`var clusterState = new ClusterState(primaryMachineId: default, machineMappings);`
`276`	`276`	`foreach (var mapping in machineMappings)`
`277`	`277`	`{`