HdfsWritableBolt.java

package resa.evaluation.migrate;

import backtype.storm.Config;
import backtype.storm.serialization.SerializationFactory;
import backtype.storm.task.OutputCollector;
import backtype.storm.task.TopologyContext;
import backtype.storm.topology.IRichBolt;
import backtype.storm.utils.Utils;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.KryoSerializable;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.serializers.DefaultSerializers;
import org.apache.curator.framework.CuratorFramework;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileAlreadyExistsException;
import org.apache.hadoop.fs.FileSystem;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import redis.clients.jedis.Jedis;
import resa.metrics.MetricNames;
import resa.topology.DelegatedBolt;
import resa.util.ConfigUtil;
import resa.util.NetUtil;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.reflect.Field;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.stream.Stream;

/**
 * Created by ding on 14-6-7.
 */
public class HdfsWritableBolt extends DelegatedBolt {

    private static final Logger LOG = LoggerFactory.getLogger(HdfsWritableBolt.class);
    private static final Map<Integer, Object> DATA_HOLDER = new ConcurrentHashMap<>();
    // this should be part of user resource, but no initialization hook in work.clj yet
    private static CuratorFramework zk = null;

    private static CuratorFramework zkInstance(Map<String, Object> conf) {
        if (zk == null) {
            synchronized (HdfsWritableBolt.class) {
                if (zk == null) {
                    zk = Utils.newCuratorStarted(conf, (List<String>) conf.get(Config.STORM_ZOOKEEPER_SERVERS),
                            conf.get(Config.STORM_ZOOKEEPER_PORT));
                }
            }
        }
        return zk;
    }


    private transient Map<String, Object> conf;
    private transient Kryo kryo;
    private transient Map<String, Object> dataRef;
    private Path loaclDataPath;
    private transient ExecutorService threadPool;
    private int myTaskId;
    private String taskZkNode;
    private Configuration hdfsConf;

    public HdfsWritableBolt(IRichBolt delegate) {
        super(delegate);
    }

    @Override
    public void prepare(Map conf, TopologyContext context, OutputCollector outputCollector) {
        this.conf = conf;
        this.hdfsConf = new Configuration();
        this.myTaskId = context.getThisTaskId();
        this.taskZkNode = String.format("%s/%s/task-%03d", conf.getOrDefault("resa.migrate.zkroot", "/resa"),
                context.getStormId(), context.getThisTaskId());
        ensureZkNode();
        kryo = SerializationFactory.getKryo(conf);
        //LOG.info("Bolt " + getDelegate().getClass() + " need persist");
        dataRef = getDataRef(context);
        loaclDataPath = Paths.get((String) conf.get(Config.STORM_LOCAL_DIR), "data", context.getStormId());
        if (!Files.exists(loaclDataPath)) {
            try {
                Files.createDirectories(loaclDataPath);
            } catch (FileAlreadyExistsException e) {
            } catch (IOException e) {
                LOG.warn("Cannot create data path: " + loaclDataPath, e);
                loaclDataPath = null;
            }
        }
        if (loaclDataPath != null) {
            loaclDataPath = loaclDataPath.resolve(String.format("task-%03d.data", context.getThisTaskId()));
        }
        long t1 = System.currentTimeMillis();
        seekAndLoadTaskData(context.getStormId(), context.getThisTaskId(), conf);
        long cost = System.currentTimeMillis() - t1;
        LOG.info("Task-{} Load data cost {}ms", context.getThisTaskId(), cost);
        this.threadPool = context.getSharedExecutor();
//        threadPool.submit(() -> writeTimeToRedis(context, cost));
        int checkpointInt = ConfigUtil.getInt(conf, "resa.comp.checkpoint.interval.sec", 180);
        if (checkpointInt > 0) {
            context.registerMetric(MetricNames.SERIALIZED_SIZE, () -> createCheckpointAndGetSize(context),
                    checkpointInt);
        }
        super.prepare(conf, context, outputCollector);
        if (!dataRef.isEmpty()) {
            reportMyData();
        }
    }

    private void writeTimeToRedis(TopologyContext context, long time) {
        Jedis jedis = new Jedis((String) conf.get("redis.host"), ConfigUtil.getInt(conf, "redis.port", 6379));
        jedis.set(context.getStormId() + "-" + context.getThisComponentId() + "-" + context.getThisTaskId()
                + "-start-time", time < 100 ? "zero" : String.valueOf(time));
        jedis.disconnect();
    }

    private void ensureZkNode() {
        try {
            if (zkInstance(conf).checkExists().forPath(taskZkNode) == null) {
                zkInstance(conf).create().creatingParentsIfNeeded().forPath(taskZkNode, "".getBytes());
            }
        } catch (Exception e) {
            throw new RuntimeException("Check task zkNode failed: " + taskZkNode, e);
        }
    }

    private String getDataLocation() {
        try {
            byte[] data = zkInstance(conf).getData().forPath(taskZkNode);
            String location = new String(data);
            return location.equals("") ? null : location;
        } catch (Exception e) {
            LOG.warn("Read host failed, taskZkNode=" + taskZkNode, e);
        }
        return null;
    }

    private void seekAndLoadTaskData(String topoId, int taskId, Map<String, Object> conf) {
        /* load data:
           1. find in memory, maybe data migration was not required
           2. retrieve data from src worker using TCP
           3. try to load checkpoint from redis or other distributed systems
        */
        Object data = DATA_HOLDER.remove(taskId);
        if (data != null) {
            dataRef.putAll((Map<String, Object>) data);
            LOG.info("Task-{} load data from memory", taskId);
            return;
        }
        try {
            FileSystem fs = FileSystem.get(hdfsConf);
            org.apache.hadoop.fs.Path file = getFilePath(conf, topoId, taskId);
            if (fs.exists(file)) {
                FSDataInputStream in = fs.open(file);
                loadData(in);
                in.close();
            }
        } catch (IOException e) {
            LOG.warn("read from hdfs failed for task-" + taskId, e);
        }
    }

    private org.apache.hadoop.fs.Path getFilePath(Map<String, Object> conf, String topoId, int taskId) {
        String topoName = topoId.split("-")[0];
        return new org.apache.hadoop.fs.Path(String.format("/resa/%s/task-%03d.data", topoName, taskId));
    }

    private static Map<String, Object> getDataRef(TopologyContext context) {
        try {
            Field f = context.getClass().getDeclaredField("_taskData");
            f.setAccessible(true);
            return (Map<String, Object>) f.get(context);
        } catch (Exception e) {
            throw new IllegalStateException("Cannot get task data ref", e);
        }
    }

    private Long createCheckpointAndGetSize(TopologyContext context) {
        Long size = -1L;
        if (!dataRef.isEmpty() && loaclDataPath != null) {
            try {
                Path tmpFile = Files.createTempFile("resa-cp-", ".tmp");
                size = writeData(Files.newOutputStream(tmpFile));
                Files.move(tmpFile, loaclDataPath, StandardCopyOption.REPLACE_EXISTING);
                reportMyData();
                threadPool.submit(() -> {
                    org.apache.hadoop.fs.Path file = getFilePath(conf, context.getStormId(), context.getThisTaskId());
                    try (OutputStream out = FileSystem.get(hdfsConf).create(file, true, 12 * 1024, (short) 2,
                            32 * 1024 * 1024L)) {
                        Files.copy(loaclDataPath, out);
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                });
            } catch (Exception e) {
                LOG.warn("Save bolt failed", e);
            }
        }
        return size;
    }

    private void reportMyData() {
        try {
            zkInstance(conf).setData().forPath(taskZkNode, NetUtil.getLocalIP().getBytes());
        } catch (Exception e) {
            LOG.warn("Report my data failed", e);
        }
    }

    private int loadData(InputStream in) {
        Input kryoIn = new Input(in);
        int size = kryoIn.readInt();
        for (int i = 0; i < size; i++) {
            String key = kryoIn.readString();
            Class c = kryo.readClass(kryoIn).getType();
            Object v;
            if (KryoSerializable.class.isAssignableFrom(c)) {
                v = new DefaultSerializers.KryoSerializableSerializer().read(kryo, kryoIn, c);
            } else {
                v = kryo.readClassAndObject(kryoIn);
            }
            dataRef.put(key, v);
        }
        kryoIn.close();
        return kryoIn.total();
    }

    private long writeData(OutputStream out) {
        Output kryoOut = new Output(out);
        kryoOut.writeInt(dataRef.size());
        dataRef.forEach((k, v) -> {
            kryoOut.writeString(k);
            if (v instanceof KryoSerializable) {
                kryo.writeClass(kryoOut, v.getClass());
                ((KryoSerializable) v).write(kryo, kryoOut);
            } else {
                kryo.writeClassAndObject(kryoOut, v);
            }
        });
        long size = kryoOut.total();
        kryoOut.close();
        return size;
    }

    @Override
    public void cleanup() {
        super.cleanup();
        if (!dataRef.isEmpty()) {
            saveData();
        }
    }

    private void saveData() {
        String workerTasks = System.getProperty("new-worker-assignment");
        LOG.info("new-worker-assignment: {}", workerTasks);
        boolean needMigration = workerTasks == null || workerTasks.equals("") ? true :
                Stream.of(workerTasks.split(",")).mapToInt(Integer::parseInt).noneMatch(i -> myTaskId == i);
        if (!needMigration) {
            DATA_HOLDER.put(myTaskId, dataRef);
            LOG.info("Put task-{} data to memory DATA_HOLDER", myTaskId);
        }
//        else if (loaclDataPath != null) {
//            try {
//                writeData(Files.newOutputStream(loaclDataPath));
//            } catch (IOException e) {
//                LOG.info("Save data failed", e);
//            }
//        }
    }

}