samza-kv/src/main/scala/org/apache/samza/storage/kv/CachedStore.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 org.apache.samza.util.Logging

 import scala.collection._
 import java.nio.file.Path
 import java.util.{Arrays, Optional}

 import com.google.common.annotations.VisibleForTesting
 import org.apache.samza.checkpoint.CheckpointId

 /**
  * A write-behind caching layer around the rocksdb store. The purpose of this cache is three-fold:
  * 1. Batch together writes to rocksdb, this turns out to be a great optimization
  * 2. Avoid duplicate writes and duplicate log entries within a commit interval. i.e. if there are two updates to the same key, log only the later.
  * 3. Avoid deserialization cost for gets on very common keys (unless the keys itself are an array of bytes,
  * in which case they don't need to be deserialized and there is no performance benefit).
  *
  * This caching does introduce a few odd corner cases :-(
  * 1. Items in the cache have pass-by-reference semantics but items in rocksdb have pass-by-value semantics. Modifying items after a put is a bad idea.
  * 2. Range queries require flushing the cache (as the ordering comes from rocksdb)
  *
  * In implementation this cache is just an LRU hash map that discards the oldest entry when full. There is an accompanying "dirty list" that references keys
  * that have not yet been written to disk. All writes go to the dirty list and when the list is long enough we call putAll on all those values at once. Dirty items
  * that time out of the cache before being written will also trigger a putAll of the dirty list.
  *
  * This class is thread safe.
  *
  * @param store The store to cache
  * @param cacheSize The number of entries to hold in the in memory-cache
  * @param writeBatchSize The number of entries to batch together before forcing a write
  * @param metrics The metrics recording object for this cached store
  */
 class CachedStore[K, V](
   val store: KeyValueStore[K, V],
   val cacheSize: Int,
   val writeBatchSize: Int,
   val metrics: CachedStoreMetrics = new CachedStoreMetrics) extends KeyValueStore[K, V] with Logging {

   /** the number of items in the dirty list */
   @volatile private var dirtyCount = 0

   /** the number of items currently in the cache */
   @volatile private var cacheCount = 0

   /** the list of items to be written out on flush from newest to oldest */
   private var dirty = new mutable.DoubleLinkedList[K]()

   /** the synchronization lock to protect access to the store from multiple threads **/
   private val lock = new Object

   /** an lru cache of values that holds cacheEntries and calls putAll() on dirty entries if necessary when discarding */
   private val cache = new java.util.LinkedHashMap[K, CacheEntry[K, V]]((cacheSize * 1.2).toInt, 1.0f, true) {
     override def removeEldestEntry(eldest: java.util.Map.Entry[K, CacheEntry[K, V]]): Boolean = {
       val evict = super.size > cacheSize
       if (evict) {
         val entry = eldest.getValue
         // if this entry hasn't been written out yet, flush it and all other dirty keys
         if (entry.dirty != null) {
           debug("Found a dirty entry. Calling putAll() on all dirty entries.")
           putAllDirtyEntries()
         }
       }
       evict
     }
   }

   /** tracks whether an array has been used as a key. since this is dangerous with LinkedHashMap, we want to warn on it. **/
   @volatile private var containsArrayKeys = false

   // Use counters here, rather than directly accessing variables using .size
   // since metrics can be accessed in other threads, and cache.size is not
   // thread safe since we're using a LinkedHashMap, and dirty.size is slow
   // since it requires a full traversal of the linked list.
   metrics.setDirtyCount(() => dirtyCount)
   metrics.setCacheSize(() => cacheCount)

   override def get(key: K) = lock.synchronized({
     metrics.gets.inc

     val c = cache.get(key)
     if (c != null) {
       metrics.cacheHits.inc
       c.value
     } else {
       val v = store.get(key)
       cache.put(key, new CacheEntry(v, null))
       cacheCount = cache.size
       v
     }
   })

   private class CachedStoreIterator(val iter: KeyValueIterator[K, V])
     extends KeyValueIterator[K, V] {

     var last: Entry[K, V] = null

     override def close(): Unit = iter.close()

     override def remove(): Unit = throw new UnsupportedOperationException("CachedStore iterator doesn't support remove")

     override def next() = {
       last = iter.next()
       last
     }

     override def hasNext: Boolean = iter.hasNext
   }

   override def range(from: K, to: K): KeyValueIterator[K, V] = lock.synchronized({
     metrics.ranges.inc
     putAllDirtyEntries()

     new CachedStoreIterator(store.range(from, to))
   })

   override def all(): KeyValueIterator[K, V] = lock.synchronized({
     metrics.alls.inc
     putAllDirtyEntries()

     new CachedStoreIterator(store.all())
   })

   override def put(key: K, value: V) {
     lock.synchronized({
       metrics.puts.inc

       checkKeyIsArray(key)

       // Add the key to the front of the dirty list (and remove any prior
       // occurrences to dedupe).
       val found = cache.get(key)
       if (found == null || found.dirty == null) {
         this.dirtyCount += 1
       } else {
         // If we are removing the head of the list, move the head to the next
         // element. See SAMZA-45 for details.
         if (found.dirty.prev == null) {
           this.dirty = found.dirty.next
           this.dirty.prev = null
         } else {
           found.dirty.remove()
         }
       }
       this.dirty = new mutable.DoubleLinkedList(key, this.dirty)

       // Add the key to the cache (but don't allocate a new cache entry if we
       // already have one).
       if (found == null) {
         cache.put(key, new CacheEntry(value, this.dirty))
         cacheCount = cache.size
       } else {
         found.value = value
         found.dirty = this.dirty
       }

       // putAll() dirty values if the write list is full.
       val purgeNeeded = if (dirtyCount >= writeBatchSize) {
         debug("Dirty count %s >= write batch size %s. Calling putAll() on all dirty entries." format (dirtyCount, writeBatchSize))
         true
       } else if (hasArrayKeys) {
         // Flush every time to support the following legacy behavior:
         // If array keys are used with a cached store, get() will always miss the cache because of array equality semantics
         // However, it will fall back to the underlying store which does support arrays.
         true
       } else {
         false
       }

       if (purgeNeeded) {
         putAllDirtyEntries()
       }
     })
   }

   override def flush() {
     trace("Purging dirty entries from CachedStore.")
     lock.synchronized({
       metrics.flushes.inc
       putAllDirtyEntries()
       store.flush()
     })
     trace("Flushed store.")
   }

   /**
    * The synchronization lock must be held before calling this method.
    */
   private def putAllDirtyEntries() {
     trace("Calling putAll() on dirty entries.")
     // write out the contents of the dirty list oldest first
     val batch = new Array[Entry[K, V]](this.dirtyCount)
     var pos : Int = this.dirtyCount - 1
     for (k <- this.dirty) {
       val entry = this.cache.get(k)
       entry.dirty = null // not dirty any more
       batch(pos) = new Entry(k, entry.value)
       pos -= 1
     }
     store.putAll(Arrays.asList(batch : _*))
     metrics.putAllDirtyEntriesBatchSize.inc(batch.size)
     // reset the dirty list
     this.dirty = new mutable.DoubleLinkedList[K]()
     this.dirtyCount = 0
   }

   override def putAll(entries: java.util.List[Entry[K, V]]) {
     lock.synchronized({
       val iter = entries.iterator
       while (iter.hasNext) {
         val curr = iter.next
         put(curr.getKey, curr.getValue)
       }
     })
   }

   override def delete(key: K) {
     lock.synchronized({
       metrics.deletes.inc
       put(key, null.asInstanceOf[V])
     })
   }

   override def close() {
     lock.synchronized({
       trace("Closing.")
       flush()
       store.close()
     })
   }

   override def deleteAll(keys: java.util.List[K]) = {
     lock.synchronized({
       super.deleteAll(keys)
     })
   }

   private def checkKeyIsArray(key: K) {
     if (!containsArrayKeys && key.isInstanceOf[Array[_]]) {
       // Warn the first time that we see an array key.
       warn("Using arrays as keys results in unpredictable behavior since cache is implemented with a map. Consider using ByteBuffer, or a different key type, or turn off the cache altogether.")
       containsArrayKeys = true
     }
   }

   override def getAll(keys: java.util.List[K]): java.util.Map[K, V] = {
     lock.synchronized({
       metrics.gets.inc(keys.size)
       val returnValue = new java.util.HashMap[K, V](keys.size)
       val misses = new java.util.ArrayList[K]
       val keysIterator = keys.iterator
       while (keysIterator.hasNext) {
         val key = keysIterator.next
         val cached = cache.get(key)
         if (cached != null) {
           metrics.cacheHits.inc
           returnValue.put(key, cached.value)
         } else {
           misses.add(key)
         }
       }
       if (!misses.isEmpty) {
         val entryIterator = store.getAll(misses).entrySet.iterator
         while (entryIterator.hasNext) {
           val entry = entryIterator.next
           returnValue.put(entry.getKey, entry.getValue)
           cache.put(entry.getKey, new CacheEntry(entry.getValue, null))
         }
         cacheCount = cache.size // update outside the loop since it's used for metrics and not for time-sensitive logic
       }
       returnValue
     })
   }

   def hasArrayKeys = containsArrayKeys

   override def snapshot(from: K, to: K): KeyValueSnapshot[K, V] = {
     store.snapshot(from, to)
   }

   override def checkpoint(id: CheckpointId): Optional[Path] = {
     store.checkpoint(id)
   }

   @VisibleForTesting
   private[kv] def getStore: KeyValueStore[K, V] = {
     store
   }
 }

 private case class CacheEntry[K, V](var value: V, var dirty: mutable.DoubleLinkedList[K])
	/*
	* 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 org.apache.samza.util.Logging

	import scala.collection._
	import java.nio.file.Path
	import java.util.{Arrays, Optional}

	import com.google.common.annotations.VisibleForTesting
	import org.apache.samza.checkpoint.CheckpointId

	/**
	* A write-behind caching layer around the rocksdb store. The purpose of this cache is three-fold:
	* 1. Batch together writes to rocksdb, this turns out to be a great optimization
	* 2. Avoid duplicate writes and duplicate log entries within a commit interval. i.e. if there are two updates to the same key, log only the later.
	* 3. Avoid deserialization cost for gets on very common keys (unless the keys itself are an array of bytes,
	* in which case they don't need to be deserialized and there is no performance benefit).
	*
	* This caching does introduce a few odd corner cases :-(
	* 1. Items in the cache have pass-by-reference semantics but items in rocksdb have pass-by-value semantics. Modifying items after a put is a bad idea.
	* 2. Range queries require flushing the cache (as the ordering comes from rocksdb)
	*
	* In implementation this cache is just an LRU hash map that discards the oldest entry when full. There is an accompanying "dirty list" that references keys
	* that have not yet been written to disk. All writes go to the dirty list and when the list is long enough we call putAll on all those values at once. Dirty items
	* that time out of the cache before being written will also trigger a putAll of the dirty list.
	*
	* This class is thread safe.
	*
	* @param store The store to cache
	* @param cacheSize The number of entries to hold in the in memory-cache
	* @param writeBatchSize The number of entries to batch together before forcing a write
	* @param metrics The metrics recording object for this cached store
	*/
	class CachedStore[K, V](
	val store: KeyValueStore[K, V],
	val cacheSize: Int,
	val writeBatchSize: Int,
	val metrics: CachedStoreMetrics = new CachedStoreMetrics) extends KeyValueStore[K, V] with Logging {

	/** the number of items in the dirty list */
	@volatile private var dirtyCount = 0

	/** the number of items currently in the cache */
	@volatile private var cacheCount = 0

	/** the list of items to be written out on flush from newest to oldest */
	private var dirty = new mutable.DoubleLinkedList[K]()

	/ the synchronization lock to protect access to the store from multiple threads /
	private val lock = new Object

	/** an lru cache of values that holds cacheEntries and calls putAll() on dirty entries if necessary when discarding */
	private val cache = new java.util.LinkedHashMap[K, CacheEntry[K, V]]((cacheSize * 1.2).toInt, 1.0f, true) {
	override def removeEldestEntry(eldest: java.util.Map.Entry[K, CacheEntry[K, V]]): Boolean = {
	val evict = super.size > cacheSize
	if (evict) {
	val entry = eldest.getValue
	// if this entry hasn't been written out yet, flush it and all other dirty keys
	if (entry.dirty != null) {
	debug("Found a dirty entry. Calling putAll() on all dirty entries.")
	putAllDirtyEntries()
	}
	}
	evict
	}
	}

	/ tracks whether an array has been used as a key. since this is dangerous with LinkedHashMap, we want to warn on it. /
	@volatile private var containsArrayKeys = false

	// Use counters here, rather than directly accessing variables using .size
	// since metrics can be accessed in other threads, and cache.size is not
	// thread safe since we're using a LinkedHashMap, and dirty.size is slow
	// since it requires a full traversal of the linked list.
	metrics.setDirtyCount(() => dirtyCount)
	metrics.setCacheSize(() => cacheCount)

	override def get(key: K) = lock.synchronized({
	metrics.gets.inc

	val c = cache.get(key)
	if (c != null) {
	metrics.cacheHits.inc
	c.value
	} else {
	val v = store.get(key)
	cache.put(key, new CacheEntry(v, null))
	cacheCount = cache.size
	v
	}
	})

	private class CachedStoreIterator(val iter: KeyValueIterator[K, V])
	extends KeyValueIterator[K, V] {

	var last: Entry[K, V] = null

	override def close(): Unit = iter.close()

	override def remove(): Unit = throw new UnsupportedOperationException("CachedStore iterator doesn't support remove")

	override def next() = {
	last = iter.next()
	last
	}

	override def hasNext: Boolean = iter.hasNext
	}

	override def range(from: K, to: K): KeyValueIterator[K, V] = lock.synchronized({
	metrics.ranges.inc
	putAllDirtyEntries()

	new CachedStoreIterator(store.range(from, to))
	})

	override def all(): KeyValueIterator[K, V] = lock.synchronized({
	metrics.alls.inc
	putAllDirtyEntries()

	new CachedStoreIterator(store.all())
	})

	override def put(key: K, value: V) {
	lock.synchronized({
	metrics.puts.inc

	checkKeyIsArray(key)

	// Add the key to the front of the dirty list (and remove any prior
	// occurrences to dedupe).
	val found = cache.get(key)
	if (found == null \|\| found.dirty == null) {
	this.dirtyCount += 1
	} else {
	// If we are removing the head of the list, move the head to the next
	// element. See SAMZA-45 for details.
	if (found.dirty.prev == null) {
	this.dirty = found.dirty.next
	this.dirty.prev = null
	} else {
	found.dirty.remove()
	}
	}
	this.dirty = new mutable.DoubleLinkedList(key, this.dirty)

	// Add the key to the cache (but don't allocate a new cache entry if we
	// already have one).
	if (found == null) {
	cache.put(key, new CacheEntry(value, this.dirty))
	cacheCount = cache.size
	} else {
	found.value = value
	found.dirty = this.dirty
	}

	// putAll() dirty values if the write list is full.
	val purgeNeeded = if (dirtyCount >= writeBatchSize) {
	debug("Dirty count %s >= write batch size %s. Calling putAll() on all dirty entries." format (dirtyCount, writeBatchSize))
	true
	} else if (hasArrayKeys) {
	// Flush every time to support the following legacy behavior:
	// If array keys are used with a cached store, get() will always miss the cache because of array equality semantics
	// However, it will fall back to the underlying store which does support arrays.
	true
	} else {
	false
	}

	if (purgeNeeded) {
	putAllDirtyEntries()
	}
	})
	}

	override def flush() {
	trace("Purging dirty entries from CachedStore.")
	lock.synchronized({
	metrics.flushes.inc
	putAllDirtyEntries()
	store.flush()
	})
	trace("Flushed store.")
	}

	/**
	* The synchronization lock must be held before calling this method.
	*/
	private def putAllDirtyEntries() {
	trace("Calling putAll() on dirty entries.")
	// write out the contents of the dirty list oldest first
	val batch = new Array[Entry[K, V]](this.dirtyCount)
	var pos : Int = this.dirtyCount - 1
	for (k <- this.dirty) {
	val entry = this.cache.get(k)
	entry.dirty = null // not dirty any more
	batch(pos) = new Entry(k, entry.value)
	pos -= 1
	}
	store.putAll(Arrays.asList(batch : _*))
	metrics.putAllDirtyEntriesBatchSize.inc(batch.size)
	// reset the dirty list
	this.dirty = new mutable.DoubleLinkedList[K]()
	this.dirtyCount = 0
	}

	override def putAll(entries: java.util.List[Entry[K, V]]) {
	lock.synchronized({
	val iter = entries.iterator
	while (iter.hasNext) {
	val curr = iter.next
	put(curr.getKey, curr.getValue)
	}
	})
	}

	override def delete(key: K) {
	lock.synchronized({
	metrics.deletes.inc
	put(key, null.asInstanceOf[V])
	})
	}

	override def close() {
	lock.synchronized({
	trace("Closing.")
	flush()
	store.close()
	})
	}

	override def deleteAll(keys: java.util.List[K]) = {
	lock.synchronized({
	super.deleteAll(keys)
	})
	}

	private def checkKeyIsArray(key: K) {
	if (!containsArrayKeys && key.isInstanceOf[Array[_]]) {
	// Warn the first time that we see an array key.
	warn("Using arrays as keys results in unpredictable behavior since cache is implemented with a map. Consider using ByteBuffer, or a different key type, or turn off the cache altogether.")
	containsArrayKeys = true
	}
	}

	override def getAll(keys: java.util.List[K]): java.util.Map[K, V] = {
	lock.synchronized({
	metrics.gets.inc(keys.size)
	val returnValue = new java.util.HashMap[K, V](keys.size)
	val misses = new java.util.ArrayList[K]
	val keysIterator = keys.iterator
	while (keysIterator.hasNext) {
	val key = keysIterator.next
	val cached = cache.get(key)
	if (cached != null) {
	metrics.cacheHits.inc
	returnValue.put(key, cached.value)
	} else {
	misses.add(key)
	}
	}
	if (!misses.isEmpty) {
	val entryIterator = store.getAll(misses).entrySet.iterator
	while (entryIterator.hasNext) {
	val entry = entryIterator.next
	returnValue.put(entry.getKey, entry.getValue)
	cache.put(entry.getKey, new CacheEntry(entry.getValue, null))
	}
	cacheCount = cache.size // update outside the loop since it's used for metrics and not for time-sensitive logic
	}
	returnValue
	})
	}

	def hasArrayKeys = containsArrayKeys

	override def snapshot(from: K, to: K): KeyValueSnapshot[K, V] = {
	store.snapshot(from, to)
	}

	override def checkpoint(id: CheckpointId): Optional[Path] = {
	store.checkpoint(id)
	}

	@VisibleForTesting
	private[kv] def getStore: KeyValueStore[K, V] = {
	store
	}
	}

	private case class CacheEntry[K, V](var value: V, var dirty: mutable.DoubleLinkedList[K])