external/storm-hdfs/src/main/java/org/apache/storm/hdfs/bolt/AbstractHdfsBolt.java - storm - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.  The ASF licenses this file to you under the Apache License, Version
  * 2.0 (the "License"); you may not use this file except in compliance with the License.  You may obtain a copy of the License at
  *
  * http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */

 package org.apache.storm.hdfs.bolt;

 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Timer;
 import java.util.TimerTask;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.storm.hdfs.bolt.format.FileNameFormat;
 import org.apache.storm.hdfs.bolt.rotation.FileRotationPolicy;
 import org.apache.storm.hdfs.bolt.rotation.TimedRotationPolicy;
 import org.apache.storm.hdfs.bolt.sync.SyncPolicy;
 import org.apache.storm.hdfs.common.NullPartitioner;
 import org.apache.storm.hdfs.common.Partitioner;
 import org.apache.storm.hdfs.common.rotation.RotationAction;
 import org.apache.storm.hdfs.security.HdfsSecurityUtil;
 import org.apache.storm.task.OutputCollector;
 import org.apache.storm.task.TopologyContext;
 import org.apache.storm.topology.OutputFieldsDeclarer;
 import org.apache.storm.topology.base.BaseRichBolt;
 import org.apache.storm.tuple.Tuple;
 import org.apache.storm.utils.TupleUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 public abstract class AbstractHdfsBolt extends BaseRichBolt {
     private static final Logger LOG = LoggerFactory.getLogger(AbstractHdfsBolt.class);
     private static final Integer DEFAULT_RETRY_COUNT = 3;
     /**
      * Half of the default Config.TOPOLOGY_MESSAGE_TIMEOUT_SECS
      */
     private static final int DEFAULT_TICK_TUPLE_INTERVAL_SECS = 15;
     private static final Integer DEFAULT_MAX_OPEN_FILES = 50;

     protected Map<String, Writer> writers;
     protected Map<String, Integer> rotationCounterMap = new HashMap<>();
     protected List<RotationAction> rotationActions = new ArrayList<>();
     protected OutputCollector collector;
     protected transient FileSystem fs;
     protected SyncPolicy syncPolicy;
     protected FileRotationPolicy rotationPolicy;
     protected FileNameFormat fileNameFormat;
     protected String fsUrl;
     protected String configKey;
     protected transient Object writeLock;
     protected transient Timer rotationTimer; // only used for TimedRotationPolicy
     protected long offset = 0;
     protected Integer fileRetryCount = DEFAULT_RETRY_COUNT;
     protected Integer tickTupleInterval = DEFAULT_TICK_TUPLE_INTERVAL_SECS;
     protected Integer maxOpenFiles = DEFAULT_MAX_OPEN_FILES;
     protected Partitioner partitioner = new NullPartitioner();
     protected transient Configuration hdfsConfig;
     private List<Tuple> tupleBatch = new LinkedList<>();

     protected void rotateOutputFile(Writer writer) throws IOException {
         LOG.info("Rotating output file...");
         long start = System.currentTimeMillis();
         synchronized (this.writeLock) {
             writer.close();

             LOG.info("Performing {} file rotation actions.", this.rotationActions.size());
             for (RotationAction action : this.rotationActions) {
                 action.execute(this.fs, writer.getFilePath());
             }
         }
         long time = System.currentTimeMillis() - start;
         LOG.info("File rotation took {} ms.", time);
     }

     /**
      * Marked as final to prevent override. Subclasses should implement the doPrepare() method.
      */
     @Override
     public final void prepare(Map<String, Object> conf, TopologyContext topologyContext, OutputCollector collector) {
         this.writeLock = new Object();
         if (this.syncPolicy == null) {
             throw new IllegalStateException("SyncPolicy must be specified.");
         }
         if (this.rotationPolicy == null) {
             throw new IllegalStateException("RotationPolicy must be specified.");
         }
         if (this.fsUrl == null) {
             throw new IllegalStateException("File system URL must be specified.");
         }

         writers = new WritersMap(this.maxOpenFiles, collector);

         this.collector = collector;
         this.fileNameFormat.prepare(conf, topologyContext);
         this.hdfsConfig = new Configuration();
         Map<String, Object> map = (Map<String, Object>) conf.get(this.configKey);
         if (map != null) {
             for (String key : map.keySet()) {
                 this.hdfsConfig.set(key, String.valueOf(map.get(key)));
             }
         }

         try {
             HdfsSecurityUtil.login(conf, hdfsConfig);
             doPrepare(conf, topologyContext, collector);
         } catch (Exception e) {
             throw new RuntimeException("Error preparing HdfsBolt: " + e.getMessage(), e);
         }

         if (this.rotationPolicy instanceof TimedRotationPolicy) {
             startTimedRotationPolicy();
         }
     }

     @Override
     public final void execute(Tuple tuple) {

         synchronized (this.writeLock) {
             boolean forceSync = false;
             Writer writer = null;
             String writerKey = null;

             if (TupleUtils.isTick(tuple)) {
                 LOG.debug("TICK! forcing a file system flush");
                 this.collector.ack(tuple);
                 forceSync = true;
             } else {

                 writerKey = getHashKeyForTuple(tuple);

                 try {
                     writer = getOrCreateWriter(writerKey, tuple);
                     this.offset = writer.write(tuple);
                     tupleBatch.add(tuple);
                 } catch (IOException e) {
                     //If the write failed, try to sync anything already written
                     LOG.info("Tuple failed to write, forcing a flush of existing data.");
                     this.collector.reportError(e);
                     forceSync = true;
                     this.collector.fail(tuple);
                 }
             }

             if (this.syncPolicy.mark(tuple, this.offset) || (forceSync && tupleBatch.size() > 0)) {
                 int attempts = 0;
                 boolean success = false;
                 IOException lastException = null;
                 // Make every attempt to sync the data we have.  If it can't be done then kill the bolt with
                 // a runtime exception.  The filesystem is presumably in a very bad state.
                 while (success == false && attempts < fileRetryCount) {
                     attempts += 1;
                     try {
                         syncAllWriters();
                         LOG.debug("Data synced to filesystem. Ack'ing [{}] tuples", tupleBatch.size());
                         for (Tuple t : tupleBatch) {
                             this.collector.ack(t);
                         }
                         tupleBatch.clear();
                         syncPolicy.reset();
                         success = true;
                     } catch (IOException e) {
                         LOG.warn("Data could not be synced to filesystem on attempt [{}]", attempts);
                         this.collector.reportError(e);
                         lastException = e;
                     }
                 }

                 // If unsuccesful fail the pending tuples
                 if (success == false) {
                     LOG.warn("Data could not be synced to filesystem, failing this batch of tuples");
                     for (Tuple t : tupleBatch) {
                         this.collector.fail(t);
                     }
                     tupleBatch.clear();

                     throw new RuntimeException("Sync failed [" + attempts + "] times.", lastException);
                 }
             }

             if (writer != null && writer.needsRotation()) {
                 doRotationAndRemoveWriter(writerKey, writer);
             }
         }
     }

     private Writer getOrCreateWriter(String writerKey, Tuple tuple) throws IOException {
         Writer writer;

         writer = writers.get(writerKey);
         if (writer == null) {
             Path pathForNextFile = getBasePathForNextFile(tuple);
             writer = makeNewWriter(pathForNextFile, tuple);
             writers.put(writerKey, writer);
         }
         return writer;
     }

     /**
      * A tuple must be mapped to a writer based on two factors.
      *  - bolt specific logic that must separate tuples into different files in the same directory (see the avro bolt
      *    for an example of this)
      *  - the directory the tuple will be partioned into
      */
     private String getHashKeyForTuple(Tuple tuple) {
         final String boltKey = getWriterKey(tuple);
         final String partitionDir = this.partitioner.getPartitionPath(tuple);
         return boltKey + "****" + partitionDir;
     }

     void doRotationAndRemoveWriter(String writerKey, Writer writer) {
         try {
             rotateOutputFile(writer);
         } catch (IOException e) {
             this.collector.reportError(e);
             LOG.error("File could not be rotated");
             //At this point there is nothing to do.  In all likelihood any filesystem operations will fail.
             //The next tuple will almost certainly fail to write and/or sync, which force a rotation.  That
             //will give rotateAndReset() a chance to work which includes creating a fresh file handle.
         } finally {
             //rotateOutputFile(writer) has closed the writer. It's safe to remove the writer from the map here.
             writers.remove(writerKey);
         }
     }

     @Override
     public Map<String, Object> getComponentConfiguration() {
         return TupleUtils.putTickFrequencyIntoComponentConfig(super.getComponentConfiguration(), tickTupleInterval);
     }

     @Override
     public void declareOutputFields(OutputFieldsDeclarer outputFieldsDeclarer) {
     }

     @Override
     public void cleanup() {
         doRotationAndRemoveAllWriters();
         if (this.rotationTimer != null) {
             this.rotationTimer.cancel();
         }
     }

     private void doRotationAndRemoveAllWriters() {
         synchronized (writeLock) {
             for (final Writer writer : writers.values()) {
                 try {
                     rotateOutputFile(writer);
                 } catch (IOException e) {
                     this.collector.reportError(e);
                     LOG.warn("IOException during scheduled file rotation.", e);
                 }
             }
             //above for-loop has closed all the writers. It's safe to clear the map here.
             writers.clear();
         }
     }

     private void syncAllWriters() throws IOException {
         for (Writer writer : writers.values()) {
             writer.sync();
         }
     }

     private void startTimedRotationPolicy() {
         long interval = ((TimedRotationPolicy) this.rotationPolicy).getInterval();
         this.rotationTimer = new Timer(true);
         TimerTask task = new TimerTask() {
             @Override
             public void run() {
                 doRotationAndRemoveAllWriters();
             }
         };
         this.rotationTimer.scheduleAtFixedRate(task, interval, interval);
     }

     protected Path getBasePathForNextFile(Tuple tuple) {

         final String partitionPath = this.partitioner.getPartitionPath(tuple);
         final int rotation;
         if (rotationCounterMap.containsKey(partitionPath)) {
             rotation = rotationCounterMap.get(partitionPath) + 1;
         } else {
             rotation = 0;
         }
         rotationCounterMap.put(partitionPath, rotation);

         return new Path(this.fsUrl + this.fileNameFormat.getPath() + partitionPath,
                         this.fileNameFormat.getName(rotation, System.currentTimeMillis()));
     }

     protected abstract void doPrepare(Map<String, Object> conf, TopologyContext topologyContext, OutputCollector collector) throws
         IOException;

     protected abstract String getWriterKey(Tuple tuple);

     protected abstract Writer makeNewWriter(Path path, Tuple tuple) throws IOException;

     static class WritersMap extends LinkedHashMap<String, Writer> {
         final long maxWriters;
         final OutputCollector collector;

         public WritersMap(long maxWriters, OutputCollector collector) {
             super((int) maxWriters, 0.75f, true);
             this.maxWriters = maxWriters;
             this.collector = collector;
         }

         @Override
         protected boolean removeEldestEntry(Map.Entry<String, Writer> eldest) {
             if (this.size() > this.maxWriters) {
                 //The writer must be closed before removed from the map.
                 //If it failed, we might lose some data.
                 try {
                     eldest.getValue().close();
                 } catch (IOException e) {
                     collector.reportError(e);
                     LOG.error("Failed to close the eldest Writer");
                 }
                 return true;
             } else {
                 return false;
             }
         }
     }
 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version
	* 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
	* and limitations under the License.
	*/

	package org.apache.storm.hdfs.bolt;

	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.LinkedHashMap;
	import java.util.LinkedList;
	import java.util.List;
	import java.util.Map;
	import java.util.Timer;
	import java.util.TimerTask;
	import org.apache.hadoop.conf.Configuration;
	import org.apache.hadoop.fs.FileSystem;
	import org.apache.hadoop.fs.Path;
	import org.apache.storm.hdfs.bolt.format.FileNameFormat;
	import org.apache.storm.hdfs.bolt.rotation.FileRotationPolicy;
	import org.apache.storm.hdfs.bolt.rotation.TimedRotationPolicy;
	import org.apache.storm.hdfs.bolt.sync.SyncPolicy;
	import org.apache.storm.hdfs.common.NullPartitioner;
	import org.apache.storm.hdfs.common.Partitioner;
	import org.apache.storm.hdfs.common.rotation.RotationAction;
	import org.apache.storm.hdfs.security.HdfsSecurityUtil;
	import org.apache.storm.task.OutputCollector;
	import org.apache.storm.task.TopologyContext;
	import org.apache.storm.topology.OutputFieldsDeclarer;
	import org.apache.storm.topology.base.BaseRichBolt;
	import org.apache.storm.tuple.Tuple;
	import org.apache.storm.utils.TupleUtils;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	public abstract class AbstractHdfsBolt extends BaseRichBolt {
	private static final Logger LOG = LoggerFactory.getLogger(AbstractHdfsBolt.class);
	private static final Integer DEFAULT_RETRY_COUNT = 3;
	/**
	* Half of the default Config.TOPOLOGY_MESSAGE_TIMEOUT_SECS
	*/
	private static final int DEFAULT_TICK_TUPLE_INTERVAL_SECS = 15;
	private static final Integer DEFAULT_MAX_OPEN_FILES = 50;

	protected Map<String, Writer> writers;
	protected Map<String, Integer> rotationCounterMap = new HashMap<>();
	protected List<RotationAction> rotationActions = new ArrayList<>();
	protected OutputCollector collector;
	protected transient FileSystem fs;
	protected SyncPolicy syncPolicy;
	protected FileRotationPolicy rotationPolicy;
	protected FileNameFormat fileNameFormat;
	protected String fsUrl;
	protected String configKey;
	protected transient Object writeLock;
	protected transient Timer rotationTimer; // only used for TimedRotationPolicy
	protected long offset = 0;
	protected Integer fileRetryCount = DEFAULT_RETRY_COUNT;
	protected Integer tickTupleInterval = DEFAULT_TICK_TUPLE_INTERVAL_SECS;
	protected Integer maxOpenFiles = DEFAULT_MAX_OPEN_FILES;
	protected Partitioner partitioner = new NullPartitioner();
	protected transient Configuration hdfsConfig;
	private List<Tuple> tupleBatch = new LinkedList<>();

	protected void rotateOutputFile(Writer writer) throws IOException {
	LOG.info("Rotating output file...");
	long start = System.currentTimeMillis();
	synchronized (this.writeLock) {
	writer.close();

	LOG.info("Performing {} file rotation actions.", this.rotationActions.size());
	for (RotationAction action : this.rotationActions) {
	action.execute(this.fs, writer.getFilePath());
	}
	}
	long time = System.currentTimeMillis() - start;
	LOG.info("File rotation took {} ms.", time);
	}

	/**
	* Marked as final to prevent override. Subclasses should implement the doPrepare() method.
	*/
	@Override
	public final void prepare(Map<String, Object> conf, TopologyContext topologyContext, OutputCollector collector) {
	this.writeLock = new Object();
	if (this.syncPolicy == null) {
	throw new IllegalStateException("SyncPolicy must be specified.");
	}
	if (this.rotationPolicy == null) {
	throw new IllegalStateException("RotationPolicy must be specified.");
	}
	if (this.fsUrl == null) {
	throw new IllegalStateException("File system URL must be specified.");
	}

	writers = new WritersMap(this.maxOpenFiles, collector);

	this.collector = collector;
	this.fileNameFormat.prepare(conf, topologyContext);
	this.hdfsConfig = new Configuration();
	Map<String, Object> map = (Map<String, Object>) conf.get(this.configKey);
	if (map != null) {
	for (String key : map.keySet()) {
	this.hdfsConfig.set(key, String.valueOf(map.get(key)));
	}
	}

	try {
	HdfsSecurityUtil.login(conf, hdfsConfig);
	doPrepare(conf, topologyContext, collector);
	} catch (Exception e) {
	throw new RuntimeException("Error preparing HdfsBolt: " + e.getMessage(), e);
	}

	if (this.rotationPolicy instanceof TimedRotationPolicy) {
	startTimedRotationPolicy();
	}
	}

	@Override
	public final void execute(Tuple tuple) {

	synchronized (this.writeLock) {
	boolean forceSync = false;
	Writer writer = null;
	String writerKey = null;

	if (TupleUtils.isTick(tuple)) {
	LOG.debug("TICK! forcing a file system flush");
	this.collector.ack(tuple);
	forceSync = true;
	} else {

	writerKey = getHashKeyForTuple(tuple);

	try {
	writer = getOrCreateWriter(writerKey, tuple);
	this.offset = writer.write(tuple);
	tupleBatch.add(tuple);
	} catch (IOException e) {
	//If the write failed, try to sync anything already written
	LOG.info("Tuple failed to write, forcing a flush of existing data.");
	this.collector.reportError(e);
	forceSync = true;
	this.collector.fail(tuple);
	}
	}

	if (this.syncPolicy.mark(tuple, this.offset) \|\| (forceSync && tupleBatch.size() > 0)) {
	int attempts = 0;
	boolean success = false;
	IOException lastException = null;
	// Make every attempt to sync the data we have. If it can't be done then kill the bolt with
	// a runtime exception. The filesystem is presumably in a very bad state.
	while (success == false && attempts < fileRetryCount) {
	attempts += 1;
	try {
	syncAllWriters();
	LOG.debug("Data synced to filesystem. Ack'ing [{}] tuples", tupleBatch.size());
	for (Tuple t : tupleBatch) {
	this.collector.ack(t);
	}
	tupleBatch.clear();
	syncPolicy.reset();
	success = true;
	} catch (IOException e) {
	LOG.warn("Data could not be synced to filesystem on attempt [{}]", attempts);
	this.collector.reportError(e);
	lastException = e;
	}
	}

	// If unsuccesful fail the pending tuples
	if (success == false) {
	LOG.warn("Data could not be synced to filesystem, failing this batch of tuples");
	for (Tuple t : tupleBatch) {
	this.collector.fail(t);
	}
	tupleBatch.clear();

	throw new RuntimeException("Sync failed [" + attempts + "] times.", lastException);
	}
	}

	if (writer != null && writer.needsRotation()) {
	doRotationAndRemoveWriter(writerKey, writer);
	}
	}
	}

	private Writer getOrCreateWriter(String writerKey, Tuple tuple) throws IOException {
	Writer writer;

	writer = writers.get(writerKey);
	if (writer == null) {
	Path pathForNextFile = getBasePathForNextFile(tuple);
	writer = makeNewWriter(pathForNextFile, tuple);
	writers.put(writerKey, writer);
	}
	return writer;
	}

	/**
	* A tuple must be mapped to a writer based on two factors.
	* - bolt specific logic that must separate tuples into different files in the same directory (see the avro bolt
	* for an example of this)
	* - the directory the tuple will be partioned into
	*/
	private String getHashKeyForTuple(Tuple tuple) {
	final String boltKey = getWriterKey(tuple);
	final String partitionDir = this.partitioner.getPartitionPath(tuple);
	return boltKey + "****" + partitionDir;
	}

	void doRotationAndRemoveWriter(String writerKey, Writer writer) {
	try {
	rotateOutputFile(writer);
	} catch (IOException e) {
	this.collector.reportError(e);
	LOG.error("File could not be rotated");
	//At this point there is nothing to do. In all likelihood any filesystem operations will fail.
	//The next tuple will almost certainly fail to write and/or sync, which force a rotation. That
	//will give rotateAndReset() a chance to work which includes creating a fresh file handle.
	} finally {
	//rotateOutputFile(writer) has closed the writer. It's safe to remove the writer from the map here.
	writers.remove(writerKey);
	}
	}

	@Override
	public Map<String, Object> getComponentConfiguration() {
	return TupleUtils.putTickFrequencyIntoComponentConfig(super.getComponentConfiguration(), tickTupleInterval);
	}

	@Override
	public void declareOutputFields(OutputFieldsDeclarer outputFieldsDeclarer) {
	}

	@Override
	public void cleanup() {
	doRotationAndRemoveAllWriters();
	if (this.rotationTimer != null) {
	this.rotationTimer.cancel();
	}
	}

	private void doRotationAndRemoveAllWriters() {
	synchronized (writeLock) {
	for (final Writer writer : writers.values()) {
	try {
	rotateOutputFile(writer);
	} catch (IOException e) {
	this.collector.reportError(e);
	LOG.warn("IOException during scheduled file rotation.", e);
	}
	}
	//above for-loop has closed all the writers. It's safe to clear the map here.
	writers.clear();
	}
	}

	private void syncAllWriters() throws IOException {
	for (Writer writer : writers.values()) {
	writer.sync();
	}
	}

	private void startTimedRotationPolicy() {
	long interval = ((TimedRotationPolicy) this.rotationPolicy).getInterval();
	this.rotationTimer = new Timer(true);
	TimerTask task = new TimerTask() {
	@Override
	public void run() {
	doRotationAndRemoveAllWriters();
	}
	};
	this.rotationTimer.scheduleAtFixedRate(task, interval, interval);
	}

	protected Path getBasePathForNextFile(Tuple tuple) {

	final String partitionPath = this.partitioner.getPartitionPath(tuple);
	final int rotation;
	if (rotationCounterMap.containsKey(partitionPath)) {
	rotation = rotationCounterMap.get(partitionPath) + 1;
	} else {
	rotation = 0;
	}
	rotationCounterMap.put(partitionPath, rotation);

	return new Path(this.fsUrl + this.fileNameFormat.getPath() + partitionPath,
	this.fileNameFormat.getName(rotation, System.currentTimeMillis()));
	}

	protected abstract void doPrepare(Map<String, Object> conf, TopologyContext topologyContext, OutputCollector collector) throws
	IOException;

	protected abstract String getWriterKey(Tuple tuple);

	protected abstract Writer makeNewWriter(Path path, Tuple tuple) throws IOException;

	static class WritersMap extends LinkedHashMap<String, Writer> {
	final long maxWriters;
	final OutputCollector collector;

	public WritersMap(long maxWriters, OutputCollector collector) {
	super((int) maxWriters, 0.75f, true);
	this.maxWriters = maxWriters;
	this.collector = collector;
	}

	@Override
	protected boolean removeEldestEntry(Map.Entry<String, Writer> eldest) {
	if (this.size() > this.maxWriters) {
	//The writer must be closed before removed from the map.
	//If it failed, we might lose some data.
	try {
	eldest.getValue().close();
	} catch (IOException e) {
	collector.reportError(e);
	LOG.error("Failed to close the eldest Writer");
	}
	return true;
	} else {
	return false;
	}
	}
	}
	}