Updates container location support

Framework/Core/CoreInterface.cs

@@ -16,6 +16,11 @@ namespace Core
             return entryPoint.GetPlugin<T>();
         }
 
+        public IKnownLocations GetKnownLocations()
+        {
+            return entryPoint.Tools.CreateWorkflow().GetAvailableLocations();
+        }
+
         public IDownloadedLog DownloadLog(IHasContainer containerSource, int? tailLines = null)
         {
             return DownloadLog(containerSource.Container, tailLines);

Framework/KubernetesWorkflow/K8sCluster.cs

@@ -11,7 +11,6 @@ namespace KubernetesWorkflow
 
         public Configuration Configuration { get; }
         public string HostAddress { get; private set; } = string.Empty;
-        public K8sNodeLabel[] AvailableK8sNodes { get; set; } = new K8sNodeLabel[0];
 
         public KubernetesClientConfiguration GetK8sClientConfig()
         {
@@ -20,20 +19,6 @@ namespace KubernetesWorkflow
             return config;
         }
 
-        public K8sNodeLabel? GetNodeLabelForLocation(Location location)
-        {
-            switch (location)
-            {
-                case Location.One:
-                    return K8sNodeIfAvailable(0);
-                case Location.Two:
-                    return K8sNodeIfAvailable(1);
-                case Location.Three:
-                    return K8sNodeIfAvailable(2);
-            }
-            return null;
-        }
-
         public TimeSpan K8sOperationTimeout()
         {
             return Configuration.OperationTimeout;
@@ -68,12 +53,6 @@ namespace KubernetesWorkflow
                 HostAddress = config.Host;
             }
         }
-
-        private K8sNodeLabel? K8sNodeIfAvailable(int index)
-        {
-            if (AvailableK8sNodes.Length <= index) return null;
-            return AvailableK8sNodes[index];
-        }
     }
 
     public class K8sNodeLabel

Framework/KubernetesWorkflow/K8sController.cs

@@ -28,11 +28,10 @@ namespace KubernetesWorkflow
         {
             client.Dispose();
         }
-        
-        public RunningPod BringOnline(ContainerRecipe[] containerRecipes, Location location)
+
+        public RunningPod BringOnline(ContainerRecipe[] containerRecipes, ILocation location)
         {
             log.Debug();
-            DiscoverK8sNodes();
             EnsureTestNamespace();
 
             var deploymentName = CreateDeployment(containerRecipes, location);
@@ -109,19 +108,7 @@ namespace KubernetesWorkflow
 
         #region Discover K8s Nodes
 
-        private void DiscoverK8sNodes()
-        {
-            if (cluster.AvailableK8sNodes == null || !cluster.AvailableK8sNodes.Any())
-            {
-                cluster.AvailableK8sNodes = GetAvailableK8sNodes();
-                if (cluster.AvailableK8sNodes.Length < 3)
-                {
-                    log.Debug($"Warning: For full location support, at least 3 Kubernetes Nodes are required in the cluster. Nodes found: '{string.Join(",", cluster.AvailableK8sNodes.Select(p => $"{p.Key}={p.Value}"))}'.");
-                }
-            }
-        }
-
-        private K8sNodeLabel[] GetAvailableK8sNodes()
+        public K8sNodeLabel[] GetAvailableK8sNodes()
         {
             var nodes = client.Run(c => c.ListNode());
 
@@ -322,7 +309,7 @@ namespace KubernetesWorkflow
 
         #region Deployment management
 
-        private string CreateDeployment(ContainerRecipe[] containerRecipes, Location location)
+        private string CreateDeployment(ContainerRecipe[] containerRecipes, ILocation location)
         {
             var deploymentSpec = new V1Deployment
             {
@@ -364,9 +351,9 @@ namespace KubernetesWorkflow
             WaitUntilDeploymentOffline(deploymentName);
         }
 
-        private IDictionary<string, string> CreateNodeSelector(Location location)
+        private IDictionary<string, string> CreateNodeSelector(ILocation location)
         {
-            var nodeLabel = cluster.GetNodeLabelForLocation(location);
+            var nodeLabel = GetNodeLabelForLocation(location);
             if (nodeLabel == null) return new Dictionary<string, string>();
 
             return new Dictionary<string, string>
@@ -375,6 +362,12 @@ namespace KubernetesWorkflow
             };
         }
 
+        private K8sNodeLabel? GetNodeLabelForLocation(ILocation location)
+        {
+            var l = (Location)location;
+            return l.NodeLabel;
+        }
+
         private IDictionary<string, string> GetSelector(ContainerRecipe[] containerRecipes)
         {
             return containerRecipes.First().PodLabels.GetLabels();

Framework/KubernetesWorkflow/KnownLocations.cs

@@ -0,0 +1,42 @@
+namespace KubernetesWorkflow
+{
+    public interface IKnownLocations
+    {
+        /// <summary>
+        /// Returns a known location given an index.
+        /// Each index guarantees a different location.
+        /// </summary>
+        ILocation Get(int index);
+        int NumberOfLocations { get; }
+
+        /// <summary>
+        /// Returns the location object for a specific kubernetes node. Throws if it doesn't exist.
+        /// </summary>
+        ILocation Get(string kubeNodeName);
+    }
+
+    public class KnownLocations : IKnownLocations
+    {
+        private readonly Location[] locations;
+
+        public KnownLocations(Location[] locations)
+        {
+            this.locations = locations;
+            if (locations.Any(l => l.NodeLabel == null)) throw new Exception("Must not contain unspecified location");
+        }
+
+        public static ILocation UnspecifiedLocation { get; } = new Location();
+
+        public int NumberOfLocations => locations.Length;
+
+        public ILocation Get(int index)
+        {
+            return locations[index];
+        }
+
+        public ILocation Get(string kubeNodeName)
+        {
+            return locations.Single(l => l.NodeLabel != null && l.NodeLabel.Value == kubeNodeName);
+        }
+    }
+}

Framework/KubernetesWorkflow/Location.cs

@@ -1,10 +1,22 @@
 namespace KubernetesWorkflow
 {
-    public enum Location
+    public interface ILocation
     {
-        Unspecified,
-        One,
-        Two,
-        Three,
+    }
+
+    public class Location : ILocation
+    {
+        internal Location(K8sNodeLabel? nodeLabel = null)
+        {
+            NodeLabel = nodeLabel;
+        }
+
+        internal K8sNodeLabel? NodeLabel { get; }
+
+        public override string ToString()
+        {
+            if (NodeLabel == null) return "Location:Unspecified";
+            return $"Location:KubeNode-'{NodeLabel.Key}:{NodeLabel.Value}'";
+        }
     }
 }

Framework/KubernetesWorkflow/LocationProvider.cs

@@ -0,0 +1,49 @@
+using Logging;
+
+namespace KubernetesWorkflow
+{
+    public class LocationProvider
+    {
+        private readonly TimeSpan locationsExpirationTime = TimeSpan.FromMinutes(10);
+        private readonly ILog log;
+        private readonly Action<Action<K8sController>> onController;
+        private Location[] knownLocations = Array.Empty<Location>();
+        private DateTime lastUpdate = DateTime.UtcNow;
+
+        public LocationProvider(ILog log, Action<Action<K8sController>> onController)
+        {
+            this.log = log;
+            this.onController = onController;
+        }
+
+        public IKnownLocations GetAvailableLocations()
+        {
+            if (ShouldUpdateKnownLocations())
+            {
+                onController(UpdateKnownLocations);
+            }
+
+            return new KnownLocations(knownLocations);
+        }
+
+        private void UpdateKnownLocations(K8sController controller)
+        {
+            knownLocations = controller.GetAvailableK8sNodes().Select(CreateLocation).ToArray();
+            lastUpdate = DateTime.UtcNow;
+
+            log.Log($"Detected {knownLocations.Length} available locations: '{string.Join(",", knownLocations.Select(l => l.ToString()))}'");
+        }
+
+        private Location CreateLocation(K8sNodeLabel k8sNode)
+        {
+            return new Location(k8sNode);
+        }
+
+        private bool ShouldUpdateKnownLocations()
+        {
+            if (!knownLocations.Any()) return true;
+            if (DateTime.UtcNow - lastUpdate > locationsExpirationTime) return true;
+            return false;
+        }
+    }
+}

Framework/KubernetesWorkflow/StartupWorkflow.cs

@@ -5,7 +5,9 @@ namespace KubernetesWorkflow
 {
     public interface IStartupWorkflow
     {
-        RunningContainers Start(int numberOfContainers, Location location, ContainerRecipeFactory recipeFactory, StartupConfig startupConfig);
+        IKnownLocations GetAvailableLocations();
+        RunningContainers Start(int numberOfContainers, ContainerRecipeFactory recipeFactory, StartupConfig startupConfig);
+        RunningContainers Start(int numberOfContainers, ILocation location, ContainerRecipeFactory recipeFactory, StartupConfig startupConfig);
         CrashWatcher CreateCrashWatcher(RunningContainer container);
         void Stop(RunningContainers runningContainers);
         void DownloadContainerLog(RunningContainer container, ILogHandler logHandler, int? tailLines = null);
@@ -22,6 +24,7 @@ namespace KubernetesWorkflow
         private readonly KnownK8sPods knownK8SPods;
         private readonly string k8sNamespace;
         private readonly RecipeComponentFactory componentFactory = new RecipeComponentFactory();
+        private readonly LocationProvider locationProvider;
 
         internal StartupWorkflow(ILog log, WorkflowNumberSource numberSource, K8sCluster cluster, KnownK8sPods knownK8SPods, string k8sNamespace)
         {
@@ -30,9 +33,21 @@ namespace KubernetesWorkflow
             this.cluster = cluster;
             this.knownK8SPods = knownK8SPods;
             this.k8sNamespace = k8sNamespace;
+
+            locationProvider = new LocationProvider(log, K8s);
         }
 
-        public RunningContainers Start(int numberOfContainers, Location location, ContainerRecipeFactory recipeFactory, StartupConfig startupConfig)
+        public IKnownLocations GetAvailableLocations()
+        {
+            return locationProvider.GetAvailableLocations();
+        }
+
+        public RunningContainers Start(int numberOfContainers, ContainerRecipeFactory recipeFactory, StartupConfig startupConfig)
+        {
+            return Start(numberOfContainers, KnownLocations.UnspecifiedLocation, recipeFactory, startupConfig);
+        }
+
+        public RunningContainers Start(int numberOfContainers, ILocation location, ContainerRecipeFactory recipeFactory, StartupConfig startupConfig)
         {
             return K8s(controller =>
             {

ProjectPlugins/CodexContractsPlugin/CodexContractsStarter.cs

@@ -24,7 +24,7 @@ namespace CodexContractsPlugin
             var startupConfig = CreateStartupConfig(gethNode);
             startupConfig.NameOverride = "codex-contracts";
 
-            var containers = workflow.Start(1, Location.Unspecified, new CodexContractsContainerRecipe(), startupConfig);
+            var containers = workflow.Start(1, new CodexContractsContainerRecipe(), startupConfig);
             if (containers.Containers.Length != 1) throw new InvalidOperationException("Expected 1 Codex contracts container to be created. Test infra failure.");
             var container = containers.Containers[0];
 

ProjectPlugins/CodexPlugin/CodexSetup.cs

@@ -8,7 +8,7 @@ namespace CodexPlugin
     public interface ICodexSetup
     {
         ICodexSetup WithName(string name);
-        ICodexSetup At(Location location);
+        ICodexSetup At(ILocation location);
         ICodexSetup WithBootstrapNode(ICodexNode node);
         ICodexSetup WithLogLevel(CodexLogLevel level);
         /// <summary>
@@ -43,7 +43,7 @@ namespace CodexPlugin
             return this;
         }
 
-        public ICodexSetup At(Location location)
+        public ICodexSetup At(ILocation location)
         {
             Location = location;
             return this;

ProjectPlugins/CodexPlugin/CodexStarter.cs

@@ -69,7 +69,7 @@ namespace CodexPlugin
             return startupConfig;
         }
 
-        private RunningContainers[] StartCodexContainers(StartupConfig startupConfig, int numberOfNodes, Location location)
+        private RunningContainers[] StartCodexContainers(StartupConfig startupConfig, int numberOfNodes, ILocation location)
         {
             var result = new List<RunningContainers>();
             for (var i = 0; i < numberOfNodes; i++)

ProjectPlugins/CodexPlugin/CodexStartupConfig.cs

@@ -6,7 +6,7 @@ namespace CodexPlugin
     public class CodexStartupConfig
     {
         public string? NameOverride { get; set; }
-        public Location Location { get; set; }
+        public ILocation Location { get; set; } = KnownLocations.UnspecifiedLocation;
         public CodexLogLevel LogLevel { get; set; }
         public string[]? LogTopics { get; set; }
         public ByteSize? StorageQuota { get; set; }

ProjectPlugins/GethPlugin/GethStarter.cs

@@ -21,7 +21,7 @@ namespace GethPlugin
             startupConfig.NameOverride = gethStartupConfig.NameOverride;
 
             var workflow = tools.CreateWorkflow();
-            var containers = workflow.Start(1, Location.Unspecified, new GethContainerRecipe(), startupConfig);
+            var containers = workflow.Start(1, new GethContainerRecipe(), startupConfig);
             if (containers.Containers.Length != 1) throw new InvalidOperationException("Expected 1 Geth bootstrap node to be created. Test infra failure.");
             var container = containers.Containers[0];
 

ProjectPlugins/MetricsPlugin/PrometheusStarter.cs

@@ -21,7 +21,7 @@ namespace MetricsPlugin
             startupConfig.Add(new PrometheusStartupConfig(GeneratePrometheusConfig(targets)));
 
             var workflow = tools.CreateWorkflow();
-            var runningContainers = workflow.Start(1, Location.Unspecified, recipe, startupConfig);
+            var runningContainers = workflow.Start(1, recipe, startupConfig);
             if (runningContainers.Containers.Length != 1) throw new InvalidOperationException("Expected only 1 Prometheus container to be created.");
 
             Log("Metrics server started.");

Tests/CodexTests/BasicTests/TwoClientTests.cs

@@ -1,6 +1,5 @@
 using CodexPlugin;
 using DistTestCore;
-using KubernetesWorkflow;
 using NUnit.Framework;
 using Utils;
 
@@ -23,8 +22,15 @@ namespace Tests.BasicTests
         [Test]
         public void TwoClientsTwoLocationsTest()
         {
-            var primary = Ci.StartCodexNode(s => s.At(Location.One));
-            var secondary = Ci.StartCodexNode(s => s.At(Location.Two));
+            var locations = Ci.GetKnownLocations();
+            if (locations.NumberOfLocations < 2)
+            {
+                Assert.Inconclusive("Two-locations test requires 2 nodes to be available in the cluster.");
+                return;
+            }
+
+            var primary = Ci.StartCodexNode(s => s.At(locations.Get(0)));
+            var secondary = Ci.StartCodexNode(s => s.At(locations.Get(1)));
 
             PerformTwoClientTest(primary, secondary);
         }