Merge branch 'master' of https://github.com/apache/samza

Author: sborya

samza-api/src/main/java/org/apache/samza/storage/kv/KeyValueStore.java

@@ -46,7 +46,19 @@ public interface KeyValueStore<K, V> {
    * @return a map of the keys that were found and their respective values.
    * @throws NullPointerException if the specified {@code keys} list, or any of the keys, is {@code null}.
    */
-  Map<K, V> getAll(List<K> keys);
+  default Map<K, V> getAll(List<K> keys) {
+    Map<K, V> map = new HashMap<>(keys.size());
+
+    for (K key : keys) {
+      V value = get(key);
+
+      if (value != null) {
+        map.put(key, value);
+      }
+    }
+
+    return map;
+  }
 
   /**
    * Updates the mapping of the specified key-value pair; Associates the specified {@code key} with the specified {@code value}.
@@ -79,7 +91,11 @@ public interface KeyValueStore<K, V> {
    * @param keys the keys for which the mappings are to be deleted.
    * @throws NullPointerException if the specified {@code keys} list, or any of the keys, is {@code null}.
    */
-  void deleteAll(List<K> keys);
+  default void deleteAll(List<K> keys) {
+    for (K key : keys) {
+      delete(key);
+    }
+  }
 
   /**
    * Returns an iterator for a sorted range of entries specified by [{@code from}, {@code to}).
@@ -111,53 +127,4 @@ public interface KeyValueStore<K, V> {
    * Flushes this key-value store, if applicable.
    */
   void flush();
-
-  /**
-   * Represents an extension for classes that implement {@link KeyValueStore}.
-   */
-  // TODO replace with default interface methods when we can use Java 8 features.
-  class Extension {
-    private Extension() {
-      // This class cannot be instantiated
-    }
-
-    /**
-     * Gets the values with which the specified {@code keys} are associated.
-     *
-     * @param store the key-value store for which this operation is to be performed.
-     * @param keys the keys with which the associated values are to be fetched.
-     * @param <K> the type of keys maintained by the specified {@code store}.
-     * @param <V> the type of values maintained by the specified {@code store}.
-     * @return a map of the keys that were found and their respective values.
-     * @throws NullPointerException if the specified {@code keys} list, or any of the keys, is {@code null}.
-     */
-    public static <K, V> Map<K, V> getAll(final KeyValueStore<K, V> store, final List<K> keys) {
-      final Map<K, V> map = new HashMap<>(keys.size());
-
-      for (final K key : keys) {
-        final V value = store.get(key);
-
-        if (value != null) {
-          map.put(key, value);
-        }
-      }
-
-      return map;
-    }
-
-    /**
-     * Deletes the mappings for the specified {@code keys} from this key-value store (if such mappings exist).
-     *
-     * @param store the key-value store for which this operation is to be performed.
-     * @param keys the keys for which the mappings are to be deleted.
-     * @param <K> the type of keys maintained by the specified {@code store}.
-     * @param <V> the type of values maintained by the specified {@code store}.
-     * @throws NullPointerException if the specified {@code keys} list, or any of the keys, is {@code null}.
-     */
-    public static <K, V> void deleteAll(final KeyValueStore<K, V> store, final List<K> keys) {
-      for (final K key : keys) {
-        store.delete(key);
-      }
-    }
-  }
 }

samza-kafka/src/main/scala/org/apache/samza/checkpoint/kafka/KafkaCheckpointManager.scala

@@ -43,6 +43,9 @@ import scala.collection.mutable
  * keyed to that taskName. If there is no such message, no checkpoint data
  * exists.  The underlying log has a single partition into which all
  * checkpoints and TaskName to changelog partition mappings are written.
+ *
+ * This class is thread safe for writing but not for reading checkpoints.
+ * This is currently OK since checkpoints are only read on the main thread.
  */
 class KafkaCheckpointManager(
                               clientId: String,
@@ -64,14 +67,14 @@ class KafkaCheckpointManager(
                               checkpointTopicProperties: Properties = new Properties) extends CheckpointManager with Logging {
 
   var taskNames = Set[TaskName]()
-  @volatile var  systemProducer: SystemProducer = null
+  @volatile var systemProducer: SystemProducer = null
+  var systemConsumer: SystemConsumer = null
   var taskNamesToOffsets: Map[TaskName, Checkpoint] = null
   val systemAdmin = getSystemAdmin()
 
   val kafkaUtil: KafkaUtil = new KafkaUtil(retryBackoff, connectZk)
 
 
-
   KafkaCheckpointLogKey.setSystemStreamPartitionGrouperFactoryString(systemStreamPartitionGrouperFactoryString)
 
   info("Creating KafkaCheckpointManager with: clientId=%s, checkpointTopic=%s, systemName=%s" format(clientId, checkpointTopic, systemName))
@@ -181,73 +184,40 @@ class KafkaCheckpointManager(
    */
   private def readLog(shouldHandleEntry: (KafkaCheckpointLogKey) => Boolean,
                       handleEntry: (ByteBuffer, KafkaCheckpointLogKey) => Unit): Unit = {
-
-    val UNKNOWN_OFFSET = "-1"
-    var attempts = 10
-    val POLL_TIMEOUT = 1000L
+    info("Reading from checkpoint system:%s topic:%s" format(systemName, checkpointTopic))
 
     val ssp: SystemStreamPartition = new SystemStreamPartition(systemName, checkpointTopic, new Partition(0))
-    val systemConsumer = getSystemConsumer()
-    val partitionMetadata = getSSPMetadata(checkpointTopic, new Partition(0))
-    // offsets returned are strings
-    val newestOffset = if (partitionMetadata.getNewestOffset == null) UNKNOWN_OFFSET else partitionMetadata.getNewestOffset
-    val oldestOffset = partitionMetadata.getOldestOffset
-    systemConsumer.register(ssp, oldestOffset) // checkpoint stream should always be read from the beginning
-    systemConsumer.start()
 
-    var msgCount = 0
-    try {
-      val emptyEnvelopes = util.Collections.emptyMap[SystemStreamPartition, java.util.List[IncomingMessageEnvelope]]
-      // convert offsets to long
-      var currentOffset = UNKNOWN_OFFSET.toLong
-      val newestOffsetLong = newestOffset.toLong
-      val sspToPoll = Collections.singleton(ssp)
-      while (currentOffset < newestOffsetLong) {
-
-        val envelopes: java.util.Map[SystemStreamPartition, java.util.List[IncomingMessageEnvelope]] =
-        try {
-          systemConsumer.poll(sspToPoll, POLL_TIMEOUT)
-        } catch {
-          case e: Exception => {
-            // these exceptions are most likely intermediate
-            warn("Got %s exception while polling the consumer for checkpoints." format e)
-            if (attempts == 0) throw new SamzaException("Multiple attempts failed while reading the checkpoints. Giving up.", e)
-            attempts -= 1
-            emptyEnvelopes
-          }
-        }
+    if (systemConsumer == null) {
+      val partitionMetadata = getSSPMetadata(checkpointTopic, new Partition(0))
+      val oldestOffset = partitionMetadata.getOldestOffset
 
-        val messages: util.List[IncomingMessageEnvelope] = envelopes.get(ssp)
-        val messagesNum = if (messages != null) messages.size else 0
-        info("CheckpointMgr read %s envelopes (%s messages) from ssp %s. Current offset is %s, newest is %s"
-                     format (envelopes.size(), messagesNum, ssp, currentOffset, newestOffset))
-        if (envelopes.isEmpty || messagesNum <= 0) {
-          info("Got empty/null list of messages")
-        } else {
-          msgCount += messages.size()
-          // check the key
-          for (msg: IncomingMessageEnvelope <- messages) {
-            val key = msg.getKey.asInstanceOf[Array[Byte]]
-            currentOffset = msg.getOffset().toLong
-            if (key == null) {
-              throw new KafkaUtilException("While reading checkpoint (currentOffset=%s) stream encountered message without key."
-                                                   format currentOffset)
-            }
+      systemConsumer = getSystemConsumer()
+      systemConsumer.register(ssp, oldestOffset)
+      systemConsumer.start()
+    }
 
-            val checkpointKey = KafkaCheckpointLogKey.fromBytes(key)
+    val iterator =  new SystemStreamPartitionIterator(systemConsumer, ssp);
+    var msgCount = 0
+    while (iterator.hasNext) {
+      val msg = iterator.next
+      msgCount += 1
+
+      val offset = msg.getOffset
+      val key = msg.getKey.asInstanceOf[Array[Byte]]
+      if (key == null) {
+        throw new KafkaUtilException(
+          "While reading checkpoint (currentOffset=%s) stream encountered message without key." format offset)
+      }
 
-            if (!shouldHandleEntry(checkpointKey)) {
-              info("Skipping checkpoint log entry at offset %s with key %s." format(currentOffset, checkpointKey))
-            } else {
-              // handleEntry requires ByteBuffer
-              val checkpointPayload = ByteBuffer.wrap(msg.getMessage.asInstanceOf[Array[Byte]])
-              handleEntry(checkpointPayload, checkpointKey)
-            }
-          }
-        }
+      val checkpointKey = KafkaCheckpointLogKey.fromBytes(key)
+
+      if (!shouldHandleEntry(checkpointKey)) {
+        info("Skipping checkpoint log entry at offset %s with key %s." format(offset, checkpointKey))
+      } else {
+        val checkpointPayload = ByteBuffer.wrap(msg.getMessage.asInstanceOf[Array[Byte]])
+        handleEntry(checkpointPayload, checkpointKey)
       }
-    } finally {
-      systemConsumer.stop()
     }
     info("Done reading %s messages from checkpoint system:%s topic:%s" format(msgCount, systemName, checkpointTopic))
   }
@@ -282,12 +252,17 @@ class KafkaCheckpointManager(
 
 
   def stop = {
-    synchronized (
+    synchronized {
       if (systemProducer != null) {
         systemProducer.stop
         systemProducer = null
       }
-    )
+
+      if (systemConsumer != null) {
+        systemConsumer.stop
+        systemConsumer = null
+      }
+    }
 
   }
 

samza-kv-inmemory/src/main/scala/org/apache/samza/storage/kv/inmemory/InMemoryKeyValueStore.scala

@@ -81,10 +81,6 @@ class InMemoryKeyValueStore(val metrics: KeyValueStoreMetrics = new KeyValueStor
     put(key, null)
   }
 
-  override def deleteAll(keys: java.util.List[Array[Byte]]) = {
-    KeyValueStore.Extension.deleteAll(this, keys)
-  }
-
   override def putAll(entries: util.List[Entry[Array[Byte], Array[Byte]]]): Unit = {
     // TreeMap's putAll requires a map, so we'd need to iterate over all the entries anyway
     // to use it, in order to putAll here.  Therefore, just iterate here.
@@ -116,8 +112,4 @@ class InMemoryKeyValueStore(val metrics: KeyValueStoreMetrics = new KeyValueStor
     }
     found
   }
-
-  override def getAll(keys: java.util.List[Array[Byte]]): java.util.Map[Array[Byte], Array[Byte]] = {
-    KeyValueStore.Extension.getAll(this, keys);
-  }
 }

samza-kv-rocksdb/src/main/scala/org/apache/samza/storage/kv/RocksDbKeyValueStore.scala

@@ -20,89 +20,71 @@
 package org.apache.samza.storage.kv
 
 import java.io.File
+import java.util.concurrent.TimeUnit
+
 import org.apache.samza.SamzaException
-import org.apache.samza.util.{ LexicographicComparator, Logging }
 import org.apache.samza.config.Config
-import org.rocksdb._
-import org.rocksdb.TtlDB
+import org.apache.samza.util.{LexicographicComparator, Logging}
+import org.rocksdb.{TtlDB, _}
 
 object RocksDbKeyValueStore extends Logging {
 
-  def openDB(dir: File, options: Options, storeConfig: Config, isLoggedStore: Boolean, storeName: String, metrics: KeyValueStoreMetrics): RocksDB = {
+  def openDB(dir: File, options: Options, storeConfig: Config, isLoggedStore: Boolean,
+             storeName: String, metrics: KeyValueStoreMetrics): RocksDB = {
     var ttl = 0L
     var useTTL = false
 
-    if (storeConfig.containsKey("rocksdb.ttl.ms"))
-    {
-      try
-      {
+    if (storeConfig.containsKey("rocksdb.ttl.ms")) {
+      try {
         ttl = storeConfig.getLong("rocksdb.ttl.ms")
 
-        // RocksDB accepts TTL in seconds, convert ms to seconds
-        if(ttl > 0) {
-          if (ttl < 1000)
-          {
-            warn("The ttl values requested for %s is %d, which is less than 1000 (minimum), using 1000 instead",
-              storeName,
-              ttl)
+        if (ttl > 0) {
+          if (ttl < 1000) {
+            warn("The ttl value requested for %s is %d which is less than 1000 (minimum). " +
+              "Using 1000 ms instead.", storeName, ttl)
             ttl = 1000
           }
-          ttl = ttl / 1000
-        }
-        else {
-          warn("Non-positive TTL for RocksDB implies infinite TTL for the data. More Info -https://github.com/facebook/rocksdb/wiki/Time-to-Live")
+          ttl = TimeUnit.MILLISECONDS.toSeconds(ttl)
+        } else {
+          warn("Non-positive TTL for RocksDB implies infinite TTL for the data. " +
+            "More Info - https://github.com/facebook/rocksdb/wiki/Time-to-Live")
         }
 
         useTTL = true
-        if (isLoggedStore)
-        {
-          warn("%s is a TTL based store, changelog is not supported for TTL based stores, use at your own discretion" format storeName)
+        if (isLoggedStore) {
+          warn("%s is a TTL based store. Changelog is not supported for TTL based stores. " +
+            "Use at your own discretion." format storeName)
         }
+      } catch {
+        case nfe: NumberFormatException =>
+          throw new SamzaException("rocksdb.ttl.ms configuration value %s for store %s is not a number."
+            format (storeConfig.get("rocksdb.ttl.ms"), storeName), nfe)
       }
-      catch
-        {
-          case nfe: NumberFormatException => throw new SamzaException("rocksdb.ttl.ms configuration is not a number, " + "value found %s" format storeConfig.get(
-            "rocksdb.ttl.ms"))
-        }
     }
 
-    try
-    {
+    try {
       val rocksDb =
-        if (useTTL)
-        {
+        if (useTTL) {
           info("Opening RocksDB store with TTL value: %s" format ttl)
           TtlDB.open(options, dir.toString, ttl.toInt, false)
-        }
-        else
-        {
+        } else {
           RocksDB.open(options, dir.toString)
         }
 
-      if (storeConfig.containsKey("rocksdb.metrics.list"))
-      {
+      if (storeConfig.containsKey("rocksdb.metrics.list")) {
         storeConfig
           .get("rocksdb.metrics.list")
           .split(",")
           .map(property => property.trim)
-          .foreach(property =>
-            metrics.newGauge(property, () => rocksDb.getProperty(property))
-          )
+          .foreach(property => metrics.newGauge(property, () => rocksDb.getProperty(property)))
       }
 
       rocksDb
+    } catch {
+      case rocksDBException: RocksDBException =>
+        throw new SamzaException("Error opening RocksDB store %s at location %s" format (storeName, dir.toString),
+          rocksDBException)
     }
-    catch
-      {
-        case rocksDBException: RocksDBException =>
-        {
-          throw new SamzaException(
-            "Error opening RocksDB store %s at location %s, received the following exception from RocksDB %s".format(
-              storeName,
-              dir.toString,
-              rocksDBException))
-        }
-      }
   }
 }
 
@@ -187,10 +169,6 @@ class RocksDbKeyValueStore(
     put(key, null)
   }
 
-  def deleteAll(keys: java.util.List[Array[Byte]]) = {
-    KeyValueStore.Extension.deleteAll(this, keys)
-  }
-
   def range(from: Array[Byte], to: Array[Byte]): KeyValueIterator[Array[Byte], Array[Byte]] = {
     metrics.ranges.inc
     require(from != null && to != null, "Null bound not allowed.")