utils/warden/src/main/java/org/onlab/warden/Warden.java - onos - Gitiles

 /*
  * 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.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.io.ByteStreams;

 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.PrintWriter;
 import java.text.SimpleDateFormat;
 import java.util.Date;
 import java.util.HashSet;
 import java.util.Random;
 import java.util.Set;
 import java.util.Timer;
 import java.util.TimerTask;
 import java.util.concurrent.TimeUnit;

 import static com.google.common.base.Preconditions.*;

 /**
  * Warden for tracking use of shared test cells.
  */
 class Warden {

     private static final String CELL_NOT_NULL = "Cell name cannot be null";
     private static final String USER_NOT_NULL = "User name cannot be null";
     private static final String KEY_NOT_NULL = "User key cannot be null";
     private static final String UTF_8 = "UTF-8";

     private static final long TIMEOUT = 10; // 10 seconds
     private static final int MAX_MINUTES = 240; // 4 hours max
     private static final int MINUTE = 60_000; // 1 minute
     private static final int DEFAULT_MINUTES = 60;

     private final File log = new File("warden.log");

     private final File cells = new File("cells");
     private final File supported = new File(cells, "supported");
     private final File reserved = new File(cells, "reserved");

     private final Random random = new Random();

     /**
      * Creates a new cell warden.
      */
     Warden() {
         random.setSeed(System.currentTimeMillis());
         Timer timer = new Timer("cell-pruner", true);
         timer.schedule(new Reposessor(), MINUTE / 4, MINUTE / 2);
     }

     /**
      * Returns list of names of supported cells.
      *
      * @return list of cell names
      */
     Set<String> getCells() {
         String[] list = supported.list();
         return list != null ? ImmutableSet.copyOf(list) : ImmutableSet.of();
     }

     /**
      * Returns list of names of available cells.
      *
      * @return list of cell names
      */
     private Set<String> getAvailableCells() {
         Set<String> available = new HashSet<>(getCells());
         available.removeAll(getReservedCells());
         return ImmutableSet.copyOf(available);
     }

     /**
      * Returns list of names of reserved cells.
      *
      * @return list of cell names
      */
     private Set<String> getReservedCells() {
         String[] list = reserved.list();
         return list != null ? ImmutableSet.copyOf(list) : ImmutableSet.of();
     }

     /**
      * 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) {
         checkNotNull(userName, USER_NOT_NULL);
         for (String cellName : getReservedCells()) {
             Reservation reservation = currentCellReservation(cellName);
             if (reservation != null && userName.equals(reservation.userName)) {
                 return reservation;
             }
         }
         return null;
     }

     /**
      * Returns the name of the user who reserved the given cell.
      *
      * @param cellName cell name
      * @return cell reservation record or null if cell is not reserved
      */
     Reservation currentCellReservation(String cellName) {
         checkNotNull(cellName, CELL_NOT_NULL);
         File cellFile = new File(reserved, cellName);
         if (!cellFile.exists()) {
             return null;
         }
         try (InputStream stream = new FileInputStream(cellFile)) {
             return new Reservation(new String(ByteStreams.toByteArray(stream), "UTF-8"));
         } catch (IOException e) {
             throw new IllegalStateException("Unable to get current user for cell " + cellName, e);
         }
     }

     /**
      * Reserves a cell for the specified user and their public access key.
      *
      * @param userName user name
      * @param sshKey   user ssh public key
      * @param minutes  number of minutes for reservation
      * @return reserved cell definition
      */
     synchronized String borrowCell(String userName, String sshKey, int minutes) {
         checkNotNull(userName, USER_NOT_NULL);
         checkNotNull(sshKey, KEY_NOT_NULL);
         checkArgument(minutes < MAX_MINUTES, "Number of minutes must be less than %d", MAX_MINUTES);
         long now = System.currentTimeMillis();
         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);
         } 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);
         }

         reserveCell(reservation);
         createCell(reservation, sshKey);
         log(userName, reservation.cellName, "borrowed for " + reservation.duration + " minutes");
         return getCellDefinition(reservation.cellName);
     }

     /**
      * Returns the specified cell for the specified user and their public access key.
      *
      * @param userName user name
      */
     synchronized void returnCell(String userName) {
         checkNotNull(userName, USER_NOT_NULL);
         Reservation reservation = currentUserReservation(userName);
         checkState(reservation != null, "User %s has no cell reservations", userName);

         unreserveCell(reservation);
         destroyCell(reservation);
         log(userName, reservation.cellName, "returned");
     }

     /**
      * Reserves the specified cell for the user the source file and writes the
      * specified content to the target file.
      *
      * @param reservation cell reservation record
      */
     private void reserveCell(Reservation reservation) {
         File cellFile = new File(reserved, reservation.cellName);
         try (FileOutputStream stream = new FileOutputStream(cellFile)) {
             stream.write(reservation.encode().getBytes(UTF_8));
         } catch (IOException e) {
             throw new IllegalStateException("Unable to reserve cell " + reservation.cellName, e);
         }
     }

     /**
      * Returns the cell definition of the specified cell.
      *
      * @param cellName cell name
      * @return cell definition
      */
     private String getCellDefinition(String cellName) {
         CellInfo cellInfo = getCellInfo(cellName);
         return exec(String.format("ssh %s warden/bin/cell-def %s",
                                   cellInfo.hostName, cellInfo.cellName));
     }

     /**
      * Cancels the specified reservation.
      *
      * @param reservation reservation record
      */
     private void unreserveCell(Reservation reservation) {
         checkState(new File(reserved, reservation.cellName).delete(),
                    "Unable to return cell %s", reservation.cellName);
     }

     /**
      * Creates the cell for the specified user SSH key.
      *
      * @param reservation cell reservation
      * @param sshKey      ssh key
      */
     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",
                                    cellInfo.hostName, cellInfo.cellName,
                                    cellInfo.ipPrefix, sshKey);
         exec(cmd);
     }

     /**
      * Destroys the specified cell.
      *
      * @param reservation reservation record
      */
     private void destroyCell(Reservation reservation) {
         CellInfo cellInfo = getCellInfo(reservation.cellName);
         exec(String.format("ssh %s warden/bin/destroy-cell %s",
                            cellInfo.hostName, cellInfo.cellName));
     }

     /**
      * Reads the information about the specified cell.
      *
      * @param cellName cell name
      * @return cell information
      */
     private CellInfo getCellInfo(String cellName) {
         File cellFile = new File(supported, cellName);
         try (InputStream stream = new FileInputStream(cellFile)) {
             String[] fields = new String(ByteStreams.toByteArray(stream), UTF_8).split(" ");
             return new CellInfo(cellName, fields[0], fields[1]);
         } catch (IOException e) {
             throw new IllegalStateException("Unable to definition for cell " + cellName, e);
         }
     }

     // Executes the specified command.
     private String exec(String command) {
         try {
             Process process = Runtime.getRuntime().exec(command);
             String output = new String(ByteStreams.toByteArray(process.getInputStream()), UTF_8);
             process.waitFor(TIMEOUT, TimeUnit.SECONDS);
             return process.exitValue() == 0 ? output : null;
         } catch (Exception e) {
             throw new IllegalStateException("Unable to execute " + command);
         }
     }

     // Creates an audit log entry.
     private void log(String userName, String cellName, 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.flush();
         } catch (IOException e) {
             throw new IllegalStateException("Unable to log reservation action", e);
         }
     }

     // Carrier of cell information
     private final class CellInfo {
         final String cellName;
         final String hostName;
         final String ipPrefix;

         private CellInfo(String cellName, String hostName, String ipPrefix) {
             this.cellName = cellName;
             this.hostName = hostName;
             this.ipPrefix = ipPrefix;
         }
     }

     // Task for re-possessing overdue cells
     private final class Reposessor extends TimerTask {
         @Override
         public void run() {
             long now = System.currentTimeMillis();
             for (String cellName : getReservedCells()) {
                 Reservation reservation = currentCellReservation(cellName);
                 if (reservation != null &&
                         (reservation.time + reservation.duration * MINUTE) < now) {
                     try {
                         returnCell(reservation.userName);
                     } catch (Exception e) {
                         e.printStackTrace();
                     }
                 }
             }
         }
     }
 }
	/*
	* 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.collect.ImmutableList;
	import com.google.common.collect.ImmutableSet;
	import com.google.common.io.ByteStreams;

	import java.io.File;
	import java.io.FileInputStream;
	import java.io.FileOutputStream;
	import java.io.IOException;
	import java.io.InputStream;
	import java.io.PrintWriter;
	import java.text.SimpleDateFormat;
	import java.util.Date;
	import java.util.HashSet;
	import java.util.Random;
	import java.util.Set;
	import java.util.Timer;
	import java.util.TimerTask;
	import java.util.concurrent.TimeUnit;

	import static com.google.common.base.Preconditions.*;

	/**
	* Warden for tracking use of shared test cells.
	*/
	class Warden {

	private static final String CELL_NOT_NULL = "Cell name cannot be null";
	private static final String USER_NOT_NULL = "User name cannot be null";
	private static final String KEY_NOT_NULL = "User key cannot be null";
	private static final String UTF_8 = "UTF-8";

	private static final long TIMEOUT = 10; // 10 seconds
	private static final int MAX_MINUTES = 240; // 4 hours max
	private static final int MINUTE = 60_000; // 1 minute
	private static final int DEFAULT_MINUTES = 60;

	private final File log = new File("warden.log");

	private final File cells = new File("cells");
	private final File supported = new File(cells, "supported");
	private final File reserved = new File(cells, "reserved");

	private final Random random = new Random();

	/**
	* Creates a new cell warden.
	*/
	Warden() {
	random.setSeed(System.currentTimeMillis());
	Timer timer = new Timer("cell-pruner", true);
	timer.schedule(new Reposessor(), MINUTE / 4, MINUTE / 2);
	}

	/**
	* Returns list of names of supported cells.
	*
	* @return list of cell names
	*/
	Set<String> getCells() {
	String[] list = supported.list();
	return list != null ? ImmutableSet.copyOf(list) : ImmutableSet.of();
	}

	/**
	* Returns list of names of available cells.
	*
	* @return list of cell names
	*/
	private Set<String> getAvailableCells() {
	Set<String> available = new HashSet<>(getCells());
	available.removeAll(getReservedCells());
	return ImmutableSet.copyOf(available);
	}

	/**
	* Returns list of names of reserved cells.
	*
	* @return list of cell names
	*/
	private Set<String> getReservedCells() {
	String[] list = reserved.list();
	return list != null ? ImmutableSet.copyOf(list) : ImmutableSet.of();
	}

	/**
	* 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) {
	checkNotNull(userName, USER_NOT_NULL);
	for (String cellName : getReservedCells()) {
	Reservation reservation = currentCellReservation(cellName);
	if (reservation != null && userName.equals(reservation.userName)) {
	return reservation;
	}
	}
	return null;
	}

	/**
	* Returns the name of the user who reserved the given cell.
	*
	* @param cellName cell name
	* @return cell reservation record or null if cell is not reserved
	*/
	Reservation currentCellReservation(String cellName) {
	checkNotNull(cellName, CELL_NOT_NULL);
	File cellFile = new File(reserved, cellName);
	if (!cellFile.exists()) {
	return null;
	}
	try (InputStream stream = new FileInputStream(cellFile)) {
	return new Reservation(new String(ByteStreams.toByteArray(stream), "UTF-8"));
	} catch (IOException e) {
	throw new IllegalStateException("Unable to get current user for cell " + cellName, e);
	}
	}

	/**
	* Reserves a cell for the specified user and their public access key.
	*
	* @param userName user name
	* @param sshKey user ssh public key
	* @param minutes number of minutes for reservation
	* @return reserved cell definition
	*/
	synchronized String borrowCell(String userName, String sshKey, int minutes) {
	checkNotNull(userName, USER_NOT_NULL);
	checkNotNull(sshKey, KEY_NOT_NULL);
	checkArgument(minutes < MAX_MINUTES, "Number of minutes must be less than %d", MAX_MINUTES);
	long now = System.currentTimeMillis();
	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);
	} 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);
	}

	reserveCell(reservation);
	createCell(reservation, sshKey);
	log(userName, reservation.cellName, "borrowed for " + reservation.duration + " minutes");
	return getCellDefinition(reservation.cellName);
	}

	/**
	* Returns the specified cell for the specified user and their public access key.
	*
	* @param userName user name
	*/
	synchronized void returnCell(String userName) {
	checkNotNull(userName, USER_NOT_NULL);
	Reservation reservation = currentUserReservation(userName);
	checkState(reservation != null, "User %s has no cell reservations", userName);

	unreserveCell(reservation);
	destroyCell(reservation);
	log(userName, reservation.cellName, "returned");
	}

	/**
	* Reserves the specified cell for the user the source file and writes the
	* specified content to the target file.
	*
	* @param reservation cell reservation record
	*/
	private void reserveCell(Reservation reservation) {
	File cellFile = new File(reserved, reservation.cellName);
	try (FileOutputStream stream = new FileOutputStream(cellFile)) {
	stream.write(reservation.encode().getBytes(UTF_8));
	} catch (IOException e) {
	throw new IllegalStateException("Unable to reserve cell " + reservation.cellName, e);
	}
	}

	/**
	* Returns the cell definition of the specified cell.
	*
	* @param cellName cell name
	* @return cell definition
	*/
	private String getCellDefinition(String cellName) {
	CellInfo cellInfo = getCellInfo(cellName);
	return exec(String.format("ssh %s warden/bin/cell-def %s",
	cellInfo.hostName, cellInfo.cellName));
	}

	/**
	* Cancels the specified reservation.
	*
	* @param reservation reservation record
	*/
	private void unreserveCell(Reservation reservation) {
	checkState(new File(reserved, reservation.cellName).delete(),
	"Unable to return cell %s", reservation.cellName);
	}

	/**
	* Creates the cell for the specified user SSH key.
	*
	* @param reservation cell reservation
	* @param sshKey ssh key
	*/
	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",
	cellInfo.hostName, cellInfo.cellName,
	cellInfo.ipPrefix, sshKey);
	exec(cmd);
	}

	/**
	* Destroys the specified cell.
	*
	* @param reservation reservation record
	*/
	private void destroyCell(Reservation reservation) {
	CellInfo cellInfo = getCellInfo(reservation.cellName);
	exec(String.format("ssh %s warden/bin/destroy-cell %s",
	cellInfo.hostName, cellInfo.cellName));
	}

	/**
	* Reads the information about the specified cell.
	*
	* @param cellName cell name
	* @return cell information
	*/
	private CellInfo getCellInfo(String cellName) {
	File cellFile = new File(supported, cellName);
	try (InputStream stream = new FileInputStream(cellFile)) {
	String[] fields = new String(ByteStreams.toByteArray(stream), UTF_8).split(" ");
	return new CellInfo(cellName, fields[0], fields[1]);
	} catch (IOException e) {
	throw new IllegalStateException("Unable to definition for cell " + cellName, e);
	}
	}

	// Executes the specified command.
	private String exec(String command) {
	try {
	Process process = Runtime.getRuntime().exec(command);
	String output = new String(ByteStreams.toByteArray(process.getInputStream()), UTF_8);
	process.waitFor(TIMEOUT, TimeUnit.SECONDS);
	return process.exitValue() == 0 ? output : null;
	} catch (Exception e) {
	throw new IllegalStateException("Unable to execute " + command);
	}
	}

	// Creates an audit log entry.
	private void log(String userName, String cellName, 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.flush();
	} catch (IOException e) {
	throw new IllegalStateException("Unable to log reservation action", e);
	}
	}

	// Carrier of cell information
	private final class CellInfo {
	final String cellName;
	final String hostName;
	final String ipPrefix;

	private CellInfo(String cellName, String hostName, String ipPrefix) {
	this.cellName = cellName;
	this.hostName = hostName;
	this.ipPrefix = ipPrefix;
	}
	}

	// Task for re-possessing overdue cells
	private final class Reposessor extends TimerTask {
	@Override
	public void run() {
	long now = System.currentTimeMillis();
	for (String cellName : getReservedCells()) {
	Reservation reservation = currentCellReservation(cellName);
	if (reservation != null &&
	(reservation.time + reservation.duration * MINUTE) < now) {
	try {
	returnCell(reservation.userName);
	} catch (Exception e) {
	e.printStackTrace();
	}
	}
	}
	}
	}
	}