1 files changed, 164 insertions, 0 deletions
diff --git a/config-model/src/main/java/com/yahoo/config/model/provision/InMemoryProvisioner.java b/config-model/src/main/java/com/yahoo/config/model/provision/InMemoryProvisioner.java
new file mode 100644
index 00000000000..c06efece329
--- /dev/null
+++ b/config-model/src/main/java/com/yahoo/config/model/provision/InMemoryProvisioner.java
@@ -0,0 +1,164 @@
+// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.config.model.provision;
+
+import com.yahoo.collections.ListMap;
+import com.yahoo.collections.Pair;
+import com.yahoo.config.model.api.HostProvisioner;
+import com.yahoo.config.provision.*;
+
+import java.util.*;
+
+/**
+ * In memory host provisioner. NB! ATM cannot be reused after allocate has been called.
+ *
+ * @author musum
+ * @author bratseth
+ */
+public class InMemoryProvisioner implements HostProvisioner {
+
+    /**
+     * If this is true an exception is thrown when all nodes are used.
+     * If false this will simply return nodes best effort, preferring to satisfy the
+     * number of groups requested when possible.
+     */
+    private final boolean failOnOutOfCapacity;
+
+    /** Hosts which should be returned as retired */
+    private final Set<String> retiredHostNames;
+
+    /** Free hosts of each flavor */
+    private final ListMap<String, Host> freeNodes = new ListMap<>();
+    private final Map<String, HostSpec> legacyMapping = new LinkedHashMap<>();
+    private final Map<ClusterSpec, List<HostSpec>> allocations = new LinkedHashMap<>();
+
+    /** Indexes must be unique across all groups in a cluster */
+    private final Map<Pair<ClusterSpec.Type,ClusterSpec.Id>, Integer> nextIndexInCluster = new HashMap<>();
+
+    /** Use this index as start index for all clusters */
+    private final int startIndexForClusters;
+
+    /** Creates this with a number of nodes of the flavor 'default' */
+    public InMemoryProvisioner(int nodeCount) {
+        this(Collections.singletonMap("default", createHostInstances(nodeCount)), true, 0);
+    }
+
+    /** Creates this with a set of host names of the flavor 'default' */
+    public InMemoryProvisioner(boolean failOnOutOfCapacity, String... hosts) {
+        this(Collections.singletonMap("default", toHostInstances(hosts)), failOnOutOfCapacity, 0);
+    }
+
+    /** Creates this with a set of hosts of the flavor 'default' */
+    public InMemoryProvisioner(Hosts hosts, boolean failOnOutOfCapacity, String ... retiredHostNames) {
+        this(Collections.singletonMap("default", hosts.getHosts()), failOnOutOfCapacity, 0, retiredHostNames);
+    }
+
+    /** Creates this with a set of hosts of the flavor 'default' */
+    public InMemoryProvisioner(Hosts hosts, boolean failOnOutOfCapacity, int startIndexForClusters, String ... retiredHostNames) {
+        this(Collections.singletonMap("default", hosts.getHosts()), failOnOutOfCapacity, startIndexForClusters, retiredHostNames);
+    }
+
+    public InMemoryProvisioner(Map<String, Collection<Host>> hosts, boolean failOnOutOfCapacity, int startIndexForClusters, String ... retiredHostNames) {
+        this.failOnOutOfCapacity = failOnOutOfCapacity;
+        for (Map.Entry<String, Collection<Host>> hostsOfFlavor : hosts.entrySet())
+            for (Host host : hostsOfFlavor.getValue())
+                freeNodes.put(hostsOfFlavor.getKey(), host);
+        this.retiredHostNames = new HashSet<>(Arrays.asList(retiredHostNames));
+        this.startIndexForClusters = startIndexForClusters;
+    }
+
+    private static Collection<Host> toHostInstances(String[] hostnames) {
+        List<Host> hosts = new ArrayList<>();
+        for (String hostname : hostnames) {
+            hosts.add(new Host(hostname));
+        }
+        return hosts;
+    }
+
+    private static Collection<Host> createHostInstances(int hostCount) {
+        List<Host> hosts = new ArrayList<>();
+        for (int i = 1; i <= hostCount; i++) {
+            hosts.add(new Host("host" + i));
+        }
+        return hosts;
+    }
+
+    @Override
+    public HostSpec allocateHost(String alias) {
+        if (legacyMapping.containsKey(alias)) return legacyMapping.get(alias);
+        List<Host> defaultHosts = freeNodes.get("default");
+        if (defaultHosts.isEmpty()) throw new IllegalArgumentException("No more hosts of default flavor available");
+        Host newHost = freeNodes.removeValue("default", 0);
+        HostSpec hostSpec = new HostSpec(newHost.getHostname(), newHost.getHostAliases());
+        legacyMapping.put(alias, hostSpec);
+        return hostSpec;
+    }
+
+    @Override
+    public List<HostSpec> prepare(ClusterSpec cluster, Capacity requestedCapacity, int groups, ProvisionLogger logger) {
+        if (cluster.group().isPresent() && groups > 1)
+            throw new IllegalArgumentException("Cannot both be specifying a group and ask for groups to be created");
+        if (requestedCapacity.nodeCount() % groups != 0)
+            throw new IllegalArgumentException("Requested " + requestedCapacity.nodeCount() + " nodes in " +
+                                               groups + " groups, but the node count is not divisible into this number of groups");
+
+        int capacity = failOnOutOfCapacity ? requestedCapacity.nodeCount() :
+                                             Math.min(requestedCapacity.nodeCount(), freeNodes.get("default").size() + totalAllocatedTo(cluster));
+        if (groups > capacity)
+            groups = capacity;
+
+        String flavor = requestedCapacity.flavor().orElse("default");
+
+        List<HostSpec> allocation = new ArrayList<>();
+        if (groups == 1) {
+            allocation.addAll(allocateHostGroup(cluster, flavor, capacity, startIndexForClusters));
+        }
+        else {
+            for (int i = 0; i < groups; i++) {
+                allocation.addAll(allocateHostGroup(cluster.changeGroup(Optional.of(ClusterSpec.Group.from(String.valueOf(i)))),
+                                                    flavor,
+                                                    capacity / groups,
+                                                    allocation.size()));
+            }
+        }
+        for (ListIterator<HostSpec> i = allocation.listIterator(); i.hasNext(); ) {
+            HostSpec host = i.next();
+            if (retiredHostNames.contains(host.hostname()))
+                i.set(retire(host));
+        }
+        return allocation;
+    }
+
+    private HostSpec retire(HostSpec host) {
+        return new HostSpec(host.hostname(), host.aliases(), host.membership().get().retire());
+    }
+
+    private List<HostSpec> allocateHostGroup(ClusterSpec clusterGroup, String flavor, int nodesInGroup, int startIndex) {
+        List<HostSpec> allocation = allocations.getOrDefault(clusterGroup, new ArrayList<>());
+        allocations.put(clusterGroup, allocation);
+
+        int nextIndex = nextIndexInCluster.getOrDefault(new Pair<>(clusterGroup.type(), clusterGroup.id()), startIndex);
+        while (allocation.size() < nodesInGroup) {
+            if (freeNodes.get(flavor).isEmpty()) throw new IllegalArgumentException("No nodes of flavor '" + flavor + "' available");
+            Host newHost = freeNodes.removeValue(flavor, 0);
+            ClusterMembership membership = ClusterMembership.from(clusterGroup, nextIndex++);
+            allocation.add(new HostSpec(newHost.getHostname(), newHost.getHostAliases(), membership));
+        }
+        nextIndexInCluster.put(new Pair<>(clusterGroup.type(), clusterGroup.id()), nextIndex);
+
+        while (allocation.size() > nodesInGroup)
+            allocation.remove(0);
+
+        return allocation;
+    }
+
+    private int totalAllocatedTo(ClusterSpec cluster) {
+        int count = 0;
+        for (Map.Entry<ClusterSpec, List<HostSpec>> allocation : allocations.entrySet()) {
+            if ( ! allocation.getKey().type().equals(cluster.type())) continue;
+            if ( ! allocation.getKey().id().equals(cluster.id())) continue;
+            count += allocation.getValue().size();
+        }
+        return count;
+    }
+
+}