public interface BackingStore {

	/**
	 * Persist the jobs to the list of scheduled jobs
	 * 
	 * @param scheduledJobs
	 *            is the list of {@linkplain Job}s assigned to the agent
	 */
	public void addScheduledJobs(List<Job> scheduledJobs);

	/**
	 * Add a job to the unscheduled jobs queue
	 */
	public void addUnscheduledJob(Job job);

	/**
	 * Clean up complete and optionally, failed jobs from the backing store
	 * 
	 * @param includeFailedJobs
	 *            determines whether failed jobs will also be cleaned up
	 * @param expirationAge
	 *            is the duration that must have elapsed before the job is
	 *            eligible for cleanup, e.g. "24 hours" means that the job won't
	 *            be cleaned up until at least 24 hours have elapsed since it
	 *            completed.
	 * @return the number of jobs that were cleaned up
	 */
	public int cleanupJobs(boolean includeFailedJobs, Duration expirationAge);

	/**
	 * The backing store makes a note of jobs that complete successfully. Get
	 * the list of these.
	 * 
	 * @return a list of complete {@linkplain Job}
	 */
	public List<Job> getCompletedJobs();

	/**
	 * The backing store makes a note of jobs that fail. Get the list of these.
	 * 
	 * @return a list of failed {@linkplain Job}
	 */
	public List<Job> getFailedJobs();

	/**
	 * Fetch the next monotonically-increasing job identifier. NOTE: This action
	 * *must* increment the job identifier in the persistent store, as well as
	 * return it!
	 */
	public long getNextJobID();

	/**
	 * Persist the details of a failed job
	 * 
	 * @param job
	 *            is the job that has failed
	 */
	public void persistJobFailure(Job job);

	/**
	 * Persist the progress made against a job
	 * 
	 * @param job
	 *            contains a record of progress made by a worker
	 */
	public void persistJobProgress(Job job);

	/**
	 * Remove a job from the list of completed jobs
	 */
	public void removeCompletedJobById(long jobId);

	/**
	 * Remove a job from the list of failed jobs
	 */
	public void removeFailedJobById(long jobId);

	/**
	 * Remove a job from the list of scheduled jobs by ID
	 * 
	 * @param jobId
	 *            the ID of the {@linkplain Job} to remove
	 */
	public void removeScheduledJobById(long jobId);

	/**
	 * Remove a job from the unscheduled jobs queue by ID
	 */
	public void removeUnscheduledJobById(long jobId);

	/**
	 * Restore the unscheduled jobs queue from both scheduled and unscheduled
	 * jobs, as we assume that all agents are dead.
	 * 
	 * @return a list of {@linkplain Job}
	 */
	public List<Job> restoreJobs();

} public abstract class AbstractWorkRequest implements Serializable {

	private static final long serialVersionUID = -3802453222882202468L;

	private final ActorRef agent;
	private final Set<String> workerTypes;

	public AbstractWorkRequest(ActorRef agent, Set<String> workerTypes) {
		this.agent = agent;
		this.workerTypes = workerTypes;
	}

	public ActorRef getAgent() {
		return agent;
	}

	public Set<String> getWorkerTypes() {
		return workerTypes;
	}
} public class EnqueueJob implements Serializable {

	private static final long serialVersionUID = -1624154938276876491L;

	protected String workerType;
	protected Map<String, String> params;

	public EnqueueJob(String workerType) {
		this(workerType, Collections.<String, String> emptyMap());
	}

	public EnqueueJob() {
	}

	public EnqueueJob(String workerType, Map<String, String> jobParams) {
		this.workerType = workerType;
		this.params = jobParams;
	}

	public Map<String, String> getParams() {
		return params;
	}

	public String getWorkerType() {
		return workerType;
	}

	@Override
	public String toString() {
		return String.format("Enqueue job for %s", workerType);
	}

} public class Job implements Serializable, Cloneable {

	public enum State {
		COMPLETE {
			public String toString() {
				return "complete";
			}
		},
		FAILED {
			public String toString() {
				return "failed";
			}
		},
		QUEUED {
			public String toString() {
				return "queued";
			}
		},
		RUNNING {
			public String toString() {
				return "running";
			}
		},
		SCHEDULED {
			public String toString() {
				return "scheduled";
			}
		},
		DELETED {
			public String toString() {
				return "deleted";
			}
		}
	}

	private static final long serialVersionUID = -2375588116753600617L;

	private DateTime enqueuedAt = new DateTime(DateTimeUtils.currentTimeMillis());
	private Map<String, String> params = Maps.newHashMap();
	private double progress = 0.0;
	private boolean rerun = false;
	private State state = State.QUEUED;

	private DateTime startedAt;
	private DateTime completedAt;
	private String errorMessage;
	private long id;
	private String workerType;

	/**
	 * This default constructor required for Jackson JSON serialization
	 */
	public Job() {
	}

	/**
	 * Create a new job. Use this constructor when you are creating a job for the first time. The
	 * job enqueued time will be set to now and the job state will be set to queued.
	 * 
	 * @param id is the unique identifier for this job
	 * 
	 * @param workerType determines which type of worker is capable of completing this job
	 */
	public Job(long id, String workerType) {
		this();
		this.id = id;
		this.workerType = workerType;
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see java.lang.Object#clone()
	 */
	public Object clone() {
		Job clone = new Job(this.getId(), this.getWorkerType());
		clone.setEnqueuedAt(this.getEnqueuedAt());
		clone.setStartedAt(this.getStartedAt());
		clone.setCompletedAt(this.getCompletedAt());
		clone.setErrorMessage(this.getErrorMessage());
		clone.setProgress(this.getProgress());
		clone.setState(this.getState());
		clone.setRerun(this.isRerun());

		for (String key : this.getParams().keySet()) {
			clone.getParams().put(key, this.params.get(key));
		}

		return clone;
	}

	// Create a clone of this job with publicly visible properties.
	public Job clonePublicView() {
		Job result = (Job) clone();
		for (Iterator<String> it = result.getParams().keySet().iterator(); it.hasNext(); ) {
			String key = it.next();
			if (key.charAt(0) == '$') {
				result.getParams().put(key, "*hidden*");
			}
		}
		return result;
	}

	public DateTime getCompletedAt() {
		return completedAt;
	}

	public DateTime getEnqueuedAt() {
		return enqueuedAt;
	}

	public String getErrorMessage() {
		return errorMessage;
	}

	public long getId() {
		return id;
	}

	public Map<String, String> getParams() {
		return params;
	}

	public double getProgress() {
		return progress;
	}

	public DateTime getStartedAt() {
		return startedAt;
	}

	public State getState() {
		return state;
	}

	public String getWorkerType() {
		return workerType;
	}

	public boolean isRerun() {
		return rerun;
	}

	public void setCompletedAt(DateTime completedAt) {
		this.completedAt = completedAt;
	}

	public void setEnqueuedAt(DateTime enqueuedAt) {
		this.enqueuedAt = enqueuedAt;
	}

	public void setErrorMessage(String errorMessage) {
		this.errorMessage = errorMessage;
	}

	public void setParams(Map<String, String> params) {
		this.params = params;
	}

	public void setProgress(double progress) {
		this.progress = progress;
	}

	public void setRerun(boolean rerun) {
		this.rerun = rerun;
	}

	public void setStartedAt(DateTime startedAt) {
		this.startedAt = startedAt;
	}

	public void setState(State state) {
		this.state = state;
	}

	// Don't print parameters that start with $
	public Map<String,String> getParams(boolean includeSecrets) {
		if (includeSecrets) {
			return params;
		}
		Map<String,String> result = new HashMap<>();
		for (String key : params.keySet()) {
			if (key.charAt(0) != '$') {
				result.put(key, params.get(key));
			} else {
				result.put(key, "*hidden*");
			}
		}
		return result;
	}

	@Override
	public String toString() {
		return String
				.format("Job %s (state=%s, enqueuedAt=%s, startedAt=%s, completedAt=%s, workerType=%s, re-run=%s, progress=%s params=%s)",
						id, state, getEnqueuedAt(), getStartedAt(), getCompletedAt(), workerType,
						rerun, progress, getParams(false));
	}

	@Override
	public boolean equals(Object obj) {
		if (obj == null) {
			return false;
		}
		if (getClass() != obj.getClass()) {
			return false;
		}
		final Job otherJob = (Job) obj;
		return Objects.equal(enqueuedAt, otherJob.getEnqueuedAt())
				&& Objects.equal(startedAt, otherJob.getStartedAt())
				&& Objects.equal(completedAt, otherJob.getCompletedAt())
				&& Objects.equal(errorMessage, otherJob.getErrorMessage())
				&& Objects.equal(id, otherJob.getId())
				&& Objects.equal(params, otherJob.getParams())
				&& Objects.equal(progress, otherJob.getProgress())
				&& Objects.equal(rerun, otherJob.isRerun())
				&& Objects.equal(state, otherJob.getState())
				&& Objects.equal(workerType, otherJob.getWorkerType());
	}

	@Override
	public int hashCode() {
		return Objects.hashCode(enqueuedAt, startedAt, completedAt, errorMessage, id, params,
				progress, rerun, state, workerType);
	}
} public class JobSummary implements Serializable {

	private static final long serialVersionUID = 8252819036997216081L;

	private List<Job> jobs = new ArrayList<>();

	/**
	 * empty constructor required for JSON mapping
	 */
	public JobSummary() {
	}

	public JobSummary(Collection<Job> jobs) {
		super();
		for (Job job : jobs) {
			this.jobs.add((Job) job.clonePublicView());
		}
	}

	public List<Job> getJobs() {
		return jobs;
	}

	@Override
	public String toString() {
		StringBuilder builder = new StringBuilder();
		Iterator<Job> iterator = jobs.iterator();
		while (iterator.hasNext()) {
			Job job = iterator.next();
			builder.append(job.toString());
			if (iterator.hasNext())
				builder.append("|");
		}
		return builder.toString();
	}
} public enum SimpleMessage {
	AGENT_HEARTBEAT, AGENT_REGISTERED, AGENT_SHUTDOWN, AGENT_DEAD, LIST_AGENTS, CHECK_AGENTS, SCHEDULE_JOBS, JOB_SUMMARY, BROADCAST_JOBS
} public class WorkResponse implements Serializable {

	private static final long serialVersionUID = 2445495404623836556L;

	private final List<Job> jobs = new ArrayList<Job>();

	public WorkResponse() {
	}

	public WorkResponse(Job job) {
		this.jobs.add(job);
	}

	public WorkResponse(List<Job> jobs) {
		this.jobs.addAll(jobs);
	}

	public List<Job> getJobs() {
		return jobs;
	}

	@Override
	public String toString() {
		if (jobs.isEmpty())
			return "No jobs";
		else
			return jobs.size() + " jobs";
	}
} public class CapacityScheduler extends AbstractScheduler<CapacityWorkRequest> {

	private class WorkerUniquenessParameters extends HashMap<String, Set<Map<String, String>>> {
		private static final long serialVersionUID = -2463700325426716368L;
	}

	private final Config config;
	private final Map<String, Set<String>> workerTypesToUniqueParameters;

	@SuppressWarnings("unchecked")
	public CapacityScheduler(Class<? extends BackingStore> backingStore) {
		super(backingStore);
		config = getContext().system().settings().config()
				.getConfig("oncue.scheduler.capacity-scheduler");

		List<? extends ConfigObject> uniquenessConstraints = config
				.getObjectList("uniqueness-constraints");
		workerTypesToUniqueParameters = Maps.newHashMap();
		for (ConfigObject value : uniquenessConstraints) {
			Set<String> keys = Sets.newHashSet();
			if (value.containsKey("uniqueness-keys")) {
				for (String key : (List<String>) value.get("uniqueness-keys").unwrapped()) {
					keys.add(key);
				}
			}
			workerTypesToUniqueParameters.put((String) value.get("worker-type").unwrapped(), keys);
		}
	}

	protected Comparator<Job> getComparator() {
		return new PriorityJobComparator();
	}

	@Override
	protected void scheduleJobs(CapacityWorkRequest workRequest) {
		WorkerUniquenessParameters runningUniquenessConstrainedJobTypes = getScheduledUniquenessConstrainedParams();
		List<Job> jobs = new ArrayList<>();
		int allocatedMemory = 0;

		Iterator<Job> iterator = unscheduledJobs.iterator();

		while (iterator.hasNext()) {
			Job job = iterator.next();
			if (workRequest.getWorkerTypes().contains(job.getWorkerType())) {
				int requiredMemory = getRequiredMemory(job);
				if (requiredMemory + allocatedMemory <= workRequest.getAvailableMemory()) {
					if (workerTypesToUniqueParameters.containsKey(job.getWorkerType())) {
						if (runningUniquenessConstrainedJobTypes.containsKey(job.getWorkerType())) {
							boolean blockedByRunningUniquenessConstrainedJob = findRunningUniquenessConstraintedJobConflicts(
									runningUniquenessConstrainedJobTypes, job);
							if (!blockedByRunningUniquenessConstrainedJob) {
								addJob(job, runningUniquenessConstrainedJobTypes);
								jobs.add(job);
								allocatedMemory += requiredMemory;
							}
						} else {
							addJob(job, runningUniquenessConstrainedJobTypes);
							jobs.add(job);
							allocatedMemory += requiredMemory;
						}
					} else {
						jobs.add(job);
						allocatedMemory += requiredMemory;
					}
				}
			}
		}

		log.debug("Scheduling {} job(s) with a total memory of {}", jobs.size(), allocatedMemory);

		// Create the schedule
		Schedule schedule = new Schedule();
		schedule.setJobs(getSender(), jobs);

		// Dispatch the schedule
		dispatchJobs(schedule);
	}

	/**
	 * Check if the uniqueness parameters for the given job conflict with any existing running job's
	 * uniqueness parameters. I.e., given a map of worker type -> (set of (map of uniqueness key ->
	 * parameter value for that uniqueness key)), ensure that for each defined uniqueness key for
	 * the job's worker type there are no entries in the running job's uniquness key/value pairs
	 * that are identical.
	 * 
	 * @param runningUniquenessConstrainedJobTypes Information on the currently running jobs and
	 * their parameter values for all uniqueness constrained parameters
	 * @param job The job
	 * @return
	 */
	private boolean findRunningUniquenessConstraintedJobConflicts(
			WorkerUniquenessParameters runningUniquenessConstrainedJobTypes, Job job) {
		boolean foundIdentical = false;
		for (Map<String, String> parameterKeyValues : runningUniquenessConstrainedJobTypes.get(job
				.getWorkerType())) {
			boolean identical = true;
			for (String parameter : parameterKeyValues.keySet()) {
				if (!job.getParams().get(parameter).equals(parameterKeyValues.get(parameter))) {
					identical = false;
					break;
				}
			}
			if (identical) {
				foundIdentical = true;
				break;
			}
		}
		return foundIdentical;
	}

	@Override
	protected void augmentJob(Job job) {
		ensureRequiredMemory(job);
	}

	/**
	 * Ensure that a job has a "memory" parameter. If it exists, leave it as it is. If it does not,
	 * attempt to populate it from the configuration. This will crash if both the job does not have
	 * the parameter and the configuration does not define a default value for this type of job.
	 * 
	 * @param job The job
	 */
	private void ensureRequiredMemory(Job job) {
		Map<String, String> params = job.getParams();
		if (!params.containsKey("memory")) {
			Config config = getContext().system().settings().config();
			params.put(
					"memory",
					String.valueOf(config.getConfig("oncue.scheduler.capacity-scheduler").getInt(
							"default-requirements." + job.getWorkerType() + ".memory")));
		}
	}

	/**
	 * Get an integer representation of the memory key for a job. Caller should ensure jobs has a
	 * memory parameter.
	 * 
	 * @param job The job
	 * @return
	 */
	private int getRequiredMemory(Job job) {
		return Integer.parseInt(job.getParams().get("memory"));
	}

	/**
	 * Return a list of all scheduled jobs that have uniqueness constraints and the parameter
	 * key/values for keys defined in those uniqueness constraints
	 * 
	 * @return
	 */
	private WorkerUniquenessParameters getScheduledUniquenessConstrainedParams() {
		List<Job> scheduledJobs = getScheduledJobs();
		WorkerUniquenessParameters scheduledUniquenessConstrainedParams = new WorkerUniquenessParameters();
		for (Job job : scheduledJobs) {
			if (workerTypesToUniqueParameters.containsKey(job.getWorkerType())) {
				addJob(job, scheduledUniquenessConstrainedParams);
			}
		}
		return scheduledUniquenessConstrainedParams;
	}

	/**
	 * Add a job and it's uniqueness parameters to a list of active uniqueness constrained jobs.
	 * 
	 * @param job
	 * @param runningUniquenessConstrainedJobTypes
	 */
	private void addJob(Job job, WorkerUniquenessParameters runningUniquenessConstrainedJobTypes) {
		Set<String> uniqueParams = workerTypesToUniqueParameters.get(job.getWorkerType());
		Map<String, String> uniqueParamValues = Maps.newHashMap();
		for (String key : job.getParams().keySet()) {
			if (uniqueParams.contains(key)) {
				uniqueParamValues.put(key, job.getParams().get(key));
			}
		}
		if (!runningUniquenessConstrainedJobTypes.containsKey(job.getWorkerType())) {
			Set<Map<String, String>> params = Sets.newHashSet();
			runningUniquenessConstrainedJobTypes.put(job.getWorkerType(), params);
		}

		runningUniquenessConstrainedJobTypes.get(job.getWorkerType()).add(uniqueParamValues);
	}

} public abstract class ActorSystemTest {

	protected Config config;

	protected LoggingAdapter log;

	protected Settings settings;

	protected ActorSystem system;

	// By default ActorSystemTest children will wait for all actors to finish before shutting down
	// the test system. If you want to shut down the test system while e.g. workers are still
	// running, set this to false.
	protected boolean waitForRunningJobs = true;

	// Required for naming agents uniquely
	private int agentCount = 0;

	/**
	 * Construct a Agent without a probe
	 */
	public ActorRef createAgent(ActorSystem system, final Set<String> workers) {
		return createAgent(system, workers, null);
	}

	/**
	 * Create an agent component, with a set of workers and an optional probe
	 * 
	 * @param probe can be null
	 */
	@SuppressWarnings("serial")
	public ActorRef createAgent(ActorSystem system, final Set<String> workers,
			final ActorRef probe) {
		agentCount++;
		return system.actorOf(new Props(new UntypedActorFactory() {

			@Override
			public Actor create() throws Exception {
				AbstractAgent agent = (AbstractAgent) Class.forName(settings.AGENT_CLASS)
						.getConstructor(Set.class).newInstance(workers);
				if (probe != null)
					agent.injectProbe(probe);
				return agent;
			}
		}), settings.AGENT_NAME + agentCount);
	}

	/**
	 * Construct a Scheduler without a probe
	 */
	public ActorRef createScheduler(ActorSystem system) {
		return createScheduler(system, null);
	}

	/**
	 * Create a scheduler component, with an optional probe
	 * 
	 * @param probe can be null
	 */
	@SuppressWarnings("serial")
	public ActorRef createScheduler(ActorSystem system, final ActorRef probe) {
		return system.actorOf(new Props(new UntypedActorFactory() {

			@Override
			public Actor create() throws Exception {
				Class<?> schedulerClass = Class.forName(settings.SCHEDULER_CLASS);
				Class<?> backingStoreClass = null;
				if (settings.SCHEDULER_BACKING_STORE_CLASS != null)
					backingStoreClass = Class.forName(settings.SCHEDULER_BACKING_STORE_CLASS);

				@SuppressWarnings("rawtypes")
				AbstractScheduler scheduler = (AbstractScheduler) schedulerClass
						.getConstructor(Class.class).newInstance(backingStoreClass);
				if (probe != null)
					scheduler.injectProbe(probe);
				return scheduler;
			}
		}), settings.SCHEDULER_NAME);
	}

	@Before
	public void startActorSystem() {
		/*
		 * Load configuration specific to this test and fall back to the reference configuration
		 */
		config = ConfigFactory.load();
		config = ConfigFactory.load(getClass().getSimpleName()).withFallback(config);

		system = ActorSystem.create("oncue-test", config);
		settings = SettingsProvider.SettingsProvider.get(system);
		log = Logging.getLogger(system, this);
	}

	@After
	public void stopActorSystem() throws Exception {
		system.shutdown();
		if (waitForRunningJobs) {
			while (!system.isTerminated()) {
				log.debug("Waiting for system to shut down...");
				Thread.sleep(500);
			}
		}
		log.debug("System shut down");
	}

	@Before
	@After
	public void cleanRedis() {
		try (RedisConnection redis = new RedisConnection()) {
			redis.flushDB();
		}
	}
} public class SimpleLoadTestWorker extends AbstractWorker {

	private static final int LOAD_FACTOR = 10000;

	@Override
	protected void doWork(Job job) {
		processJob();
	}

	private void processJob() {
		int count = 0;
		while (count < LOAD_FACTOR) {
			count++;

			/*
			 * Yield occasionally, so we don't starve the Agent of CPU and
			 * prevent its heart beat.
			 */
			if (count % 10 == 0)
				Thread.yield();
		}
	}

	@Override
	protected void redoWork(Job job) {
		processJob();
	}
} public class TestWorker extends AbstractWorker {

	@Override
	public void doWork(Job job) throws InterruptedException {
		processJob();
	}

	private void processJob() throws InterruptedException {
		double progress = 0.0;
		for (int i = 0; i < 3; i++) {
			progress += 0.25;
			Thread.sleep(500);
			reportProgress(progress);
		}
		Thread.sleep(500);
	}

	@Override
	protected void redoWork(Job job) throws Exception {
		processJob();
	}

} public class TestWorker2 extends AbstractWorker {

	@Override
	public void doWork(Job job) throws InterruptedException {
		processJob();
	}

	private void processJob() throws InterruptedException {
		double progress = 0.0;
		for (int i = 0; i < 3; i++) {
			progress += 0.25;
			Thread.sleep(500);
			reportProgress(progress);
		}
		Thread.sleep(500);
	}

	@Override
	protected void redoWork(Job job) throws Exception {
		processJob();
	}

}
import static junit.framework.Assert.assertEquals;
import static junit.framework.Assert.fail;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import oncue.backingstore.BackingStore;
import oncue.common.messages.AbstractWorkRequest;
import oncue.common.messages.EnqueueJob;
import oncue.common.messages.Job;
import oncue.common.messages.JobSummary;
import oncue.common.messages.SimpleMessages.SimpleMessage;
import oncue.common.messages.WorkResponse;
import oncue.scheduler.CapacityScheduler;
import oncue.tests.base.ActorSystemTest;
import oncue.tests.load.workers.SimpleLoadTestWorker;
import oncue.tests.workers.TestWorker;
import oncue.tests.workers.TestWorker2;
import org.junit.Test;
import scala.concurrent.duration.FiniteDuration;
import akka.actor.Actor;
import akka.actor.ActorRef;
import akka.actor.Props;
import akka.actor.UntypedActorFactory;
import akka.testkit.JavaTestKit;
import akka.testkit.TestActorRef;
import com.google.api.client.util.Maps;
import com.google.common.collect.Sets;
package oncue.tests.strategies;







public class CapacityStrategyTest extends ActorSystemTest {

	private static final String TEST_WORKER = TestWorker.class.getName();

	@Test
	public void doesNotScheduleJobsThatCannotBeHandledByTheAgent() {
		new JavaTestKit(system) {
			{
				// Create a scheduler
				ActorRef scheduler = createScheduler(system);

				// Enqueue jobs
				scheduler.tell(new EnqueueJob(TestWorker2.class.getName(), memory("2600")),