package models.com.mc.workers;
import akka.actor.ActorRef;
import akka.actor.Cancellable;
import akka.actor.Props;
import akka.actor.UntypedActor;
import akka.contrib.pattern.DistributedPubSubExtension;
import akka.contrib.pattern.DistributedPubSubMediator;
import akka.contrib.pattern.DistributedPubSubMediator.Put;
import akka.event.Logging;
import akka.event.LoggingAdapter;
import scala.concurrent.duration.Deadline;
import scala.concurrent.duration.FiniteDuration;
import models.com.mc.workers.MasterWorkerProtocol.*;
import java.io.Serializable;
import java.util.*;
public class Master extends UntypedActor {
public static String ResultsTopic = "results";
public static Props props(FiniteDuration workTimeout) {
return Props.create(Master.class, workTimeout);
}
private final FiniteDuration workTimeout;
private final ActorRef mediator = DistributedPubSubExtension.get(getContext().system()).mediator();
private final LoggingAdapter log = Logging.getLogger(getContext().system(), this);
private final Cancellable cleanupTask;
private HashMap<String, WorkerState> workers = new HashMap<String, WorkerState>();
private Queue<Work> pendingWork = new LinkedList<Work>();
private Set<String> workIds = new LinkedHashSet<String>();
public Master(FiniteDuration workTimeout) {
this.workTimeout = workTimeout;
mediator.tell(new Put(getSelf()), getSelf());
this.cleanupTask = getContext().system().scheduler().schedule(
workTimeout.div(2), workTimeout.div(2), getSelf(), CleanupTick, getContext().dispatcher(), getSelf());
}
@Override
public void postStop() {
cleanupTask.cancel();
}
@Override
public void onReceive(Object message) {
if (message instanceof RegisterWorker) {
RegisterWorker msg =
(RegisterWorker) message;
String workerId = msg.workerId;
if (workers.containsKey(workerId)) {
workers.put(workerId, workers.get(workerId).copyWithRef(getSender()));
} else {
log.debug("Worker registered: {}", workerId);
workers.put(workerId, new WorkerState(getSender(), Idle.instance));
if (!pendingWork.isEmpty())
getSender().tell(WorkIsReady.getInstance(), getSelf());
}
}
else if (message instanceof WorkerRequestsWork) {
WorkerRequestsWork msg = (WorkerRequestsWork) message;
String workerId = msg.workerId;
if (!pendingWork.isEmpty()) {
WorkerState state = workers.get(workerId);
if (state != null && state.status.isIdle()) {
Work work = pendingWork.remove();
log.debug("Giving worker {} some work {}", workerId, work.getJob());
// TODO store in Eventsourced
getSender().tell(work, getSelf());
workers.put(workerId, state.copyWithStatus(new Busy(work, workTimeout.fromNow())));
}
}
}
else if (message instanceof WorkIsDone) {
WorkIsDone msg = (WorkIsDone) message;
String workerId = msg.workerId;
String workId = msg.workId;
WorkerState state = workers.get(workerId);
if (state != null && state.status.isBusy() && state.status.getWork().getWorkId().equals(workId)) {
Work work = state.status.getWork();
Object result = msg.result;
log.debug("Work is done: {} => {} by worker {}", work, result, workerId);
// TODO store in Eventsourced
workers.put(workerId, state.copyWithStatus(Idle.instance));
mediator.tell(new DistributedPubSubMediator.Publish(ResultsTopic,
new WorkResult(workId, result)), getSelf());
getSender().tell(new Ack(workId), getSelf());
} else {
if (workIds.contains(workId)) {
// previous Ack was lost, confirm again that this is done
getSender().tell(new Ack(workId), getSelf());
}
}
}
else if (message instanceof WorkFailed) {
WorkFailed msg = (WorkFailed) message;
String workerId = msg.workerId;
String workId = msg.workId;
WorkerState state = workers.get(workerId);
if (state != null && state.status.isBusy() && state.status.getWork().getWorkId().equals(workId)) {
log.info("Work failed: {}", state.status.getWork());
// TODO store in Eventsourced
workers.put(workerId, state.copyWithStatus(Idle.instance));
pendingWork.add(state.status.getWork());
notifyWorkers();
}
}
else if (message instanceof Work) {
Work work = (Work) message;
// idempotent
if (workIds.contains(work.getWorkId())) {
getSender().tell(new Ack(work.getWorkId()), getSelf());
} else {
log.debug("Accepted work: {}", work);
// TODO store in Eventsourced
pendingWork.add(work);
workIds.add(work.getWorkId());
getSender().tell(new Ack(work.getWorkId()), getSelf());
notifyWorkers();
}
}
else if (message == CleanupTick) {
Iterator<Map.Entry<String, WorkerState>> iterator =
workers.entrySet().iterator();
while (iterator.hasNext()) {
Map.Entry<String, WorkerState> entry = iterator.next();
String workerId = entry.getKey();
WorkerState state = entry.getValue();
if (state.status.isBusy()) {
if (state.status.getDeadLine().isOverdue()) {
Work work = state.status.getWork();
log.info("Work timed out: {}", work);
// TODO store in Eventsourced
iterator.remove();
pendingWork.add(work);
notifyWorkers();
}
}
}
}
else {
unhandled(message);
}
}
private void notifyWorkers() {
if (!pendingWork.isEmpty()) {
// could pick a few random instead of all
for (WorkerState state: workers.values()) {
if (state.status.isIdle())
state.ref.tell(WorkIsReady.getInstance(), getSelf());
}
}
}
private static abstract class WorkerStatus {
protected abstract boolean isIdle();
private boolean isBusy() {
return !isIdle();
};
protected abstract Work getWork();
protected abstract Deadline getDeadLine();
}
private static final class Idle extends WorkerStatus {
private static final Idle instance = new Idle();
public static Idle getInstance() {
return instance;
}
@Override
protected boolean isIdle() {
return true;
}
@Override
protected Work getWork() {
throw new IllegalAccessError();
}
@Override
protected Deadline getDeadLine() {
throw new IllegalAccessError();
}
@Override
public String toString() {
return "Idle";
}
}
private static final class Busy extends WorkerStatus {
private final Work work;
private final Deadline deadline;
private Busy(Work work, Deadline deadline) {
this.work = work;
this.deadline = deadline;
}
@Override
protected boolean isIdle() {
return false;
}
@Override
protected Work getWork() {
return work;
}
@Override
protected Deadline getDeadLine() {
return deadline;
}
@Override
public String toString() {
return "Busy{" +
"work=" + work +
", deadline=" + deadline +
'}';
}
}
private static final class WorkerState {
public final ActorRef ref;
public final WorkerStatus status;
private WorkerState(ActorRef ref, WorkerStatus status) {
this.ref = ref;
this.status = status;
}
private WorkerState copyWithRef(ActorRef ref) {
return new WorkerState(ref, this.status);
}
private WorkerState copyWithStatus(WorkerStatus status) {
return new WorkerState(this.ref, status);
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
WorkerState that = (WorkerState) o;
if (!ref.equals(that.ref)) return false;
if (!status.equals(that.status)) return false;
return true;
}
@Override
public int hashCode() {
int result = ref.hashCode();
result = 31 * result + status.hashCode();
return result;
}
@Override
public String toString() {
return "WorkerState{" +
"ref=" + ref +
", status=" + status +
'}';
}
}
private static final Object CleanupTick = new Object() {
@Override
public String toString() {
return "CleanupTick";
}
};
// public static final class Work implements Serializable {
// public final String workId;
// public final Object job;
//
// public Work(String workId, Object job) {
// this.workId = workId;
// this.job = job;
// }
//
// @Override
// public String toString() {
// return "Work{" +
// "workId='" + workId + '\'' +
// ", job=" + job +
// '}';
// }
// }
public static final class WorkResult implements Serializable {
public final String workId;
public final Object result;
public WorkResult(String workId, Object result) {
this.workId = workId;
this.result = result;
}
@Override
public String toString() {
return "WorkResult{" +
"workId='" + workId + '\'' +
", result=" + result +
'}';
}
public String getWorkId() {
return workId;
}
public Object getResult() {
return result;
}
}
public static final class Ack implements Serializable {
final String workId;
public Ack(String workId) {
this.workId = workId;
}
@Override
public String toString() {
return "Ack{" +
"workId='" + workId + '\'' +
'}';
}
}
// TODO cleanup old workers
// TODO cleanup old workIds
}