com.github.euler.core.PooledExecution Maven / Gradle / Ivy

Go to download
package com.github.euler.core;

import java.net.URI;
import java.util.ArrayList;
import java.util.List;

import akka.actor.typed.ActorRef;
import akka.actor.typed.Behavior;
import akka.actor.typed.SupervisorStrategy;
import akka.actor.typed.javadsl.AbstractBehavior;
import akka.actor.typed.javadsl.ActorContext;
import akka.actor.typed.javadsl.Behaviors;
import akka.actor.typed.javadsl.Receive;
import akka.actor.typed.javadsl.ReceiveBuilder;

public class PooledExecution extends AbstractBehavior {

    public static Behavior create(int size, Task task) {
        return Behaviors.setup((ctx) -> new PooledExecution(ctx, size, task));
    }

    private List> pool;
    private int position;
    private int size;
    private Task task;

    public PooledExecution(ActorContext context, int size, Task task) {
        super(context);
        pool = new ArrayList>(size);
        this.size = size;
        this.task = task;
        position = 0;
    }

    @Override
    public Receive createReceive() {
        ReceiveBuilder builder = newReceiveBuilder();
        builder.onMessage(JobTaskToProcess.class, this::onJobTaskToProcess);
        builder.onMessage(Flush.class, this::onFlush);
        return builder.build();
    }

    public Behavior onFlush(Flush msg) {
        for (ActorRef taskRef : this.pool) {
            taskRef.tell(msg);
        }
        return Behaviors.same();
    }

    private Behavior superviseTaskBehavior(Task t) {
        Behavior behavior = Behaviors.supervise(MiddleManagement.create(t)).onFailure(SupervisorStrategy.restart());
        return behavior;
    }

    private Behavior onJobTaskToProcess(JobTaskToProcess msg) {
        ActorRef taskRef = getNextTaskRef();
        taskRef.tell(msg);
        return Behaviors.same();
    }

    private ActorRef getNextTaskRef() {
        if (pool.size() - 1 < position) {
            pool.add(getContext().spawn(superviseTaskBehavior(task), task.name() + "-" + position));
        }
        ActorRef taskRef = pool.get(position);
        if (position + 1 == size) {
            position = 0;
        } else {
            position++;
        }
        return taskRef;
    }

    private static class MiddleManagement extends AbstractBehavior {

        public static Behavior create(Task task) {
            return Behaviors.setup((ctx) -> new MiddleManagement(ctx, task));
        }

        private final Task task;
        private ActorRef taskRef;

        public MiddleManagement(ActorContext context, Task task) {
            super(context);
            this.task = task;
        }

        @Override
        public Receive createReceive() {
            ReceiveBuilder builder = newReceiveBuilder();
            builder.onMessage(JobTaskToProcess.class, this::onJobTaskToProcess);
            builder.onMessage(InternalJobTaskFailed.class, this::onInternalJobTaskFailed);
            builder.onMessage(Flush.class, this::onFlush);
            builder.onAnyMessage((m) -> {
                System.out.println(m);
                return Behaviors.same();
            });
            return builder.build();
        }

        public Behavior onFlush(Flush msg) {
            if (taskRef != null) {
                taskRef.tell(msg);
            }
            return Behaviors.same();
        }

        private Behavior onJobTaskToProcess(JobTaskToProcess msg) {
            ActorRef ref = getTaskRef(msg);
            ref.tell(msg);
            return Behaviors.same();
        }

        private ActorRef getTaskRef(JobTaskToProcess msg) {
            if (taskRef == null) {
                taskRef = getContext().spawn(superviseTaskBehavior(this.task), "middle-management");
                getContext().watchWith(taskRef, new InternalJobTaskFailed(msg));
            }
            return taskRef;
        }

        private Behavior superviseTaskBehavior(Task t) {
            Behavior behavior = Behaviors.supervise(t.behavior()).onFailure(SupervisorStrategy.stop());
            return behavior;
        }

        private Behavior onInternalJobTaskFailed(InternalJobTaskFailed msg) {
            taskRef = null;
            msg.replyTo.tell(new JobTaskFailed(msg.uri, msg.itemURI));
            return Behaviors.same();
        }

    }

    private static class InternalJobTaskFailed implements TaskCommand {

        public final URI uri;
        public final URI itemURI;
        public final ActorRef replyTo;

        public InternalJobTaskFailed(JobTaskToProcess msg) {
            this.uri = msg.uri;
            this.itemURI = msg.itemURI;
            this.replyTo = msg.replyTo;
        }

    }

}