org.jtrim2.taskgraph.basic.WeakLeafsOfEndNodeRestrictingStrategy Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jtrim-task-graph Show documentation
JTrim Task Graph
There is a newer version: 2.0.7
package org.jtrim2.taskgraph.basic;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Function;
import org.jtrim2.concurrent.Tasks;
import org.jtrim2.taskgraph.TaskNodeKey;
import org.jtrim2.utils.ExceptionHelper;

final class WeakLeafsOfEndNodeRestrictingStrategy implements TaskExecutionRestrictionStrategyFactory {
    private final int maxRetainedLeafNodes;

    // This is simply here to allow creating deterministic tests.
    private Function>, Collection>> queueSorter;

    public WeakLeafsOfEndNodeRestrictingStrategy(int maxRetainedLeafNodes) {
        ExceptionHelper.checkArgumentInRange(maxRetainedLeafNodes, 1, Integer.MAX_VALUE, "maxRetainedLeafNodes");
        this.maxRetainedLeafNodes = maxRetainedLeafNodes;
        this.queueSorter = queue -> queue;
    }

    @Override
    public TaskExecutionRestrictionStrategy buildStrategy(
            DependencyDag> taskGraph,
            Iterable restrictableNodes) {

        StrategyImpl strategy = new StrategyImpl(
                maxRetainedLeafNodes,
                taskGraph,
                restrictableNodes,
                queueSorter);

        strategy.scheduleUnsafe();
        return strategy;
    }

    void setQueueSorter(Function>, Collection>> queueSorter) {
        Objects.requireNonNull(queueSorter, "queueSorter");
        this.queueSorter = queueSorter;
    }

    private static void selectLeafAndEndNodes(
            DependencyDag> graph,
            Iterable restrictableNodes,
            Map, Runnable> leafNodes,
            Collection> endNodes) {

        restrictableNodes.forEach((restrictableNode) -> {
            TaskNodeKey nodeKey = restrictableNode.getNodeKey();
            if (graph.getDependencyGraph().hasChildren(nodeKey)) {
                restrictableNode.release();
            } else {
                leafNodes.put(nodeKey, Tasks.runOnceTask(restrictableNode.getReleaseAction()));
            }

            if (!graph.getForwardGraph().hasChildren(nodeKey)) {
                endNodes.add(nodeKey);
            }
        });
    }

    private static  void addMissingEndNodes(DependencyDag taskGraph, Collection endNodes) {
        DirectedGraph forwardGraph = taskGraph.getForwardGraph();
        taskGraph.getDependencyGraph().getRawGraph().keySet().forEach((nodeKey) -> {
            if (!forwardGraph.hasChildren(nodeKey)) {
                endNodes.add(nodeKey);
            }
        });
    }

    private static final class StrategyImpl implements TaskExecutionRestrictionStrategy {
        private final int maxRetainedLeafNodes;

        private final Lock mainLock;
        private final Map, Runnable> leafNodes;
        private final Map, Set>> endNodesToLeafs;
        private final Map, Set>> leafsToEndNodes;

        private final Set> endNodeQueue;
        private final Set> computedEndNodes;
        private final Set> releasedNotComputedEndNodes;
        private final Set> retainingNotComputedEndNodes;
        private final Map, Set>> scheduledLeafNodes;

        public StrategyImpl(
                int maxRetainedLeafNodes,
                DependencyDag> taskGraph,
                Iterable restrictableNodes,
                Function>, Collection>> queueSorter) {

            this.leafNodes = new HashMap<>();

            Collection> endNodes = new LinkedHashSet<>();
            selectLeafAndEndNodes(taskGraph, restrictableNodes, this.leafNodes, endNodes);
            addMissingEndNodes(taskGraph, endNodes);

            List> sortedLeafs
                    = GraphUtils.sortRecursively(taskGraph.getDependencyGraph(), endNodes, leafNodes.keySet());

            this.endNodesToLeafs = taskGraph.getForwardGraph().getAllLeafToRootNodes(sortedLeafs);
            this.leafsToEndNodes = taskGraph.getDependencyGraph().getAllLeafToRootNodes(endNodes);

            this.mainLock = new ReentrantLock();
            this.maxRetainedLeafNodes = maxRetainedLeafNodes;

            this.endNodeQueue = new LinkedHashSet<>(queueSorter.apply(endNodesToLeafs.keySet()));
            this.computedEndNodes = new HashSet<>();
            this.retainingNotComputedEndNodes = new LinkedHashSet<>();
            this.releasedNotComputedEndNodes = new HashSet<>();
            this.scheduledLeafNodes = new HashMap<>();
        }

        private static  N pollCollection(Collection src) {
            if (src.isEmpty()) {
                return null;
            }

            Iterator itr = src.iterator();
            N result = itr.next();
            itr.remove();
            return result;
        }

        private TaskNodeKey pollNextEndNode() {
            TaskNodeKey candidateEndNode = pollCollection(retainingNotComputedEndNodes);
            if (candidateEndNode == null) {
                return pollCollection(endNodeQueue);
            } else {
                endNodeQueue.remove(candidateEndNode);
                return candidateEndNode;
            }
        }

        private void scheduleOne(List releaseTasks) {
            TaskNodeKey candidateEndNode = pollNextEndNode();
            if (candidateEndNode == null) {
                return;
            }

            if (!computedEndNodes.contains(candidateEndNode)) {
                releasedNotComputedEndNodes.add(candidateEndNode);
            }

            Set> candidateLeafs
                    = endNodesToLeafs.getOrDefault(candidateEndNode, Collections.emptySet());
            addScheduledLeafs(candidateLeafs);
            candidateLeafs.forEach((leaf) -> {
                Runnable releaseTask = leafNodes.remove(leaf);
                if (releaseTask != null) {
                    releaseTasks.add(releaseTask);
                }
            });
        }

        private void addScheduledLeafs(Set> newLeafs) {
            newLeafs.forEach(this::addScheduledLeaf);
        }

        private void addScheduledLeaf(TaskNodeKey leafKey) {
            Set> scheduledRetainingEndNodes
                    = scheduledLeafNodes.computeIfAbsent(leafKey, (key) -> new HashSet<>());

            Set> retainingEndNodes
                    = leafsToEndNodes.getOrDefault(leafKey, Collections.emptySet());
            retainingEndNodes.forEach((retainingEndNode) -> {
                if (!computedEndNodes.contains(retainingEndNode)) {
                    if (!releasedNotComputedEndNodes.contains(retainingEndNode)) {
                        retainingNotComputedEndNodes.add(retainingEndNode);
                    }
                    scheduledRetainingEndNodes.add(retainingEndNode);
                }
            });

            if (scheduledRetainingEndNodes.isEmpty()) {
                scheduledLeafNodes.remove(leafKey);
            }
        }

        private void removeLeafNode(TaskNodeKey endNode, TaskNodeKey leaf) {
            Set> retainingEndNodes = scheduledLeafNodes.get(leaf);
            if (retainingEndNodes == null) {
                // This should never happen.
                return;
            }

            retainingEndNodes.remove(endNode);
            if (retainingEndNodes.isEmpty()) {
                scheduledLeafNodes.remove(leaf);
            }
        }

        private void removeLeafNodes(
                TaskNodeKey endNode,
                Set> leafs) {
            leafs.forEach((leaf) -> {
                removeLeafNode(endNode, leaf);
            });
        }

        private void scheduleUnsafe(List releaseTasks) {
            if (releasedNotComputedEndNodes.isEmpty()) {
                // If we do not have any end node released to be executed,
                // we must release at least one to avoid failing to release anything.
                scheduleOne(releaseTasks);
            }

            // FIXME: This may result in nodes never being released
            while (!endNodeQueue.isEmpty() && scheduledLeafNodes.size() < maxRetainedLeafNodes) {
                scheduleOne(releaseTasks);
            }
        }

        public void scheduleUnsafe() {
            List releaseTasks = new ArrayList<>();
            scheduleUnsafe(releaseTasks);
            releaseTasks.forEach(Runnable::run);
        }

        @Override
        public void setNodeComputed(TaskNodeKey nodeKey) {
            Set> leafs = endNodesToLeafs.get(nodeKey);
            if (leafs == null) {
                return;
            }

            List releaseTasks = new ArrayList<>();

            mainLock.lock();
            try {
                computedEndNodes.add(nodeKey);
                releasedNotComputedEndNodes.remove(nodeKey);
                retainingNotComputedEndNodes.remove(nodeKey);

                removeLeafNodes(nodeKey, leafs);
                scheduleUnsafe(releaseTasks);
            } finally {
                mainLock.unlock();
            }

            releaseTasks.forEach(Runnable::run);
        }
    }
}