com.ondeck.datapipes.flow.ParallelExecuteVisitor.java Source code

Java tutorial

Introduction

Here is the source code for com.ondeck.datapipes.flow.ParallelExecuteVisitor.java

Source

/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2015 OnDeck
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 * Except as contained in this notice, the name of a copyright holder shall not be used in advertising or otherwise to promote the sale, use or other dealings in this Software without prior written authorization of the copyright holder.
 */

package com.ondeck.datapipes.flow;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;

import org.apache.commons.pool2.BaseKeyedPooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;

/**
 * An execute visitor allowing the simultaneous execution of parallel flow branches.
 * When the visitor encounters a {@code MergeStep}, it executes each parent branch into a separate thread. If the
 * {@code MergeStep} has only one parent, the execution of that branch is done in the same thread.
 * When the visitor encouters encounter a {@code MultiStep}, each of the execution is done in a separate thread. If the
 * {@code MultiStep} input has only one element, the execution is done in the same thread.
 * @param <R> The return type of the flow to execute.
 */
public class ParallelExecuteVisitor<R> extends SimpleExecuteVisitor<R> {
    private static class BeaconPool extends GenericKeyedObjectPool<FlowStep<?>, Object> {
        private static class BeaconPoolFactory extends BaseKeyedPooledObjectFactory<FlowStep<?>, Object> {

            @Override
            public Object create(FlowStep<?> key) throws Exception {
                return new Object();
            }

            @Override
            public PooledObject<Object> wrap(Object value) {
                return new DefaultPooledObject<>(value);
            }
        }

        private static final GenericKeyedObjectPoolConfig poolConfig = new GenericKeyedObjectPoolConfig();
        static {
            poolConfig.setMaxTotalPerKey(1);
        }

        public BeaconPool() {
            super(new BeaconPoolFactory(), poolConfig);
        }
    }

    private final BeaconPool beaconPool = new BeaconPool();
    protected final ExecutorService executor;

    /**
     * Creates a {@code ParallelExecuteVisitor} with the specified {@code ExecutorService}. If using an
     * {@code ExecutorService} with a limited number of threads, make sure to provide enough threads for the flow or it
     * will result in deadlocks.
     * @param executor
     */
    public ParallelExecuteVisitor(ExecutorService executor) {
        super(Collections.synchronizedMap(new HashMap<FlowStep<?>, Object>()));
        this.executor = executor;
    }

    private Object getBeacon(FlowStep<?> step) throws FlowException {
        try {
            return beaconPool.borrowObject(step);
        } catch (Exception e) {
            throw new FlowException("error getting beacon for step " + step, e);
        }
    }

    private void returnBeacon(FlowStep<?> step, Object beacon) {
        beaconPool.returnObject(step, beacon);
    }

    @Override
    public <T> void visit(DataProviderStep<T> step) throws FlowException {
        Object beacon = getBeacon(step);
        try {
            super.visit(step);
        } finally {
            returnBeacon(step, beacon);
        }
    }

    @Override
    public <S, T> void visit(DataPipeStep<S, T> step) throws FlowException {
        Object beacon = getBeacon(step);
        try {
            super.visit(step);
        } finally {
            returnBeacon(step, beacon);
        }
    }

    @Override
    public <S, T> void visit(FunctionStep<S, T> step) throws FlowException {
        Object beacon = getBeacon(step);
        try {
            super.visit(step);
        } finally {
            returnBeacon(step, beacon);
        }
    }

    private static class BranchRunner<R> implements Callable<Object> {
        private final ParallelExecuteVisitor<R> visitor;
        private final FlowStep<?> step;

        public BranchRunner(ParallelExecuteVisitor<R> visitor, FlowStep<?> step) {
            this.visitor = visitor;
            this.step = step;
        }

        @Override
        public Object call() throws Exception {
            step.accept(visitor);
            return visitor.stepResults.get(step);
        }
    }

    @Override
    public <T> void visit(MergeStep<T> step) throws FlowException {
        Object beacon = getBeacon(step);
        try {
            if (!stepResults.containsKey(step)) {
                FlowStep<?>[] parentSteps = step.getParentSteps();
                if (parentSteps.length < 2) {
                    super.visit(step);
                } else {
                    Map<FlowStep<?>, Future<?>> futures = new HashMap<>();
                    for (FlowStep<?> parentStep : parentSteps) {
                        BranchRunner<R> branchRunner = new BranchRunner<>(this, parentStep);
                        Future<?> future = executor.submit(branchRunner);
                        futures.put(parentStep, future);
                    }
                    Map<Object, Object> parentResults = new HashMap<>();
                    for (Map.Entry<FlowStep<?>, Future<?>> execution : futures.entrySet()) {
                        Object parentResult;
                        try {
                            parentResult = execution.getValue().get();
                        } catch (InterruptedException | ExecutionException e) {
                            throw new FlowException("flow branch execution failed: " + e.getMessage(), e);
                        }
                        if (parentResults.containsKey(execution.getKey().getIndex())) {
                            throw new FlowException(
                                    String.format("data already exists for index %s, specify a unique name",
                                            execution.getKey().getIndex()));
                        }
                        parentResults.put(execution.getKey().getIndex(), parentResult);
                    }

                    T result = null;
                    boolean success = false;
                    List<Object> listenerData = ListenerHelper.before(listeners, step);
                    try {
                        result = step.execute(parentResults);
                        success = true;
                    } finally {
                        ListenerHelper.after(listeners, step, success, listenerData);
                        stepResults.put(step, result);
                    }
                }
            }
        } finally {
            returnBeacon(step, beacon);
        }
    }

    @Override
    public <I, O, S extends Collection<I>, T extends Collection<O>> void visit(MultiStep<I, O, S, T> step)
            throws FlowException {
        Object beacon = getBeacon(step);
        try {
            if (!stepResults.containsKey(step)) {
                step.getParentStep().accept(this);
                @SuppressWarnings("unchecked")
                S parentResult = (S) stepResults.get(step.getParentStep());
                if (parentResult != null) {
                    T result = null;
                    try {
                        result = step.execute(parentResult, executor, listeners);
                    } finally {
                        stepResults.put(step, result);
                    }
                }
            }
        } finally {
            returnBeacon(step, beacon);
        }
    }

    @Override
    public <T, U> void visit(ConditionalStep<T, U> step) throws FlowException {
        Object beacon = getBeacon(step);
        try {
            super.visit(step);
        } finally {
            returnBeacon(step, beacon);
        }
    }

}