Adding ability to balance load between different cell servers.

Adding ability to specify structure/size of the cell.

Change-Id: I5e87c99fe8812ba0a974d7815ab8ddc64193a608

@@ -113,7 +113,10 @@ function cell {
     case "$cell" in
     "borrow")
         aux="/tmp/cell-$$"
-        curl -sS -X POST "http://$CELL_WARDEN:4321/?duration=${2:-0}&user=${3:-$(id -un)}" \
+        spec=${3:-3+1}
+        spec=${spec/+/%2B}
+        user=${4:-$(id -un)}
+        curl -sS -X POST "http://$CELL_WARDEN:4321/?duration=${2:-0}&spec=${spec}&user=${user}" \
             -d "$(cat ~/.ssh/id_rsa.pub)" > $aux
         . $aux
         rm -f $aux

@@ -3,15 +3,23 @@
 
 name="$1"
 ipx="$2"
-shift 2
+spec="$3"
+shift 3
 key="$@"
 
 cd $(dirname $0)
 
+nodes=${spec%+*}
+mininet=${spec#*+}
+
 sudo lxc-attach -n bit-proxy -- bash -c "grep -qF \"$key\" /home/sdn/.ssh/authorized_keys || echo $key >> /home/sdn/.ssh/authorized_keys"
 
-./clone-node base-mininet ${ipx/x/0} $name-n "$key"
+if [ $mininet -ge 1 ]; then
+    ./clone-node base-mininet ${ipx/x/0} $name-n "$key"
+fi
 
-for n in {1..3}; do
+let n=1
+while [ $n -le $nodes ]; do
     ./clone-node base-onos ${ipx/x/$n} $name-$n "$key"
+    let n=n+1
 done

@@ -2,11 +2,19 @@
 # Destroys an LXC cell with the specified name.
 
 name=$1
+spec=$2
+
+nodes=${spec%+*}
+mininet=${spec#*+}
 
 cd $(dirname $0)
 
-./destroy-node $name-n
+if [ $mininet -ge 1 ]; then
+    ./destroy-node $name-n
+fi
 
-for n in {1..3}; do
+let n=1
+while [ $n -le $nodes ]; do
     ./destroy-node $name-$n
+    let n=n+1
 done

@@ -42,6 +42,23 @@
             <artifactId>jetty-servlet</artifactId>
             <version>8.1.18.v20150929</version>
         </dependency>
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <version>4.12</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.onosproject</groupId>
+            <artifactId>onlab-misc</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.onosproject</groupId>
+            <artifactId>onlab-junit</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>
@@ -77,6 +94,11 @@
                     </execution>
                 </executions>
             </plugin>
+
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-surefire-plugin</artifactId>
+            </plugin>
         </plugins>
     </build>
 

@@ -27,13 +27,15 @@ final class Reservation {
     final String userName;
     final long time;
     final int duration;
+    final String cellSpec;
 
     // Creates a new reservation record
-    Reservation(String cellName, String userName, long time, int duration) {
+    Reservation(String cellName, String userName, long time, int duration, String cellSpec) {
         this.cellName = cellName;
         this.userName = userName;
         this.time = time;
         this.duration = duration;
+        this.cellSpec = cellSpec;
     }
 
     /**
@@ -43,11 +45,12 @@ final class Reservation {
      */
     Reservation(String line) {
         String[] fields = line.trim().split("\t");
-        checkState(fields.length == 4, "Incorrect reservation encoding");
+        checkState(fields.length == 5, "Incorrect reservation encoding");
         this.cellName = fields[0];
         this.userName = fields[1];
         this.time = Long.parseLong(fields[2]);
         this.duration = Integer.parseInt(fields[3]);
+        this.cellSpec = fields[4];
     }
 
     /**
@@ -56,7 +59,7 @@ final class Reservation {
      * @return encoded string
      */
     String encode() {
-        return String.format("%s\t%s\t%s\t%s\n", cellName, userName, time, duration);
+        return String.format("%s\t%s\t%s\t%s\t%s\n", cellName, userName, time, duration, cellSpec);
     }
 
 }

@@ -16,8 +16,9 @@
 
 package org.onlab.warden;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import com.google.common.io.ByteStreams;
 
 import java.io.File;
@@ -27,8 +28,11 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.PrintWriter;
 import java.text.SimpleDateFormat;
+import java.util.ArrayList;
 import java.util.Date;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.Random;
 import java.util.Set;
 import java.util.Timer;
@@ -52,9 +56,15 @@ class Warden {
     private static final int MINUTE = 60_000; // 1 minute
     private static final int DEFAULT_MINUTES = 60;
 
+    private static final String DEFAULT_SPEC = "3+1";
+
     private final File log = new File("warden.log");
 
-    private final File cells = new File("cells");
+    // Allow overriding these for unit tests.
+    static String cmdPrefix = "";
+    static File root = new File(".");
+
+    private final File cells = new File(root, "cells");
     private final File supported = new File(cells, "supported");
     private final File reserved = new File(cells, "reserved");
 
@@ -64,6 +74,7 @@ class Warden {
      * Creates a new cell warden.
      */
     Warden() {
+        reserved.mkdirs();
         random.setSeed(System.currentTimeMillis());
         Timer timer = new Timer("cell-pruner", true);
         timer.schedule(new Reposessor(), MINUTE / 4, MINUTE / 2);
@@ -84,7 +95,7 @@ class Warden {
      *
      * @return list of cell names
      */
-    private Set<String> getAvailableCells() {
+    Set<String> getAvailableCells() {
         Set<String> available = new HashSet<>(getCells());
         available.removeAll(getReservedCells());
         return ImmutableSet.copyOf(available);
@@ -95,18 +106,28 @@ class Warden {
      *
      * @return list of cell names
      */
-    private Set<String> getReservedCells() {
+    Set<String> getReservedCells() {
         String[] list = reserved.list();
         return list != null ? ImmutableSet.copyOf(list) : ImmutableSet.of();
     }
 
     /**
+     * Returns the host name on which the specified cell is hosted.
+     *
+     * @param cellName cell name
+     * @return host name where the cell runs
+     */
+    String getCellHost(String cellName) {
+        return getCellInfo(cellName).hostName;
+    }
+
+    /**
      * Returns reservation for the specified user.
      *
      * @param userName user name
      * @return cell reservation record or null if user does not have one
      */
-    private Reservation currentUserReservation(String userName) {
+    Reservation currentUserReservation(String userName) {
         checkNotNull(userName, USER_NOT_NULL);
         for (String cellName : getReservedCells()) {
             Reservation reservation = currentCellReservation(cellName);
@@ -141,35 +162,66 @@ class Warden {
      *
      * @param userName user name
      * @param sshKey   user ssh public key
-     * @param minutes  number of minutes for reservation
+     * @param minutes  optional number of minutes for reservation
+     * @param cellSpec optional cell specification string
      * @return reserved cell definition
      */
-    synchronized String borrowCell(String userName, String sshKey, int minutes) {
+    synchronized String borrowCell(String userName, String sshKey, int minutes,
+                                   String cellSpec) {
         checkNotNull(userName, USER_NOT_NULL);
+        checkArgument(userName.matches("[\\w]+"), "Invalid user name %s", userName);
         checkNotNull(sshKey, KEY_NOT_NULL);
         checkArgument(minutes < MAX_MINUTES, "Number of minutes must be less than %d", MAX_MINUTES);
-        long now = System.currentTimeMillis();
+        checkArgument(minutes >= 0, "Number of minutes must be non-negative");
+        checkArgument(cellSpec == null || cellSpec.matches("[\\d]+\\+[0-1]"),
+                      "Invalid cell spec string %s", cellSpec);
         Reservation reservation = currentUserReservation(userName);
         if (reservation == null) {
-            checkArgument(minutes >= 0, "Number of minutes must be non-negative");
-            Set<String> cells = getAvailableCells();
-            checkState(!cells.isEmpty(), "No cells are presently available");
-            String cellName = ImmutableList.copyOf(cells).get(random.nextInt(cells.size()));
-            reservation = new Reservation(cellName, userName, now, minutes == 0 ? DEFAULT_MINUTES : minutes);
+            // If there is no reservation for the user, create one
+            String cellName = findAvailableCell();
+            reservation = new Reservation(cellName, userName, System.currentTimeMillis(),
+                                          minutes == 0 ? DEFAULT_MINUTES : minutes,
+                                          cellSpec == null ? DEFAULT_SPEC : cellSpec);
         } else if (minutes == 0) {
             // If minutes are 0, simply return the cell definition
             return getCellDefinition(reservation.cellName);
         } else {
-            reservation = new Reservation(reservation.cellName, userName, now, minutes);
+            // If minutes are > 0, update the existing cell reservation
+            reservation = new Reservation(reservation.cellName, userName,
+                                          System.currentTimeMillis(), minutes,
+                                          reservation.cellSpec);
         }
 
         reserveCell(reservation);
         createCell(reservation, sshKey);
-        log(userName, reservation.cellName, "borrowed for " + reservation.duration + " minutes");
+        log(userName, reservation.cellName, reservation.cellSpec,
+            "borrowed for " + reservation.duration + " minutes");
         return getCellDefinition(reservation.cellName);
     }
 
     /**
+     * Returns name of an available cell. Cell is chosen based on the load
+     * of its hosting server; a random one will be chosen from the set of
+     * cells hosted by the least loaded server.
+     *
+     * @return name of an available cell
+     */
+    private String findAvailableCell() {
+        Set<String> cells = getAvailableCells();
+        checkState(!cells.isEmpty(), "No cells are presently available");
+        Map<String, ServerInfo> load = Maps.newHashMap();
+
+        cells.stream().map(this::getCellInfo)
+                .forEach(info -> load.compute(info.hostName, (k, v) -> v == null ?
+                        new ServerInfo(info.hostName) : v.bumpLoad(info)));
+
+        List<ServerInfo> servers = new ArrayList<>(load.values());
+        servers.sort((a, b) -> a.load - b.load);
+        ServerInfo server = servers.get(0);
+        return server.cells.get(random.nextInt(server.cells.size())).cellName;
+    }
+
+    /**
      * Returns the specified cell for the specified user and their public access key.
      *
      * @param userName user name
@@ -181,7 +233,7 @@ class Warden {
 
         unreserveCell(reservation);
         destroyCell(reservation);
-        log(userName, reservation.cellName, "returned");
+        log(userName, reservation.cellName, reservation.cellSpec, "returned");
     }
 
     /**
@@ -229,9 +281,9 @@ class Warden {
      */
     private void createCell(Reservation reservation, String sshKey) {
         CellInfo cellInfo = getCellInfo(reservation.cellName);
-        String cmd = String.format("ssh %s warden/bin/create-cell %s %s %s",
+        String cmd = String.format("ssh %s warden/bin/create-cell %s %s %s %s",
                                    cellInfo.hostName, cellInfo.cellName,
-                                   cellInfo.ipPrefix, sshKey);
+                                   cellInfo.ipPrefix, reservation.cellSpec, sshKey);
         exec(cmd);
     }
 
@@ -242,8 +294,8 @@ class Warden {
      */
     private void destroyCell(Reservation reservation) {
         CellInfo cellInfo = getCellInfo(reservation.cellName);
-        exec(String.format("ssh %s warden/bin/destroy-cell %s",
-                           cellInfo.hostName, cellInfo.cellName));
+        exec(String.format("ssh %s warden/bin/destroy-cell %s %s",
+                           cellInfo.hostName, cellInfo.cellName, reservation.cellSpec));
     }
 
     /**
@@ -265,7 +317,7 @@ class Warden {
     // Executes the specified command.
     private String exec(String command) {
         try {
-            Process process = Runtime.getRuntime().exec(command);
+            Process process = Runtime.getRuntime().exec(cmdPrefix + command);
             String output = new String(ByteStreams.toByteArray(process.getInputStream()), UTF_8);
             process.waitFor(TIMEOUT, TimeUnit.SECONDS);
             return process.exitValue() == 0 ? output : null;
@@ -275,12 +327,12 @@ class Warden {
     }
 
     // Creates an audit log entry.
-    private void log(String userName, String cellName, String action) {
+    private void log(String userName, String cellName, String cellSpec, String action) {
         try (FileOutputStream fos = new FileOutputStream(log, true);
              PrintWriter pw = new PrintWriter(fos)) {
             SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
-            pw.println(String.format("%s\t%s\t%s\t%s", format.format(new Date()),
-                                     userName, cellName, action));
+            pw.println(String.format("%s\t%s\t%s-%s\t%s", format.format(new Date()),
+                                     userName, cellName, cellSpec, action));
             pw.flush();
         } catch (IOException e) {
             throw new IllegalStateException("Unable to log reservation action", e);
@@ -300,6 +352,23 @@ class Warden {
         }
     }
 
+    // Carrier of cell server information
+    private final class ServerInfo {
+        final String hostName;
+        int load = 0;
+        List<CellInfo> cells = Lists.newArrayList();
+
+        private ServerInfo(String hostName) {
+            this.hostName = hostName;
+        }
+
+        private ServerInfo bumpLoad(CellInfo info) {
+            cells.add(info);
+            load++;     // TODO: bump by cell size later
+            return this;
+        }
+    }
+
     // Task for re-possessing overdue cells
     private final class Reposessor extends TimerTask {
         @Override

@@ -49,8 +49,8 @@ public class WardenServlet extends HttpServlet {
                 if (reservation != null) {
                     long expiration = reservation.time + reservation.duration * 60_000;
                     long remaining = (expiration - System.currentTimeMillis()) / 60_000;
-                    out.println(String.format("%-10s\t%-10s\t%s\t%s\t%s mins (%s remaining)",
-                                              cellName,
+                    out.println(String.format("%-14s\t%-10s\t%s\t%s\t%s mins (%s remaining)",
+                                              cellName + "-" + reservation.cellSpec,
                                               reservation.userName,
                                               fmt.format(new Date(reservation.time)),
                                               fmt.format(new Date(expiration)),
@@ -72,8 +72,9 @@ public class WardenServlet extends HttpServlet {
             String sshKey = new String(ByteStreams.toByteArray(req.getInputStream()), "UTF-8");
             String userName = req.getParameter("user");
             String sd = req.getParameter("duration");
+            String spec = req.getParameter("spec");
             int duration = isNullOrEmpty(sd) ? 0 : Integer.parseInt(sd);
-            String cellDefinition = warden.borrowCell(userName, sshKey, duration);
+            String cellDefinition = warden.borrowCell(userName, sshKey, duration, spec);
             out.println(cellDefinition);
         } catch (Exception e) {
             resp.setStatus(Response.SC_INTERNAL_SERVER_ERROR);

+/*
+ * Copyright 2016 Open Networking Laboratory
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.onlab.warden;
+
+import com.google.common.io.Files;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.onlab.util.Tools;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Objects;
+
+import static org.junit.Assert.*;
+
+/**
+ * Suite of tests for the cell warden.
+ */
+public class WardenTest {
+
+    private Warden warden;
+    private File cells;
+    private File supportedCells;
+
+    @Before
+    public void setUp() throws IOException {
+        // Setup warden to be tested
+        Warden.root = Files.createTempDir();
+        Warden.cmdPrefix = "echo ";
+        cells = new File(Warden.root, "cells");
+        supportedCells = new File(cells, "supported");
+        warden = new Warden();
+
+        // Setup test cell information
+        createCell("alpha", "foo");
+        createCell("bravo", "foo");
+        createCell("charlie", "foo");
+        createCell("delta", "bar");
+        createCell("echo", "bar");
+        createCell("foxtrot", "bar");
+
+        new File("warden.log").deleteOnExit();
+    }
+
+    private void createCell(String cellName, String hostName) throws IOException {
+        File cellFile = new File(supportedCells, cellName);
+        Files.createParentDirs(cellFile);
+        Files.write((hostName + " " + cellName).getBytes(), cellFile);
+    }
+
+    @After
+    public void tearDown() throws IOException {
+        Tools.removeDirectory(Warden.root);
+    }
+
+    @Test
+    public void basics() {
+        assertEquals("incorrect number of cells", 6, warden.getCells().size());
+        validateSizes(6, 0);
+
+        String cellDefinition = warden.borrowCell("dude", "the-key", 0, null);
+        assertTrue("incorrect definition", cellDefinition.contains("cell-def"));
+        validateSizes(5, 1);
+
+        Reservation dudeCell = warden.currentUserReservation("dude");
+        validateCellState(dudeCell);
+
+        warden.borrowCell("dolt", "a-key", 0, "4+1");
+        Reservation doltCell = warden.currentUserReservation("dolt");
+        validateCellState(doltCell);
+        validateSizes(4, 2);
+
+        assertTrue("cells should not be on the same host",
+                   Objects.equals(warden.getCellHost(dudeCell.cellName),
+                                  warden.getCellHost(doltCell.cellName)));
+
+        warden.returnCell("dude");
+        validateSizes(5, 1);
+
+        warden.borrowCell("dolt", "a-key", 30, null);
+        validateSizes(5, 1);
+
+        warden.returnCell("dolt");
+        validateSizes(6, 0);
+    }
+
+    private void validateSizes(int available, int reserved) {
+        assertEquals("incorrect number of available cells", available,
+                     warden.getAvailableCells().size());
+        assertEquals("incorrect number of reserved cells", reserved,
+                     warden.getReservedCells().size());
+    }
+
+    private void validateCellState(Reservation reservation) {
+        assertFalse("cell should not be available",
+                    warden.getAvailableCells().contains(reservation.cellName));
+        assertTrue("cell should be reserved",
+                   warden.getReservedCells().contains(reservation.cellName));
+    }
+
+}
\ No newline at end of file