samza-test/src/test/scala/org/apache/samza/storage/kv/TestKeyValueStores.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.storage.kv

 import java.io.File
 import java.util.concurrent.CountDownLatch
 import java.util.{Arrays, Random}

 import org.apache.samza.config.MapConfig
 import org.apache.samza.serializers.Serde
 import org.apache.samza.storage.kv.inmemory.InMemoryKeyValueStore
 import org.junit.Assert._
 import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
 import org.junit.runners.Parameterized.Parameters
 import org.junit.{After, Before, Test}
 import org.scalatest.Assertions.intercept

 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer

 /**
   * Test suite to check different key value store operations
   *
   * @param typeOfStore Defines type of key-value store (Eg: "rocksdb" / "inmemory")
   * @param storeConfig Defines whether we're using caching / serde / both / or none in front of the store
   */
 @RunWith(value = classOf[Parameterized])
 class TestKeyValueStores(typeOfStore: String, storeConfig: String) {

   import TestKeyValueStores._

   val letters = "abcdefghijklmnopqrstuvwxyz".map(_.toString)
   val dir = new File(System.getProperty("java.io.tmpdir"), "rocksdb-test-" + new Random().nextInt(Int.MaxValue))
   var store: KeyValueStore[Array[Byte], Array[Byte]] = null
   var cache = false
   var serde = false

   @Before
   def setup() {
     val kvStore: KeyValueStore[Array[Byte], Array[Byte]] = typeOfStore match {
       case "inmemory" =>
         new InMemoryKeyValueStore(new KeyValueStoreMetrics)
       case "rocksdb" =>
         new RocksDbKeyValueStore(
           dir,
           new org.rocksdb.Options()
             .setCreateIfMissing(true)
             .setCompressionType(org.rocksdb.CompressionType.SNAPPY_COMPRESSION),
           new MapConfig(),
           false,
           "someStore")
       case _ =>
         throw new IllegalArgumentException("Type of store undefined: " + typeOfStore)
     }

     val passThroughSerde = new Serde[Array[Byte]] {
       def toBytes(obj: Array[Byte]) = obj

       def fromBytes(bytes: Array[Byte]) = bytes
     }

     store = storeConfig match {
       case "cache" =>
         cache = true
         new CachedStore(kvStore, CacheSize, BatchSize)
       case "serde" =>
         serde = true
         new SerializedKeyValueStore(kvStore, passThroughSerde, passThroughSerde)
       case "cache-and-serde" =>
         val serializedStore = new SerializedKeyValueStore(kvStore, passThroughSerde, passThroughSerde)
         serde = true
         cache = true
         new CachedStore(serializedStore, CacheSize, BatchSize)
       case _ =>
         kvStore
     }
     store = new NullSafeKeyValueStore(store)
   }

   @After
   def teardown() {
     store.close
     if (dir != null && dir.listFiles() != null) {
       for (file <- dir.listFiles)
         file.delete()
       dir.delete()
     }
   }

   @Test
   def getNonExistentIsNull() {
     assertNull(store.get(b("hello")))
   }

   @Test
   def testGetAllWhenZeroMatch() {
     store.put(b("hello"), b("world"))
     val keys = List(b("foo"), b("bar"))
     val actual = store.getAll(keys.asJava)
     keys.foreach(k => assertNull("Key: " + k, actual.get(k)))
   }

   @Test
   def testGetAllWhenFullMatch() {
     val expected = Map(b("k0") -> b("v0"), b("k1") -> b("v1"))
     expected.foreach(e => store.put(e._1, e._2))
     val actual = store.getAll(expected.keys.toList.asJava)
     assertEquals("Size", expected.size, actual.size)
     expected.foreach(e => assertArrayEquals("Value at: " + s(e._1), e._2, actual.get(e._1)))
   }

   @Test
   def testGetAllWhenPartialMatch() {
     val all = Map(b("k0") -> b("v0"), b("k1") -> b("v1"), b("k2") -> b("v2")).asJava
     val found = all.entrySet.asScala.head
     val notFound = all.entrySet.asScala.last
     store.put(found.getKey, found.getValue)
     val actual = store.getAll(List(notFound.getKey, found.getKey).asJava)
     assertNull(actual.get(notFound.getKey))
     assertArrayEquals(found.getValue, actual.get(found.getKey))
   }

   @Test
   def putAndGet() {
     val k = b("k")
     store.put(k, b("v"))
     assertArrayEquals(b("v"), store.get(k))
   }

   @Test
   def doublePutAndGet() {
     val k = b("k2")
     store.put(k, b("v1"))
     store.put(k, b("v3"))
     assertArrayEquals(b("v3"), store.get(k))
   }

   @Test
   def testNullsWithSerde() {
     if (serde) {
       val a = b("a")

       intercept[NullPointerException] {
         store.get(null)
       }
       intercept[NullPointerException] {
         store.getAll(null)
       }
       intercept[NullPointerException] {
         store.getAll(List(a, null).asJava)
       }
       intercept[NullPointerException] {
         store.delete(null)
       }
       intercept[NullPointerException] {
         store.deleteAll(null)
       }
       intercept[NullPointerException] {
         store.deleteAll(List(a, null).asJava)
       }
       intercept[NullPointerException] {
         store.put(null, a)
       }
       intercept[NullPointerException] {
         store.put(a, null)
       }
       intercept[NullPointerException] {
         store.putAll(List(new Entry(a, a), new Entry[Array[Byte], Array[Byte]](a, null)).asJava)
       }
       intercept[NullPointerException] {
         store.putAll(List(new Entry[Array[Byte], Array[Byte]](null, a)).asJava)
       }
       intercept[NullPointerException] {
         store.range(a, null)
       }
       intercept[NullPointerException] {
         store.range(null, a)
       }
     }
   }

   @Test
   def testPutAll() {
     // Use CacheSize - 1 so we fully fill the cache, but don't write any data
     // out. Our check (below) uses == for cached entries, and using
     // numEntires >= CacheSize would result in the LRU cache dropping some
     // entries. The result would be that we get the correct byte array back
     // from the cache's underlying store (rocksdb), but that == would fail.
     val numEntries = CacheSize - 1
     val entries = (0 until numEntries).map(i => new Entry(b("k" + i), b("v" + i)))
     store.putAll(entries.asJava)
     if (cache) {
       assertTrue("All values should be found and cached.", entries.forall(e => store.get(e.getKey) == e.getValue))
     } else {
       assertTrue("All values should be found.", entries.forall(e => Arrays.equals(store.get(e.getKey), e.getValue)))
     }
   }

   @Test
   def testIterateAll() {
     for (letter <- letters)
       store.put(b(letter.toString), b(letter.toString))
     val iter = store.all
     checkRange(letters, iter)
     iter.close()
   }

   @Test
   def testRange() {
     val from = 5
     val to = 20
     for (letter <- letters)
       store.put(b(letter.toString), b(letter.toString))

     val iter = store.range(b(letters(from)), b(letters(to)))
     checkRange(letters.slice(from, to), iter)
     iter.close()
   }

   @Test
   def testDelete() {
     val a = b("a")
     assertNull(store.get(a))
     store.put(a, a)
     assertArrayEquals(a, store.get(a))
     store.delete(a)
     assertNull(store.get(a))
   }

   @Test
   def testDeleteAllWhenZeroMatch() {
     val foo = b("foo")
     store.put(foo, foo)
     store.deleteAll(List(b("bar")).asJava)
     assertArrayEquals(foo, store.get(foo))
   }

   @Test
   def testDeleteAllWhenFullMatch() {
     val all = Map(b("k0") -> b("v0"), b("k1") -> b("v1"))
     all.foreach(e => store.put(e._1, e._2))
     assertEquals(all.size, store.getAll(all.keys.toList.asJava).size)
     store.deleteAll(all.keys.toList.asJava)
     all.keys.foreach(key => assertNull("Value at: " + s(key), store.get(key)))
   }

   @Test
   def testDeleteAllWhenPartialMatch() {
     val all = Map(b("k0") -> b("v0"), b("k1") -> b("v1")).asJava
     val found = all.entrySet.asScala.head
     val leftAlone = all.entrySet.asScala.last
     all.asScala.foreach(e => store.put(e._1, e._2))
     assertArrayEquals(found.getValue, store.get(found.getKey))
     store.deleteAll(List(b("not found"), found.getKey).asJava)
     store.flush()
     val allIterator = store.all
     try {
       assertEquals(1, allIterator.asScala.size)
       assertArrayEquals(leftAlone.getValue, store.get(leftAlone.getKey))
     } finally {
       allIterator.close()
     }
   }

   @Test
   def testSimpleScenario() {
     val vals = letters.map(b(_))
     for (v <- vals) {
       assertNull(store.get(v))
       store.put(v, v)
       assertArrayEquals(v, store.get(v))
     }
     vals.foreach(v => assertArrayEquals(v, store.get(v)))
     vals.foreach(v => store.delete(v))
     vals.foreach(v => assertNull(store.get(v)))
   }

   /**
     * This test specifically targets an issue in Scala 2.8.1's DoubleLinkedList
     * implementation. The issue is that it doesn't work. More specifically,
     * creating a DoubleLinkedList from an existing list does not update the
     * "prev" field of the existing list's head to point to the new head. As a
     * result, in Scala 2.8.1, every DoubleLinkedList node's prev field is null.
     * Samza gets around this by manually updating the field itself. See SAMZA-80
     * for details.
     *
     * This issue is exposed in Samza's KV cache implementation, which uses
     * DoubleLinkedList, so all comments in this method are discussing the cached
     * implementation, but the test is still useful as a sanity check for
     * non-cached stores.
     */
   @Test
   def testBrokenScalaDoubleLinkedList() {
     val something = b("")
     val keys = letters
       .map(b(_))
       .toArray

     // Load the cache to capacity.
     letters
       .slice(0, TestKeyValueStores.CacheSize)
       .map(b(_))
       .foreach(store.put(_, something))

     // Now keep everything in the cache, but with an empty dirty list.
     store.flush

     // Dirty list is now empty, and every CacheEntry has dirty=null.

     // Corrupt the dirty list by creating two dirty lists that toggle back and
     // forth depending on whether the last dirty write was to 1 or 0. The trick
     // here is that every element in the cache is treated as the "head" of the
     // DoulbeLinkedList (prev==null), even though it's not necessarily. Thus,
     // You can end up with multiple nodes each having their own version of the
     // dirty list with different elements in them.
     store.put(keys(1), something)
     store.put(keys(0), something)
     store.put(keys(1), something)
     store.flush
     // The dirty list is now empty, but 0's dirty field actually has 0 and 1.
     store.put(keys(0), something)
     // The dirty list now has 0 and 1, but 1's dirty field is null in the
     // cache because it was just flushed.

     // Get rid of 1 from the cache by reading every other element, and then
     // putting one new element.
     letters
       .slice(2, TestKeyValueStores.CacheSize)
       .map(b(_))
       .foreach(store.get(_))
     store.put(keys(10), something)

     // Now try and trigger an NPE since the dirty list has an element (1)
     // that's no longer in the cache.
     store.flush
   }

   /**
     * A little test that tries to simulate a few common patterns:
     * read-modify-write, and do-some-stuff-then-delete (windowing).
     */
   @Test
   def testRandomReadWriteRemove() {
     // Make test deterministic by seeding the random number generator.
     val rand = new Random(12345)
     val keys = letters
       .map(b(_))
       .toArray

     // Map from letter to key byte array used for letter, and expected value.
     // We have to go through some acrobatics here since Java's byte array uses
     // object identity for .equals. Two byte arrays can have identical byte
     // elements, but not be equal.
     var expected = Map[String, (Array[Byte], String)]()

     (0 until 100).foreach(loop => {
       (0 until 30).foreach(i => {
         val idx = rand.nextInt(keys.length)
         val randomValue = letters(rand.nextInt(keys.length))
         val key = keys(idx)
         val currentVal = store.get(key)
         store.put(key, b(randomValue))
         expected += letters(idx) -> (key, randomValue)
       })

       for ((k, v) <- expected) {
         val bytes = store.get(v._1)
         assertNotNull(bytes)
         assertEquals(v._2, new String(bytes, "UTF-8"))
       }

       // Iterating and making structural modifications (deletion) does not look right.
       // Separating these two steps
       val iterator = store.all
       val allKeys = new ArrayBuffer[Array[Byte]]()

       // First iterate
       while (iterator.hasNext) {
         allKeys += iterator.next.getKey
       }
       iterator.close

       // And now delete
       for (key <- allKeys) {
         store.delete(key)
         expected -= new String(key, "UTF-8")
       }

       assertEquals(0, expected.size)
     })
   }

   @Test
   def testParallelReadWriteDiffKeys(): Unit = {
     // Make test deterministic by seeding the random number generator.
     val key1 = b("key1")
     val key2 = b("key2")
     val val1 = "val1"
     val val2 = "val2"

     val runner1 = new Thread(new Runnable {
       override def run(): Unit = {
         store.put(key1, b(val1))
       }
     })

     val runner2 = new Thread(new Runnable {
       override def run(): Unit = {
         while (!val1.equals({
           store.get(key1) match {
             case null => ""
             case _ => {
               new String(store.get(key1), "UTF-8")
             }
           }
         })) {}
         store.delete(key1)
       }
     })

     val runner3 = new Thread(new Runnable {
       override def run(): Unit = {
         store.put(key2, b(val2))
       }
     })

     val runner4 = new Thread(new Runnable {
       override def run(): Unit = {
         while (!val2.equals({
           store.get(key2) match {
             case null => ""
             case _ => {
               new String(store.get(key2), "UTF-8")
             }
           }
         })) {}
         store.delete(key2)
       }
     })

     runner2.start()
     runner1.start()
     runner3.start()
     runner4.start()

     runner2.join(1000)
     runner1.join(1000)
     runner3.join(1000)
     runner4.join(1000)

     assertNull(store.get(key1))
     assertNull(store.get(key2))

     store.flush()
   }

   @Test
   def testParallelIteratorAndWrite(): Unit = {
     // Make test deterministic by seeding the random number generator.
     val key1 = b("key1")
     val key2 = b("key2")
     val val1 = "val1"
     val val2 = "val2"
     @volatile var throwable: Throwable = null

     store.put(key1, b(val1))
     store.put(key2, b(val2))

     val runner1StartLatch = new CountDownLatch(1)
     val runner2StartLatch = new CountDownLatch(1)

     val runner1 = new Thread(new Runnable {
       override def run(): Unit = {
         runner1StartLatch.await()
         store.put(key1, b("val1-2"))
         store.delete(key2)
         store.flush()
         runner2StartLatch.countDown()
       }
     })

     val runner2 = new Thread(new Runnable {
       override def run(): Unit = {
         runner2StartLatch.await()
         val iter = store.all() //snapshot after change
         try {
           while (iter.hasNext) {
             val e = iter.next()
             if ("key1".equals(new String(e.getKey, "UTF-8"))) {
               assertEquals("val1-2", new String(e.getValue, "UTF-8"))
             }
             System.out.println(String.format("iterator1: key: %s, value: %s", new String(e.getKey, "UTF-8"), new String(e.getValue, "UTF-8")))
           }
           iter.close()
         } catch {
           case t: Throwable => throwable = t
         }
       }
     })

     val runner3 = new Thread(new Runnable {
       override def run(): Unit = {
         val iter = store.all() //snapshot
         runner1StartLatch.countDown()
         try {
           while (iter.hasNext) {
             val e = iter.next()
             val key = new String(e.getKey, "UTF-8")
             val value = new String(e.getValue, "UTF-8")
             if (key.equals("key1")) {
               assertEquals(val1, value)
             }
             else if (key.equals("key2") && !val2.equals(value)) {
               assertEquals(val2, value)
             }
             else if (!key.equals("key1") && !key.equals("key2")) {
               throw new Exception("unknow key " + new String(e.getKey, "UTF-8") + ", value: " + new String(e.getValue, "UTF-8"))
             }
             System.out.println(String.format("iterator2: key: %s, value: %s", new String(e.getKey, "UTF-8"), new String(e.getValue, "UTF-8")))
           }
           iter.close()
         } catch {
           case t: Throwable => throwable = t
         }
       }
     })

     runner2.start()
     runner3.start()
     runner1.start()

     runner2.join()
     runner1.join()
     runner3.join()

     if (throwable != null) throw throwable

     store.flush()
   }

   def checkRange(vals: IndexedSeq[String], iter: KeyValueIterator[Array[Byte], Array[Byte]]) {
     for (v <- vals) {
       assertTrue(iter.hasNext)
       val entry = iter.next()
       assertEquals(v, s(entry.getKey))
       assertEquals(v, s(entry.getValue))
     }
     assertFalse(iter.hasNext)
     intercept[NoSuchElementException] {
       iter.next()
     }
   }

   /**
     * Convert string to byte buffer
     */
   def b(s: String) =
     s.getBytes

   /**
     * Convert byte buffer to string
     */
   def s(b: Array[Byte]) =
     new String(b)
 }

 object TestKeyValueStores {
   val CacheSize = 1024
   val BatchSize = 1024

   @Parameters
   def parameters: java.util.Collection[Array[String]] = Arrays.asList(
     //Inmemory
     Array("inmemory", "cache"),
     Array("inmemory", "serde"),
     Array("inmemory", "cache-and-serde"),
     Array("inmemory", "none"),
     //RocksDB
     Array("rocksdb", "cache"),
     Array("rocksdb", "serde"),
     Array("rocksdb", "cache-and-serde"),
     Array("rocksdb", "none")
   )
 }
	/*
	* 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.storage.kv

	import java.io.File
	import java.util.concurrent.CountDownLatch
	import java.util.{Arrays, Random}

	import org.apache.samza.config.MapConfig
	import org.apache.samza.serializers.Serde
	import org.apache.samza.storage.kv.inmemory.InMemoryKeyValueStore
	import org.junit.Assert._
	import org.junit.runner.RunWith
	import org.junit.runners.Parameterized
	import org.junit.runners.Parameterized.Parameters
	import org.junit.{After, Before, Test}
	import org.scalatest.Assertions.intercept

	import scala.collection.JavaConverters._
	import scala.collection.mutable.ArrayBuffer

	/**
	* Test suite to check different key value store operations
	*
	* @param typeOfStore Defines type of key-value store (Eg: "rocksdb" / "inmemory")
	* @param storeConfig Defines whether we're using caching / serde / both / or none in front of the store
	*/
	@RunWith(value = classOf[Parameterized])
	class TestKeyValueStores(typeOfStore: String, storeConfig: String) {

	import TestKeyValueStores._

	val letters = "abcdefghijklmnopqrstuvwxyz".map(_.toString)
	val dir = new File(System.getProperty("java.io.tmpdir"), "rocksdb-test-" + new Random().nextInt(Int.MaxValue))
	var store: KeyValueStore[Array[Byte], Array[Byte]] = null
	var cache = false
	var serde = false

	@Before
	def setup() {
	val kvStore: KeyValueStore[Array[Byte], Array[Byte]] = typeOfStore match {
	case "inmemory" =>
	new InMemoryKeyValueStore(new KeyValueStoreMetrics)
	case "rocksdb" =>
	new RocksDbKeyValueStore(
	dir,
	new org.rocksdb.Options()
	.setCreateIfMissing(true)
	.setCompressionType(org.rocksdb.CompressionType.SNAPPY_COMPRESSION),
	new MapConfig(),
	false,
	"someStore")
	case _ =>
	throw new IllegalArgumentException("Type of store undefined: " + typeOfStore)
	}

	val passThroughSerde = new Serde[Array[Byte]] {
	def toBytes(obj: Array[Byte]) = obj

	def fromBytes(bytes: Array[Byte]) = bytes
	}

	store = storeConfig match {
	case "cache" =>
	cache = true
	new CachedStore(kvStore, CacheSize, BatchSize)
	case "serde" =>
	serde = true
	new SerializedKeyValueStore(kvStore, passThroughSerde, passThroughSerde)
	case "cache-and-serde" =>
	val serializedStore = new SerializedKeyValueStore(kvStore, passThroughSerde, passThroughSerde)
	serde = true
	cache = true
	new CachedStore(serializedStore, CacheSize, BatchSize)
	case _ =>
	kvStore
	}
	store = new NullSafeKeyValueStore(store)
	}

	@After
	def teardown() {
	store.close
	if (dir != null && dir.listFiles() != null) {
	for (file <- dir.listFiles)
	file.delete()
	dir.delete()
	}
	}

	@Test
	def getNonExistentIsNull() {
	assertNull(store.get(b("hello")))
	}

	@Test
	def testGetAllWhenZeroMatch() {
	store.put(b("hello"), b("world"))
	val keys = List(b("foo"), b("bar"))
	val actual = store.getAll(keys.asJava)
	keys.foreach(k => assertNull("Key: " + k, actual.get(k)))
	}

	@Test
	def testGetAllWhenFullMatch() {
	val expected = Map(b("k0") -> b("v0"), b("k1") -> b("v1"))
	expected.foreach(e => store.put(e._1, e._2))
	val actual = store.getAll(expected.keys.toList.asJava)
	assertEquals("Size", expected.size, actual.size)
	expected.foreach(e => assertArrayEquals("Value at: " + s(e._1), e._2, actual.get(e._1)))
	}

	@Test
	def testGetAllWhenPartialMatch() {
	val all = Map(b("k0") -> b("v0"), b("k1") -> b("v1"), b("k2") -> b("v2")).asJava
	val found = all.entrySet.asScala.head
	val notFound = all.entrySet.asScala.last
	store.put(found.getKey, found.getValue)
	val actual = store.getAll(List(notFound.getKey, found.getKey).asJava)
	assertNull(actual.get(notFound.getKey))
	assertArrayEquals(found.getValue, actual.get(found.getKey))
	}

	@Test
	def putAndGet() {
	val k = b("k")
	store.put(k, b("v"))
	assertArrayEquals(b("v"), store.get(k))
	}

	@Test
	def doublePutAndGet() {
	val k = b("k2")
	store.put(k, b("v1"))
	store.put(k, b("v3"))
	assertArrayEquals(b("v3"), store.get(k))
	}

	@Test
	def testNullsWithSerde() {
	if (serde) {
	val a = b("a")

	intercept[NullPointerException] {
	store.get(null)
	}
	intercept[NullPointerException] {
	store.getAll(null)
	}
	intercept[NullPointerException] {
	store.getAll(List(a, null).asJava)
	}
	intercept[NullPointerException] {
	store.delete(null)
	}
	intercept[NullPointerException] {
	store.deleteAll(null)
	}
	intercept[NullPointerException] {
	store.deleteAll(List(a, null).asJava)
	}
	intercept[NullPointerException] {
	store.put(null, a)
	}
	intercept[NullPointerException] {
	store.put(a, null)
	}
	intercept[NullPointerException] {
	store.putAll(List(new Entry(a, a), new Entry[Array[Byte], Array[Byte]](a, null)).asJava)
	}
	intercept[NullPointerException] {
	store.putAll(List(new Entry[Array[Byte], Array[Byte]](null, a)).asJava)
	}
	intercept[NullPointerException] {
	store.range(a, null)
	}
	intercept[NullPointerException] {
	store.range(null, a)
	}
	}
	}

	@Test
	def testPutAll() {
	// Use CacheSize - 1 so we fully fill the cache, but don't write any data
	// out. Our check (below) uses == for cached entries, and using
	// numEntires >= CacheSize would result in the LRU cache dropping some
	// entries. The result would be that we get the correct byte array back
	// from the cache's underlying store (rocksdb), but that == would fail.
	val numEntries = CacheSize - 1
	val entries = (0 until numEntries).map(i => new Entry(b("k" + i), b("v" + i)))
	store.putAll(entries.asJava)
	if (cache) {
	assertTrue("All values should be found and cached.", entries.forall(e => store.get(e.getKey) == e.getValue))
	} else {
	assertTrue("All values should be found.", entries.forall(e => Arrays.equals(store.get(e.getKey), e.getValue)))
	}
	}

	@Test
	def testIterateAll() {
	for (letter <- letters)
	store.put(b(letter.toString), b(letter.toString))
	val iter = store.all
	checkRange(letters, iter)
	iter.close()
	}

	@Test
	def testRange() {
	val from = 5
	val to = 20
	for (letter <- letters)
	store.put(b(letter.toString), b(letter.toString))

	val iter = store.range(b(letters(from)), b(letters(to)))
	checkRange(letters.slice(from, to), iter)
	iter.close()
	}

	@Test
	def testDelete() {
	val a = b("a")
	assertNull(store.get(a))
	store.put(a, a)
	assertArrayEquals(a, store.get(a))
	store.delete(a)
	assertNull(store.get(a))
	}

	@Test
	def testDeleteAllWhenZeroMatch() {
	val foo = b("foo")
	store.put(foo, foo)
	store.deleteAll(List(b("bar")).asJava)
	assertArrayEquals(foo, store.get(foo))
	}

	@Test
	def testDeleteAllWhenFullMatch() {
	val all = Map(b("k0") -> b("v0"), b("k1") -> b("v1"))
	all.foreach(e => store.put(e._1, e._2))
	assertEquals(all.size, store.getAll(all.keys.toList.asJava).size)
	store.deleteAll(all.keys.toList.asJava)
	all.keys.foreach(key => assertNull("Value at: " + s(key), store.get(key)))
	}

	@Test
	def testDeleteAllWhenPartialMatch() {
	val all = Map(b("k0") -> b("v0"), b("k1") -> b("v1")).asJava
	val found = all.entrySet.asScala.head
	val leftAlone = all.entrySet.asScala.last
	all.asScala.foreach(e => store.put(e._1, e._2))
	assertArrayEquals(found.getValue, store.get(found.getKey))
	store.deleteAll(List(b("not found"), found.getKey).asJava)
	store.flush()
	val allIterator = store.all
	try {
	assertEquals(1, allIterator.asScala.size)
	assertArrayEquals(leftAlone.getValue, store.get(leftAlone.getKey))
	} finally {
	allIterator.close()
	}
	}

	@Test
	def testSimpleScenario() {
	val vals = letters.map(b(_))
	for (v <- vals) {
	assertNull(store.get(v))
	store.put(v, v)
	assertArrayEquals(v, store.get(v))
	}
	vals.foreach(v => assertArrayEquals(v, store.get(v)))
	vals.foreach(v => store.delete(v))
	vals.foreach(v => assertNull(store.get(v)))
	}

	/**
	* This test specifically targets an issue in Scala 2.8.1's DoubleLinkedList
	* implementation. The issue is that it doesn't work. More specifically,
	* creating a DoubleLinkedList from an existing list does not update the
	* "prev" field of the existing list's head to point to the new head. As a
	* result, in Scala 2.8.1, every DoubleLinkedList node's prev field is null.
	* Samza gets around this by manually updating the field itself. See SAMZA-80
	* for details.
	*
	* This issue is exposed in Samza's KV cache implementation, which uses
	* DoubleLinkedList, so all comments in this method are discussing the cached
	* implementation, but the test is still useful as a sanity check for
	* non-cached stores.
	*/
	@Test
	def testBrokenScalaDoubleLinkedList() {
	val something = b("")
	val keys = letters
	.map(b(_))
	.toArray

	// Load the cache to capacity.
	letters
	.slice(0, TestKeyValueStores.CacheSize)
	.map(b(_))
	.foreach(store.put(_, something))

	// Now keep everything in the cache, but with an empty dirty list.
	store.flush

	// Dirty list is now empty, and every CacheEntry has dirty=null.

	// Corrupt the dirty list by creating two dirty lists that toggle back and
	// forth depending on whether the last dirty write was to 1 or 0. The trick
	// here is that every element in the cache is treated as the "head" of the
	// DoulbeLinkedList (prev==null), even though it's not necessarily. Thus,
	// You can end up with multiple nodes each having their own version of the
	// dirty list with different elements in them.
	store.put(keys(1), something)
	store.put(keys(0), something)
	store.put(keys(1), something)
	store.flush
	// The dirty list is now empty, but 0's dirty field actually has 0 and 1.
	store.put(keys(0), something)
	// The dirty list now has 0 and 1, but 1's dirty field is null in the
	// cache because it was just flushed.

	// Get rid of 1 from the cache by reading every other element, and then
	// putting one new element.
	letters
	.slice(2, TestKeyValueStores.CacheSize)
	.map(b(_))
	.foreach(store.get(_))
	store.put(keys(10), something)

	// Now try and trigger an NPE since the dirty list has an element (1)
	// that's no longer in the cache.
	store.flush
	}

	/**
	* A little test that tries to simulate a few common patterns:
	* read-modify-write, and do-some-stuff-then-delete (windowing).
	*/
	@Test
	def testRandomReadWriteRemove() {
	// Make test deterministic by seeding the random number generator.
	val rand = new Random(12345)
	val keys = letters
	.map(b(_))
	.toArray

	// Map from letter to key byte array used for letter, and expected value.
	// We have to go through some acrobatics here since Java's byte array uses
	// object identity for .equals. Two byte arrays can have identical byte
	// elements, but not be equal.
	var expected = Map[String, (Array[Byte], String)]()

	(0 until 100).foreach(loop => {
	(0 until 30).foreach(i => {
	val idx = rand.nextInt(keys.length)
	val randomValue = letters(rand.nextInt(keys.length))
	val key = keys(idx)
	val currentVal = store.get(key)
	store.put(key, b(randomValue))
	expected += letters(idx) -> (key, randomValue)
	})

	for ((k, v) <- expected) {
	val bytes = store.get(v._1)
	assertNotNull(bytes)
	assertEquals(v._2, new String(bytes, "UTF-8"))
	}

	// Iterating and making structural modifications (deletion) does not look right.
	// Separating these two steps
	val iterator = store.all
	val allKeys = new ArrayBuffer[Array[Byte]]()

	// First iterate
	while (iterator.hasNext) {
	allKeys += iterator.next.getKey
	}
	iterator.close

	// And now delete
	for (key <- allKeys) {
	store.delete(key)
	expected -= new String(key, "UTF-8")
	}

	assertEquals(0, expected.size)
	})
	}

	@Test
	def testParallelReadWriteDiffKeys(): Unit = {
	// Make test deterministic by seeding the random number generator.
	val key1 = b("key1")
	val key2 = b("key2")
	val val1 = "val1"
	val val2 = "val2"

	val runner1 = new Thread(new Runnable {
	override def run(): Unit = {
	store.put(key1, b(val1))
	}
	})

	val runner2 = new Thread(new Runnable {
	override def run(): Unit = {
	while (!val1.equals({
	store.get(key1) match {
	case null => ""
	case _ => {
	new String(store.get(key1), "UTF-8")
	}
	}
	})) {}
	store.delete(key1)
	}
	})

	val runner3 = new Thread(new Runnable {
	override def run(): Unit = {
	store.put(key2, b(val2))
	}
	})

	val runner4 = new Thread(new Runnable {
	override def run(): Unit = {
	while (!val2.equals({
	store.get(key2) match {
	case null => ""
	case _ => {
	new String(store.get(key2), "UTF-8")
	}
	}
	})) {}
	store.delete(key2)
	}
	})

	runner2.start()
	runner1.start()
	runner3.start()
	runner4.start()

	runner2.join(1000)
	runner1.join(1000)
	runner3.join(1000)
	runner4.join(1000)

	assertNull(store.get(key1))
	assertNull(store.get(key2))

	store.flush()
	}

	@Test
	def testParallelIteratorAndWrite(): Unit = {
	// Make test deterministic by seeding the random number generator.
	val key1 = b("key1")
	val key2 = b("key2")
	val val1 = "val1"
	val val2 = "val2"
	@volatile var throwable: Throwable = null

	store.put(key1, b(val1))
	store.put(key2, b(val2))

	val runner1StartLatch = new CountDownLatch(1)
	val runner2StartLatch = new CountDownLatch(1)

	val runner1 = new Thread(new Runnable {
	override def run(): Unit = {
	runner1StartLatch.await()
	store.put(key1, b("val1-2"))
	store.delete(key2)
	store.flush()
	runner2StartLatch.countDown()
	}
	})

	val runner2 = new Thread(new Runnable {
	override def run(): Unit = {
	runner2StartLatch.await()
	val iter = store.all() //snapshot after change
	try {
	while (iter.hasNext) {
	val e = iter.next()
	if ("key1".equals(new String(e.getKey, "UTF-8"))) {
	assertEquals("val1-2", new String(e.getValue, "UTF-8"))
	}
	System.out.println(String.format("iterator1: key: %s, value: %s", new String(e.getKey, "UTF-8"), new String(e.getValue, "UTF-8")))
	}
	iter.close()
	} catch {
	case t: Throwable => throwable = t
	}
	}
	})

	val runner3 = new Thread(new Runnable {
	override def run(): Unit = {
	val iter = store.all() //snapshot
	runner1StartLatch.countDown()
	try {
	while (iter.hasNext) {
	val e = iter.next()
	val key = new String(e.getKey, "UTF-8")
	val value = new String(e.getValue, "UTF-8")
	if (key.equals("key1")) {
	assertEquals(val1, value)
	}
	else if (key.equals("key2") && !val2.equals(value)) {
	assertEquals(val2, value)
	}
	else if (!key.equals("key1") && !key.equals("key2")) {
	throw new Exception("unknow key " + new String(e.getKey, "UTF-8") + ", value: " + new String(e.getValue, "UTF-8"))
	}
	System.out.println(String.format("iterator2: key: %s, value: %s", new String(e.getKey, "UTF-8"), new String(e.getValue, "UTF-8")))
	}
	iter.close()
	} catch {
	case t: Throwable => throwable = t
	}
	}
	})

	runner2.start()
	runner3.start()
	runner1.start()

	runner2.join()
	runner1.join()
	runner3.join()

	if (throwable != null) throw throwable

	store.flush()
	}

	def checkRange(vals: IndexedSeq[String], iter: KeyValueIterator[Array[Byte], Array[Byte]]) {
	for (v <- vals) {
	assertTrue(iter.hasNext)
	val entry = iter.next()
	assertEquals(v, s(entry.getKey))
	assertEquals(v, s(entry.getValue))
	}
	assertFalse(iter.hasNext)
	intercept[NoSuchElementException] {
	iter.next()
	}
	}

	/**
	* Convert string to byte buffer
	*/
	def b(s: String) =
	s.getBytes

	/**
	* Convert byte buffer to string
	*/
	def s(b: Array[Byte]) =
	new String(b)
	}

	object TestKeyValueStores {
	val CacheSize = 1024
	val BatchSize = 1024

	@Parameters
	def parameters: java.util.Collection[Array[String]] = Arrays.asList(
	//Inmemory
	Array("inmemory", "cache"),
	Array("inmemory", "serde"),
	Array("inmemory", "cache-and-serde"),
	Array("inmemory", "none"),
	//RocksDB
	Array("rocksdb", "cache"),
	Array("rocksdb", "serde"),
	Array("rocksdb", "cache-and-serde"),
	Array("rocksdb", "none")
	)
	}