samza-test/src/main/scala/org/apache/samza/test/performance/TestKeyValuePerformance.scala - samza - 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.samza.test.performance

 import java.io.File
 import java.util
 import java.util.concurrent.TimeUnit
 import java.util.Collections
 import java.util.UUID

 import com.google.common.base.Stopwatch
 import com.google.common.collect.ImmutableList
 import com.google.common.collect.ImmutableMap
 import org.apache.commons.lang3.RandomStringUtils
 import org.apache.samza.config.{Config, JobConfig, MapConfig, StorageConfig}
 import org.apache.samza.container.TaskName
 import org.apache.samza.context.ContainerContextImpl
 import org.apache.samza.context.JobContextImpl
 import org.apache.samza.job.model.ContainerModel
 import org.apache.samza.job.model.TaskModel
 import org.apache.samza.metrics.MetricsRegistryMap
 import org.apache.samza.serializers.ByteSerde
 import org.apache.samza.serializers.SerdeManager
 import org.apache.samza.serializers.UUIDSerde
 import org.apache.samza.storage.StorageEngineFactory
 import org.apache.samza.storage.StorageEngineFactory.StoreMode
 import org.apache.samza.storage.kv.KeyValueStorageEngine
 import org.apache.samza.storage.kv.KeyValueStore
 import org.apache.samza.system.SystemProducer
 import org.apache.samza.system.SystemProducers
 import org.apache.samza.system.SystemStreamPartition
 import org.apache.samza.task.TaskInstanceCollector
 import org.apache.samza.util.{CommandLine, FileUtil, Logging, ReflectionUtil}
 import org.apache.samza.Partition
 import org.apache.samza.SamzaException
 import org.apache.samza.util.ScalaJavaUtil.JavaOptionals

 import scala.collection.JavaConverters._
 import scala.util.Random

 /**
  * A simple CLI-based tool for running various key-value performance tests.
  *
  * List of KeyValuePerformance tests must be defined in 'test.methods' configuration as a comma-separated value.
  * The tool splits this list to determine which tests to run.
  *
  * Each test should define its own set of configuration for partition count, stores etc.
  * using the "test.<test-name>.<config-string>=<config-value>" pattern
  *
  * Each test may define one or more test parameterss.
  * For example, test1 can define 2 sets of parameters by specifying "test.test1.set.count=2" and
  * define each set as:
  * "test.test1.set-1.<param-name>=<param-value>"
  * "test.test1.set-2.<param-name>=<param-value>"
  */

 object TestKeyValuePerformance extends Logging {
   val Encoding = "UTF-8"
   val JobId = RandomStringUtils.random(10)

   val testMethods: Map[String, (KeyValueStorageEngine[Array[Byte], Array[Byte]], Config) => Unit] = Map(
     "all-with-deletes" -> runTestAllWithDeletes,
     "rocksdb-write-performance" -> runTestMsgWritePerformance,
     "rocksdb-concurrent-write-performance" -> runTestConcurrentMsgWritePerformance,
     "get-all-vs-get-write-many-read-many" -> runTestGetAllVsGetWriteManyReadMany,
     "get-all-vs-get-write-once-read-many" -> runTestGetAllVsGetWriteOnceReadMany)

   def main(args: Array[String]) {
     val cmdline = new CommandLine
     val options = cmdline.parser.parse(args: _*)
     val config = cmdline.loadConfig(options)
     val tests = config.get("test.methods").split(",")

     tests.foreach{ test =>
       info("Running test: %s" format test)
       if(testMethods.contains(test)) {
         val testConfig: util.Map[String, String] = new MapConfig(config.subset("test." + test + ".", true))
         val jobConfig: util.Map[String, String] = ImmutableMap.of(JobConfig.JOB_NAME, test, JobConfig.JOB_ID, JobId)
         val combinedConfig: Config = new MapConfig(ImmutableList.of(testConfig, jobConfig))
         invokeTest(test, testMethods(test), combinedConfig)
       } else {
         error("Invalid test method. valid methods are: %s" format testMethods.keys)
         throw new SamzaException("Unknown test method: %s" format test)
       }
     }
   }

   def invokeTest(testName: String, testMethod: (KeyValueStorageEngine[Array[Byte], Array[Byte]], Config) => Unit, config: Config) {
     val partitionCount = config.getInt("partition.count", 1)
     val tasks = (0 until partitionCount)
       .map(i => new Partition(i))
       .map(partition => (new TaskName(partition.toString),
         new TaskModel(new TaskName(partition.toString),
           Collections.singleton(new SystemStreamPartition("system", "stream", partition)),
           partition)))
       .toMap

     val producerMultiplexer = new SystemProducers(
       Map[String, SystemProducer](),
       new SerdeManager
     )
     val storageConfig = new StorageConfig(config)
     // Build a Map[String, StorageEngineFactory]. The key is the store name.
     val storageEngineMappings = storageConfig
       .getStoreNames.asScala
       .map(storeName => {
         val storageFactoryClassName =
           JavaOptionals.toRichOptional(storageConfig.getStorageFactoryClassName(storeName)).toOption
                 .getOrElse(throw new SamzaException("Missing storage factory for %s." format storeName))
         (storeName, ReflectionUtil.getObj(storageFactoryClassName,
           classOf[StorageEngineFactory[Array[Byte], Array[Byte]]]))
     })

     for((storeName, storageEngine) <- storageEngineMappings) {
       val testSetCount = storageConfig.getInt("set.count", 1)
       (1 to testSetCount).foreach(testSet => {
         //Create a new DB instance for each test set
         val output = new File("/tmp/" + UUID.randomUUID())
         val byteSerde = new ByteSerde
         val engine = storageEngine.getStorageEngine(
           storeName,
           output,
           byteSerde,
           byteSerde,
           new TaskInstanceCollector(producerMultiplexer),
           new MetricsRegistryMap,
           null,
           JobContextImpl.fromConfigWithDefaults(storageConfig, null),
           new ContainerContextImpl(new ContainerModel("0", tasks.asJava), new MetricsRegistryMap), StoreMode.ReadWrite
         )

         val db = if(!engine.isInstanceOf[KeyValueStorageEngine[_,_]]) {
           throw new SamzaException("This test can only run with KeyValueStorageEngine configured as store factory.")
         } else {
           engine.asInstanceOf[KeyValueStorageEngine[Array[Byte], Array[Byte]]]
         }

         // Run the test method
         testMethod(db, storageConfig.subset("set-" + testSet + ".", true))

         new FileUtil().rm(output)
       })
     }
   }

   def runTestAllWithDeletes(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
     val numLoops = config.getInt("num.loops", 100)
     val messagesPerBatch = config.getInt("messages.per.batch", 10000)
     val messageSizeBytes = config.getInt("message.size.bytes", 200)

     info("Using (num loops, messages per batch, message size in bytes) => (%s, %s, %s)" format (numLoops, messagesPerBatch, messageSizeBytes))
     new TestKeyValuePerformance().testAllWithDeletes(db, numLoops, messagesPerBatch, messageSizeBytes)
   }

   def runTestMsgWritePerformance(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
     val messageSizeBytes = config.getInt("message.size", 200)
     val messageCount = config.getInt("message.count", 10000)

     info("Using (message count, message size in bytes) => (%s, %s)" format (messageCount, messageSizeBytes))
     new TestKeyValuePerformance().testMsgWritePerformance(db, messageCount, messageSizeBytes)
   }

   def runTestConcurrentMsgWritePerformance(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
     val messageSizeBytes = config.getInt("message.size", 200)
     val messageCount = config.getInt("message.count", 100000)
     val numThreads = config.getInt("num.threads", 4)

     new TestKeyValuePerformance().testConcurrentMsgWritePerformance(db, messageCount, messageSizeBytes, numThreads)
   }

   def runTestGetAllVsGetWriteManyReadMany(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
     new TestKeyValuePerformance().testGetAllVsGetWriteManyReadMany(db, config)
   }

   def runTestGetAllVsGetWriteOnceReadMany(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
     new TestKeyValuePerformance().testGetAllVsGetWriteOnceReadMany(db, config)
   }
 }

 class TestKeyValuePerformance extends Logging {
   import TestKeyValuePerformance._

   /**
    * A test that writes messagesPerBatch messages, deletes them all, then calls
    * store.all. The test periodically outputs the time it takes to complete
    * these operations. This test is useful to trouble shoot issues with LevleDB
    * such as the issue documented in SAMZA-254.
    */
   def testAllWithDeletes(
     store: KeyValueStore[Array[Byte], Array[Byte]],

     /**
      * How many times a batch of messages should be written and deleted.
      */
     numLoops: Int = 100,

     /**
      * The number of messages to write and delete per-batch.
      */
     messagesPerBatch: Int = 10000,

     /**
      * The size of the messages to write.
      */
     messageSizeBytes: Int = 200) {

     val stuff = (0 until messageSizeBytes).map(i => "a").mkString.getBytes(Encoding)
     val start = System.currentTimeMillis

     (0 until numLoops).foreach(i => {
       info("(%sms) Total written to store: %s" format (System.currentTimeMillis - start, i * messagesPerBatch))

       (0 until messagesPerBatch).foreach(j => {
         val k = (i * j).toString.getBytes(Encoding)
         store.put(k, stuff)
         store.delete(k)
       })

       val allStart = System.currentTimeMillis
       val iter = store.all
       info("(%sms) all() took %sms." format (System.currentTimeMillis - start, System.currentTimeMillis - allStart))
       iter.close
     })

     info("Total time: %ss" format ((System.currentTimeMillis - start) * .001))
   }

   /**
    * Test that successively writes a set of fixed-size messages to the KV store
    * and computes the total time for the operations
    * @param store Key-Value store instance that is being tested
    * @param numMsgs Total number of messages to write to the store
    * @param msgSizeInBytes Size of each message in Bytes
    */
   def testMsgWritePerformance(
     store: KeyValueStore[Array[Byte], Array[Byte]],
     numMsgs: Int = 10000,
     msgSizeInBytes: Int = 200) {

     val msg = (0 until msgSizeInBytes).map(i => "x").mkString.getBytes(Encoding)

     val start = System.currentTimeMillis
     (0 until numMsgs).foreach(i => {
       store.put(i.toString.getBytes(Encoding), msg)
     })
     val timeTaken = System.currentTimeMillis - start
     info("Total time to write %d msgs of size %d bytes : %s s" format (numMsgs, msgSizeInBytes, timeTaken * .001))
   }

   def testConcurrentMsgWritePerformance(
     store: KeyValueStore[Array[Byte], Array[Byte]],
     numMsgs: Int = 100000,
     msgSizeInBytes: Int = 200,
     numThreads: Int = 4) {

     val msg = (0 until msgSizeInBytes).map(i => "x").mkString.getBytes(Encoding)
     def createThread(name: String): Thread = new Thread(new Runnable {
       override def run() = {
         (0 until numMsgs).foreach(i => {
           store.put(i.toString.getBytes(Encoding), msg)
         })
       }
     }, name)

     val threads = (0 until numThreads).map(i => createThread(s"Writer $i"))

     val start = System.currentTimeMillis
     threads.foreach(_.start())
     threads.foreach(_.join())
     val timeTaken = System.currentTimeMillis - start

     info("Total time to write %d msgs of size %d bytes with %s threads: %s sec"
       format (numMsgs, msgSizeInBytes, numThreads, timeTaken * .001))
   }

   /**
    * Test that ::getAll performance is better than that of ::get (test when there are many writes and many reads).
    * @param store key-value store instance that is being tested
    * @param config the test case's config
    */
   def testGetAllVsGetWriteManyReadMany(store: KeyValueStore[Array[Byte],Array[Byte]], config: Config): Unit = {
     val iterationsCount = config.getInt("iterations.count", 100)
     val maxMessagesCountPerBatch = config.getInt("message.max-count-per-batch", 100000)
     val maxMessageSizeBytes = config.getInt("message.max-size.bytes", 1024)
     val timer = Stopwatch.createUnstarted
     val uuidSerde = new UUIDSerde

     info("iterations count: " + iterationsCount)
     info("max messages count per batch: " + maxMessagesCountPerBatch)
     info("max message size in bytes: " + maxMessageSizeBytes)
     info("%12s%12s%12s%12s".format("Msg Count", "Bytes/Msg", "get ms", "getAll ms"))

     try {
       (0 until iterationsCount).foreach(i => {
         val messageSizeBytes = Random.nextInt(maxMessageSizeBytes)
         val messagesCountPerBatch = Random.nextInt(maxMessagesCountPerBatch)
         val keys = (0 until messagesCountPerBatch).map(k => uuidSerde.toBytes(UUID.randomUUID)).toList
         val shuffledKeys = Random.shuffle(keys) // to reduce locality of reference -- sequential access may be unfair

         keys.foreach(k => store.put(k, Random.nextString(messageSizeBytes).getBytes(Encoding)))
         store.flush()

         timer.reset().start()
         assert(store.getAll(shuffledKeys.asJava).size == shuffledKeys.size)
         val getAllTime = timer.stop().elapsed(TimeUnit.MILLISECONDS)

         // Restore cache, in case it's enabled, to a state similar to the one above when the getAll test started
         keys.foreach(k => store.put(k, Random.nextString(messageSizeBytes).getBytes(Encoding)))
         store.flush()

         timer.reset().start()
         shuffledKeys.foreach(store.get)
         val getTime = timer.stop().elapsed(TimeUnit.MILLISECONDS)

         info("%12d%12d%12d%12d".format(messagesCountPerBatch, messageSizeBytes, getTime, getAllTime))
         if (getAllTime > getTime) {
           error("getAll was slower than get!")
         }
       })
     } finally {
       store.close()
     }
   }

   /**
    * Test that ::getAll performance is better than that of ::get (test when data are written once and read many times);
    * load is usually greater than the storage engine's cache size (not to be confused with Samza's cache layer),
    * and keys are randomly selected from the stored entries to perform a fair comparison of ::get vs. ::getAll (in case
    * the underlying storage engine caches data in blocks and ::getAll causes a block to be loaded into the cache --
    * one can argue that ::get should trigger the same behavior, but it's worth testing this WORM scenario regardless)
    * @param store key-value store instance that is being tested
    * @param config the test case's config
    */
   def testGetAllVsGetWriteOnceReadMany(store: KeyValueStore[Array[Byte],Array[Byte]], config: Config): Unit = {
     val iterationsCount = config.getInt("iterations.count", 100)
     val maxMessagesCountPerBatch = config.getInt("message.max-count-per-batch", 10000 + Random.nextInt(20000))
     val maxMessageSizeBytes = config.getInt("message.max-size.bytes", 1024)
     val totalMessagesCount = iterationsCount * maxMessagesCountPerBatch
     val timer = Stopwatch.createUnstarted
     val uuidSerde = new UUIDSerde

     info("write once -- putting %d messages in store".format(totalMessagesCount))
     val keys = (0 until totalMessagesCount).map(k => uuidSerde.toBytes(UUID.randomUUID)).toList
     keys.foreach(k => store.put(k, Random.nextString(Random.nextInt(maxMessageSizeBytes)).getBytes(Encoding)))
     store.flush()

     info("iterations count: " + iterationsCount)
     info("max messages count per batch: " + maxMessagesCountPerBatch)
     info("max message size in bytes: " + maxMessageSizeBytes)
     info("%12s%12s%12s%12s".format("Msg Count", "Total Size", "get ms", "getAll ms"))

     try {
       (0 until iterationsCount).foreach(i => {
         val messagesCountPerBatch = Random.nextInt(maxMessagesCountPerBatch)
         val shuffledKeys = Random.shuffle(keys).take(messagesCountPerBatch)

         // We want to measure ::getAll when called many times, so populate the cache because first call is a cache-miss
         val totalSize = store.getAll(shuffledKeys.asJava).values.asScala.map(_.length).sum
         timer.reset().start()
         assert(store.getAll(shuffledKeys.asJava).size == shuffledKeys.size)
         val getAllTime = timer.stop().elapsed(TimeUnit.MILLISECONDS)

         // We want to measure ::get when called many times, so populate the cache because first call is a cache-miss
         shuffledKeys.foreach(store.get)
         timer.reset().start()
         shuffledKeys.foreach(store.get)
         val getTime = timer.stop().elapsed(TimeUnit.MILLISECONDS)

         info("%12d%12d%12d%12d".format(messagesCountPerBatch, totalSize, getTime, getAllTime))
         if (getAllTime > getTime) {
           error("getAll was slower than get!")
         }
       })
     } finally {
       store.close()
     }
   }
 }
	/*
	* 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.samza.test.performance

	import java.io.File
	import java.util
	import java.util.concurrent.TimeUnit
	import java.util.Collections
	import java.util.UUID

	import com.google.common.base.Stopwatch
	import com.google.common.collect.ImmutableList
	import com.google.common.collect.ImmutableMap
	import org.apache.commons.lang3.RandomStringUtils
	import org.apache.samza.config.{Config, JobConfig, MapConfig, StorageConfig}
	import org.apache.samza.container.TaskName
	import org.apache.samza.context.ContainerContextImpl
	import org.apache.samza.context.JobContextImpl
	import org.apache.samza.job.model.ContainerModel
	import org.apache.samza.job.model.TaskModel
	import org.apache.samza.metrics.MetricsRegistryMap
	import org.apache.samza.serializers.ByteSerde
	import org.apache.samza.serializers.SerdeManager
	import org.apache.samza.serializers.UUIDSerde
	import org.apache.samza.storage.StorageEngineFactory
	import org.apache.samza.storage.StorageEngineFactory.StoreMode
	import org.apache.samza.storage.kv.KeyValueStorageEngine
	import org.apache.samza.storage.kv.KeyValueStore
	import org.apache.samza.system.SystemProducer
	import org.apache.samza.system.SystemProducers
	import org.apache.samza.system.SystemStreamPartition
	import org.apache.samza.task.TaskInstanceCollector
	import org.apache.samza.util.{CommandLine, FileUtil, Logging, ReflectionUtil}
	import org.apache.samza.Partition
	import org.apache.samza.SamzaException
	import org.apache.samza.util.ScalaJavaUtil.JavaOptionals

	import scala.collection.JavaConverters._
	import scala.util.Random

	/**
	* A simple CLI-based tool for running various key-value performance tests.
	*
	* List of KeyValuePerformance tests must be defined in 'test.methods' configuration as a comma-separated value.
	* The tool splits this list to determine which tests to run.
	*
	* Each test should define its own set of configuration for partition count, stores etc.
	* using the "test.<test-name>.<config-string>=<config-value>" pattern
	*
	* Each test may define one or more test parameterss.
	* For example, test1 can define 2 sets of parameters by specifying "test.test1.set.count=2" and
	* define each set as:
	* "test.test1.set-1.<param-name>=<param-value>"
	* "test.test1.set-2.<param-name>=<param-value>"
	*/

	object TestKeyValuePerformance extends Logging {
	val Encoding = "UTF-8"
	val JobId = RandomStringUtils.random(10)

	val testMethods: Map[String, (KeyValueStorageEngine[Array[Byte], Array[Byte]], Config) => Unit] = Map(
	"all-with-deletes" -> runTestAllWithDeletes,
	"rocksdb-write-performance" -> runTestMsgWritePerformance,
	"rocksdb-concurrent-write-performance" -> runTestConcurrentMsgWritePerformance,
	"get-all-vs-get-write-many-read-many" -> runTestGetAllVsGetWriteManyReadMany,
	"get-all-vs-get-write-once-read-many" -> runTestGetAllVsGetWriteOnceReadMany)

	def main(args: Array[String]) {
	val cmdline = new CommandLine
	val options = cmdline.parser.parse(args: _*)
	val config = cmdline.loadConfig(options)
	val tests = config.get("test.methods").split(",")

	tests.foreach{ test =>
	info("Running test: %s" format test)
	if(testMethods.contains(test)) {
	val testConfig: util.Map[String, String] = new MapConfig(config.subset("test." + test + ".", true))
	val jobConfig: util.Map[String, String] = ImmutableMap.of(JobConfig.JOB_NAME, test, JobConfig.JOB_ID, JobId)
	val combinedConfig: Config = new MapConfig(ImmutableList.of(testConfig, jobConfig))
	invokeTest(test, testMethods(test), combinedConfig)
	} else {
	error("Invalid test method. valid methods are: %s" format testMethods.keys)
	throw new SamzaException("Unknown test method: %s" format test)
	}
	}
	}

	def invokeTest(testName: String, testMethod: (KeyValueStorageEngine[Array[Byte], Array[Byte]], Config) => Unit, config: Config) {
	val partitionCount = config.getInt("partition.count", 1)
	val tasks = (0 until partitionCount)
	.map(i => new Partition(i))
	.map(partition => (new TaskName(partition.toString),
	new TaskModel(new TaskName(partition.toString),
	Collections.singleton(new SystemStreamPartition("system", "stream", partition)),
	partition)))
	.toMap

	val producerMultiplexer = new SystemProducers(
	Map[String, SystemProducer](),
	new SerdeManager
	)
	val storageConfig = new StorageConfig(config)
	// Build a Map[String, StorageEngineFactory]. The key is the store name.
	val storageEngineMappings = storageConfig
	.getStoreNames.asScala
	.map(storeName => {
	val storageFactoryClassName =
	JavaOptionals.toRichOptional(storageConfig.getStorageFactoryClassName(storeName)).toOption
	.getOrElse(throw new SamzaException("Missing storage factory for %s." format storeName))
	(storeName, ReflectionUtil.getObj(storageFactoryClassName,
	classOf[StorageEngineFactory[Array[Byte], Array[Byte]]]))
	})

	for((storeName, storageEngine) <- storageEngineMappings) {
	val testSetCount = storageConfig.getInt("set.count", 1)
	(1 to testSetCount).foreach(testSet => {
	//Create a new DB instance for each test set
	val output = new File("/tmp/" + UUID.randomUUID())
	val byteSerde = new ByteSerde
	val engine = storageEngine.getStorageEngine(
	storeName,
	output,
	byteSerde,
	byteSerde,
	new TaskInstanceCollector(producerMultiplexer),
	new MetricsRegistryMap,
	null,
	JobContextImpl.fromConfigWithDefaults(storageConfig, null),
	new ContainerContextImpl(new ContainerModel("0", tasks.asJava), new MetricsRegistryMap), StoreMode.ReadWrite
	)

	val db = if(!engine.isInstanceOf[KeyValueStorageEngine[_,_]]) {
	throw new SamzaException("This test can only run with KeyValueStorageEngine configured as store factory.")
	} else {
	engine.asInstanceOf[KeyValueStorageEngine[Array[Byte], Array[Byte]]]
	}

	// Run the test method
	testMethod(db, storageConfig.subset("set-" + testSet + ".", true))

	new FileUtil().rm(output)
	})
	}
	}

	def runTestAllWithDeletes(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
	val numLoops = config.getInt("num.loops", 100)
	val messagesPerBatch = config.getInt("messages.per.batch", 10000)
	val messageSizeBytes = config.getInt("message.size.bytes", 200)

	info("Using (num loops, messages per batch, message size in bytes) => (%s, %s, %s)" format (numLoops, messagesPerBatch, messageSizeBytes))
	new TestKeyValuePerformance().testAllWithDeletes(db, numLoops, messagesPerBatch, messageSizeBytes)
	}

	def runTestMsgWritePerformance(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
	val messageSizeBytes = config.getInt("message.size", 200)
	val messageCount = config.getInt("message.count", 10000)

	info("Using (message count, message size in bytes) => (%s, %s)" format (messageCount, messageSizeBytes))
	new TestKeyValuePerformance().testMsgWritePerformance(db, messageCount, messageSizeBytes)
	}

	def runTestConcurrentMsgWritePerformance(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
	val messageSizeBytes = config.getInt("message.size", 200)
	val messageCount = config.getInt("message.count", 100000)
	val numThreads = config.getInt("num.threads", 4)

	new TestKeyValuePerformance().testConcurrentMsgWritePerformance(db, messageCount, messageSizeBytes, numThreads)
	}

	def runTestGetAllVsGetWriteManyReadMany(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
	new TestKeyValuePerformance().testGetAllVsGetWriteManyReadMany(db, config)
	}

	def runTestGetAllVsGetWriteOnceReadMany(db: KeyValueStore[Array[Byte], Array[Byte]], config: Config) {
	new TestKeyValuePerformance().testGetAllVsGetWriteOnceReadMany(db, config)
	}
	}

	class TestKeyValuePerformance extends Logging {
	import TestKeyValuePerformance._

	/**
	* A test that writes messagesPerBatch messages, deletes them all, then calls
	* store.all. The test periodically outputs the time it takes to complete
	* these operations. This test is useful to trouble shoot issues with LevleDB
	* such as the issue documented in SAMZA-254.
	*/
	def testAllWithDeletes(
	store: KeyValueStore[Array[Byte], Array[Byte]],

	/**
	* How many times a batch of messages should be written and deleted.
	*/
	numLoops: Int = 100,

	/**
	* The number of messages to write and delete per-batch.
	*/
	messagesPerBatch: Int = 10000,

	/**
	* The size of the messages to write.
	*/
	messageSizeBytes: Int = 200) {

	val stuff = (0 until messageSizeBytes).map(i => "a").mkString.getBytes(Encoding)
	val start = System.currentTimeMillis

	(0 until numLoops).foreach(i => {
	info("(%sms) Total written to store: %s" format (System.currentTimeMillis - start, i * messagesPerBatch))

	(0 until messagesPerBatch).foreach(j => {
	val k = (i * j).toString.getBytes(Encoding)
	store.put(k, stuff)
	store.delete(k)
	})

	val allStart = System.currentTimeMillis
	val iter = store.all
	info("(%sms) all() took %sms." format (System.currentTimeMillis - start, System.currentTimeMillis - allStart))
	iter.close
	})

	info("Total time: %ss" format ((System.currentTimeMillis - start) * .001))
	}

	/**
	* Test that successively writes a set of fixed-size messages to the KV store
	* and computes the total time for the operations
	* @param store Key-Value store instance that is being tested
	* @param numMsgs Total number of messages to write to the store
	* @param msgSizeInBytes Size of each message in Bytes
	*/
	def testMsgWritePerformance(
	store: KeyValueStore[Array[Byte], Array[Byte]],
	numMsgs: Int = 10000,
	msgSizeInBytes: Int = 200) {

	val msg = (0 until msgSizeInBytes).map(i => "x").mkString.getBytes(Encoding)

	val start = System.currentTimeMillis
	(0 until numMsgs).foreach(i => {
	store.put(i.toString.getBytes(Encoding), msg)
	})
	val timeTaken = System.currentTimeMillis - start
	info("Total time to write %d msgs of size %d bytes : %s s" format (numMsgs, msgSizeInBytes, timeTaken * .001))
	}

	def testConcurrentMsgWritePerformance(
	store: KeyValueStore[Array[Byte], Array[Byte]],
	numMsgs: Int = 100000,
	msgSizeInBytes: Int = 200,
	numThreads: Int = 4) {

	val msg = (0 until msgSizeInBytes).map(i => "x").mkString.getBytes(Encoding)
	def createThread(name: String): Thread = new Thread(new Runnable {
	override def run() = {
	(0 until numMsgs).foreach(i => {
	store.put(i.toString.getBytes(Encoding), msg)
	})
	}
	}, name)

	val threads = (0 until numThreads).map(i => createThread(s"Writer $i"))

	val start = System.currentTimeMillis
	threads.foreach(_.start())
	threads.foreach(_.join())
	val timeTaken = System.currentTimeMillis - start

	info("Total time to write %d msgs of size %d bytes with %s threads: %s sec"
	format (numMsgs, msgSizeInBytes, numThreads, timeTaken * .001))
	}

	/**
	* Test that ::getAll performance is better than that of ::get (test when there are many writes and many reads).
	* @param store key-value store instance that is being tested
	* @param config the test case's config
	*/
	def testGetAllVsGetWriteManyReadMany(store: KeyValueStore[Array[Byte],Array[Byte]], config: Config): Unit = {
	val iterationsCount = config.getInt("iterations.count", 100)
	val maxMessagesCountPerBatch = config.getInt("message.max-count-per-batch", 100000)
	val maxMessageSizeBytes = config.getInt("message.max-size.bytes", 1024)
	val timer = Stopwatch.createUnstarted
	val uuidSerde = new UUIDSerde

	info("iterations count: " + iterationsCount)
	info("max messages count per batch: " + maxMessagesCountPerBatch)
	info("max message size in bytes: " + maxMessageSizeBytes)
	info("%12s%12s%12s%12s".format("Msg Count", "Bytes/Msg", "get ms", "getAll ms"))

	try {
	(0 until iterationsCount).foreach(i => {
	val messageSizeBytes = Random.nextInt(maxMessageSizeBytes)
	val messagesCountPerBatch = Random.nextInt(maxMessagesCountPerBatch)
	val keys = (0 until messagesCountPerBatch).map(k => uuidSerde.toBytes(UUID.randomUUID)).toList
	val shuffledKeys = Random.shuffle(keys) // to reduce locality of reference -- sequential access may be unfair

	keys.foreach(k => store.put(k, Random.nextString(messageSizeBytes).getBytes(Encoding)))
	store.flush()

	timer.reset().start()
	assert(store.getAll(shuffledKeys.asJava).size == shuffledKeys.size)
	val getAllTime = timer.stop().elapsed(TimeUnit.MILLISECONDS)

	// Restore cache, in case it's enabled, to a state similar to the one above when the getAll test started
	keys.foreach(k => store.put(k, Random.nextString(messageSizeBytes).getBytes(Encoding)))
	store.flush()

	timer.reset().start()
	shuffledKeys.foreach(store.get)
	val getTime = timer.stop().elapsed(TimeUnit.MILLISECONDS)

	info("%12d%12d%12d%12d".format(messagesCountPerBatch, messageSizeBytes, getTime, getAllTime))
	if (getAllTime > getTime) {
	error("getAll was slower than get!")
	}
	})
	} finally {
	store.close()
	}
	}

	/**
	* Test that ::getAll performance is better than that of ::get (test when data are written once and read many times);
	* load is usually greater than the storage engine's cache size (not to be confused with Samza's cache layer),
	* and keys are randomly selected from the stored entries to perform a fair comparison of ::get vs. ::getAll (in case
	* the underlying storage engine caches data in blocks and ::getAll causes a block to be loaded into the cache --
	* one can argue that ::get should trigger the same behavior, but it's worth testing this WORM scenario regardless)
	* @param store key-value store instance that is being tested
	* @param config the test case's config
	*/
	def testGetAllVsGetWriteOnceReadMany(store: KeyValueStore[Array[Byte],Array[Byte]], config: Config): Unit = {
	val iterationsCount = config.getInt("iterations.count", 100)
	val maxMessagesCountPerBatch = config.getInt("message.max-count-per-batch", 10000 + Random.nextInt(20000))
	val maxMessageSizeBytes = config.getInt("message.max-size.bytes", 1024)
	val totalMessagesCount = iterationsCount * maxMessagesCountPerBatch
	val timer = Stopwatch.createUnstarted
	val uuidSerde = new UUIDSerde

	info("write once -- putting %d messages in store".format(totalMessagesCount))
	val keys = (0 until totalMessagesCount).map(k => uuidSerde.toBytes(UUID.randomUUID)).toList
	keys.foreach(k => store.put(k, Random.nextString(Random.nextInt(maxMessageSizeBytes)).getBytes(Encoding)))
	store.flush()

	info("iterations count: " + iterationsCount)
	info("max messages count per batch: " + maxMessagesCountPerBatch)
	info("max message size in bytes: " + maxMessageSizeBytes)
	info("%12s%12s%12s%12s".format("Msg Count", "Total Size", "get ms", "getAll ms"))

	try {
	(0 until iterationsCount).foreach(i => {
	val messagesCountPerBatch = Random.nextInt(maxMessagesCountPerBatch)
	val shuffledKeys = Random.shuffle(keys).take(messagesCountPerBatch)

	// We want to measure ::getAll when called many times, so populate the cache because first call is a cache-miss
	val totalSize = store.getAll(shuffledKeys.asJava).values.asScala.map(_.length).sum
	timer.reset().start()
	assert(store.getAll(shuffledKeys.asJava).size == shuffledKeys.size)
	val getAllTime = timer.stop().elapsed(TimeUnit.MILLISECONDS)

	// We want to measure ::get when called many times, so populate the cache because first call is a cache-miss
	shuffledKeys.foreach(store.get)
	timer.reset().start()
	shuffledKeys.foreach(store.get)
	val getTime = timer.stop().elapsed(TimeUnit.MILLISECONDS)

	info("%12d%12d%12d%12d".format(messagesCountPerBatch, totalSize, getTime, getAllTime))
	if (getAllTime > getTime) {
	error("getAll was slower than get!")
	}
	})
	} finally {
	store.close()
	}
	}
	}