io.urmia.job.pub.InputAwareZkJobQueue.java Source code

Java tutorial

Introduction

Here is the source code for io.urmia.job.pub.InputAwareZkJobQueue.java

Source

package io.urmia.job.pub;

/**
 *
 * Copyright 2014 by Amin Abbaspour
 *
 * This file is part of Urmia.io
 *
 * Urmia.io is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Urmia.io is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Urmia.io.  If not, see <http://www.gnu.org/licenses/>.
 */

import com.google.common.base.Function;
import com.google.common.base.Optional;
import com.google.common.collect.*;
import io.urmia.md.model.job.*;
import io.urmia.md.model.storage.ObjectName;
import io.urmia.md.service.MetadataService;
import io.urmia.naming.model.NodeType;
import io.urmia.naming.service.NamingService;
import io.urmia.util.StringUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.api.CuratorWatcher;
import org.apache.curator.framework.recipes.queue.DistributedQueue;
import org.apache.curator.framework.recipes.queue.QueueBuilder;
import org.apache.curator.framework.recipes.queue.QueueSerializer;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

import static io.urmia.md.model.job.JobInput.END;

/**
 * todo: review algorithm to reuse existing queues. not always first storage node.
 */
public class InputAwareZkJobQueue implements JobQueue {

    private static final Logger log = LoggerFactory.getLogger(InputAwareZkJobQueue.class);

    private final CuratorFramework client;
    private final MetadataService mds;
    private final JobDefinition jobDef;
    private final String jobId;
    private final JobStatusUpdateHandler statHandler;
    private final NamingService ns;

    private final Map<String, DistributedQueue<JobInput>> queueMap = new LinkedHashMap<String, DistributedQueue<JobInput>>();

    //private DistributedQueue<JobInput> reduceQueue = null;

    private final Map<String, Boolean> finishedMap = new LinkedHashMap<String, Boolean>();

    private final Map<String, AtomicInteger> inputCounter = new LinkedHashMap<String, AtomicInteger>();

    public InputAwareZkJobQueue(CuratorFramework client, MetadataService mds, NamingService ns,
            JobCreateRequest request, JobStatusUpdateHandler statHandler) throws Exception {
        this.client = client;
        this.mds = mds;
        this.jobDef = request.job;
        this.jobId = request.getId();
        this.statHandler = statHandler;
        this.ns = ns;
    }

    @Override
    public void put(JobInput input) throws Exception {

        final Map<String, JobInput> inputPerNode = split(input);

        log.info("split input {} to map: {}", input, inputPerNode);

        for (Map.Entry<String, JobInput> e : inputPerNode.entrySet()) {
            getQueue(e.getKey()).put(e.getValue());
            inputCounter.get(e.getKey()).addAndGet(e.getValue().getCount());
        }

    }

    @Override
    public void close() throws Exception {
        statHandler.markInputDone();

        for (DistributedQueue<JobInput> q : queueMap.values()) {
            log.info("closing DistributedQueue: {}", q);
            q.put(END);
            q.close();
        }
    }

    private Map<String, JobInput> split(JobInput input) throws Exception {
        ImmutableMultimap.Builder<String, ObjectName> b = ImmutableMultimap.builder();

        for (/*ObjectName*/String in : input) {
            Optional<ObjectName> optOn = ObjectName.of(in);
            if (!optOn.isPresent()) {
                log.warn("no ObjectName for input: {}", in);
                continue;
            }
            ObjectName on = optOn.get();
            ServiceInstance<NodeType> s = whereIs(on);
            log.info("looking where is on: {} -> {}", on, s);
            if (s != null)
                b.put(s.getId(), on);
        }

        Map<String, Collection<ObjectName>> m = b.build().asMap();

        return Maps.transformValues(m, transformer());
    }

    private Function<Collection<ObjectName>, JobInput> transformer() {
        return new Function<Collection<ObjectName>, JobInput>() {
            public JobInput apply(Collection<ObjectName> input) {
                return new LineJobInput(input);
            }
        };
    }

    private ServiceInstance<NodeType> findRunnerInstanceUp(List<String> ids) throws Exception {
        for (String id : ids) {
            ServiceInstance<NodeType> si = ns.discover(NodeType.JRS, id);
            if (si != null) {
                log.debug("found running jrs node of id: {}", id);
                return si;
            } else {
                log.debug("jrs node is down: {}", id);
            }
        }
        return null;
    }

    private ServiceInstance<NodeType> whereIs(ObjectName on) throws Exception {
        List<String> storedNodes = mds.storedNodes(on);
        log.info("stored nodes on {} -> {}", on, storedNodes);
        return findRunnerInstanceUp(storedNodes);
    }

    private DistributedQueue<JobInput> getQueue(String instanceId) throws Exception {

        DistributedQueue<JobInput> cachedQ = queueMap.get(instanceId);
        if (cachedQ != null)
            return cachedQ;

        String json = jobDef.toString();
        byte[] bytes = json.getBytes();

        String jobs = "/urmia/1/" + instanceId + "/jobs";

        if (client.checkExists().forPath(jobs) == null)
            client.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(jobs);

        String p = jobs + "/" + jobId;
        log.info("creating job: {} -> {}", p, json);

        client.create().withMode(CreateMode.PERSISTENT).forPath(p, bytes);

        String ip = p + "/live/in";

        QueueBuilder<JobInput> builder = QueueBuilder.builder(client, null, serializer, ip);

        client.checkExists().usingWatcher(new QueueExistsWatcher()).forPath(p);

        Thread.sleep(500); // todo: why?

        DistributedQueue<JobInput> q = builder.buildQueue();
        q.start();

        queueMap.put(instanceId, q);
        finishedMap.put(instanceId, Boolean.FALSE);
        inputCounter.put(instanceId, new AtomicInteger(0));
        statHandler.incTasks(1);

        return q;
    }

    private DistributedQueue<JobInput> getReduceQueue(String instanceId) throws Exception {

        String json = jobDef.toString();
        byte[] bytes = json.getBytes();

        String jobs = "/urmia/1/" + instanceId + "/jobs";

        if (client.checkExists().forPath(jobs) == null)
            client.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(jobs);

        String p = jobs + "/" + jobId + "-reduce";
        log.info("creating job: {} -> {}", p, json);

        client.create().withMode(CreateMode.PERSISTENT).forPath(p, bytes);

        String ip = p + "/live/in";

        QueueBuilder<JobInput> builder = QueueBuilder.builder(client, null, serializer, ip);

        client.checkExists().usingWatcher(new QueueExistsWatcher()).forPath(p);

        Thread.sleep(500); // todo: why?

        DistributedQueue<JobInput> q = builder.buildQueue();
        q.start();

        return q;
    }

    private class QueueExistsWatcher implements CuratorWatcher {

        @Override
        public void process(WatchedEvent event) throws Exception {
            log.info("QueueExistsWatcher received event: {}", event);

            if (event.getType() == Watcher.Event.EventType.NodeDeleted) {
                String qPath = event.getPath();
                String hostId = getNodeId(qPath);

                if (isReduceQueuePath(qPath)) {
                    log.info("reduce job finished on q: {}", qPath);
                    addToOutputs(hostId, false, true);

                } else {
                    log.info("map job finished on node: {}", hostId);

                    boolean allFinished = markFinished(hostId);

                    addToOutputs(hostId, jobDef.hasReduce(), false);

                    if (allFinished && jobDef.hasReduce()) {

                        if (mapOutputs.isEmpty()) {
                            log.warn("map outputs empty. skipping reduce phase.");
                            return;
                        }

                        String reduceHostId = getReduceNodeId();

                        log.info("process finished in all workers. best reduce node: {}", reduceHostId);

                        log.info("reduce input(s): {}", mapOutputs);

                        buildReduceQueue(reduceHostId);
                    }
                }
            }
        }
    }

    private boolean isReduceQueuePath(String path) {
        return path.endsWith("-reduce");
    }

    private void buildReduceQueue(String hostId) throws Exception {
        DistributedQueue<JobInput> reduceQ = getReduceQueue(hostId);

        Thread.sleep(500);

        JobInput reduceInput = new LineJobInput(mapOutputs);

        log.info("putting map results into reduce q: {}", reduceInput);
        reduceQ.put(reduceInput);
        reduceQ.put(JobInput.END);
    }

    private List<String> mapOutputs = new LinkedList<String>();

    private void addToOutputs(String hostId, boolean queue, boolean reducePhase) throws Exception {
        if (queue) {
            String path = '/' + jobDef.getOwner() + "/jobs/" + jobId + "/stor" + '/' + jobDef.getOwner() + '/'
                    + ObjectName.Namespace.JOBS.path + '/' + jobId + '/' + hostId;
            mapOutputs.add(path);
            return;
        }
        addToOutputsZk(hostId, reducePhase);
    }

    private void addToOutputsZk(String hostId, boolean reducePhase) throws Exception {

        String outPath = "/urmia/1/jobs/" + jobDef.getOwner() + "/" + jobId + "/live/out.txt";

        byte[] existing = new byte[0];
        try {
            existing = client.getData().forPath(outPath);
        } catch (Exception ignored) {
        }
        String path = '/' + jobDef.getOwner() + "/jobs/" + jobId + "/stor" + '/' + jobDef.getOwner() + '/'
                + ObjectName.Namespace.JOBS.path + '/' + jobId + '/' + hostId;

        if (reducePhase)
            path += "-reduce";

        log.info("adding to outputs of job {}: {}", jobId, path);

        client.create().creatingParentsIfNeeded().inBackground().forPath(outPath,
                StringUtils.append(existing, path));
    }

    /*
    private String getMapPhaseOutputs() {
    StringBuilder b = new StringBuilder();
        
    for (String host : queueMap.keySet()) {
        try {
            String d = new String(client.getData().forPath("/urmia/1/" + host + "/jobs/" + jobId + "/live/out.txt")).trim();
            b.append(d).append('\n');
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
        
    return b.toString();
    }
    */

    private String getReduceNodeId() {
        int min = -1;
        String id = "";
        for (Map.Entry<String, AtomicInteger> e : inputCounter.entrySet())
            if (min < e.getValue().get()) {
                min = e.getValue().get();
                id = e.getKey();
            }
        return id;
    }

    private boolean markFinished(String id) {
        finishedMap.put(id, Boolean.TRUE);
        for (Boolean b : finishedMap.values())
            if (!b)
                return false;
        return true;
    }

    // /urmia/1/2c8a34b9-6bcb-4e8d-982c-35bdb9931ff1/jobs/c1481511-1870-4a28-9a09-9f31133b12fc(-reduce)
    private static String getNodeId(String qPath) {
        int index = qPath.indexOf("/urmia/1/") + 9;
        int end = qPath.indexOf('/', index);
        return qPath.substring(index, end);
    }

    private static final QueueSerializer<JobInput> serializer = new JobInputQueueSerializer();

    private static class JobInputQueueSerializer implements QueueSerializer<JobInput> {

        @Override
        public byte[] serialize(JobInput s) {
            log.info("serialize job: {}", s);
            return s.toBytes();
        }

        @Override
        public JobInput deserialize(byte[] bytes) {
            return new LineJobInput(new String(bytes).trim());
        }
    }

}