package kafka.utils
import java.util.ArrayList
import java.util.concurrent._
import collection.mutable
import collection.JavaConversions
import kafka.common.KafkaException
import java.lang.Object
class Pool[K,V](valueFactory: Option[(K) => V] = None) extends Iterable[(K, V)] {
private val pool = new ConcurrentHashMap[K, V]
private val createLock = new Object
def this(m: collection.Map[K, V]) {
this()
m.foreach(kv => pool.put(kv._1, kv._2))
}
def put(k: K, v: V) = pool.put(k, v)
def putIfNotExists(k: K, v: V) = pool.putIfAbsent(k, v)
def getAndMaybePut(key: K) = {
if (valueFactory.isEmpty)
throw new KafkaException("Empty value factory in pool.")
val curr = pool.get(key)
if (curr == null) {
createLock synchronized {
val curr = pool.get(key)
if (curr == null)
pool.put(key, valueFactory.get(key))
pool.get(key)
}
}
else
curr
}
def contains(id: K) = pool.containsKey(id)
def get(key: K): V = pool.get(key)
def remove(key: K): V = pool.remove(key)
def keys: mutable.Set[K] = {
import JavaConversions._
pool.keySet()
}
def values: Iterable[V] = {
import JavaConversions._
new ArrayList[V](pool.values())
}
def clear() { pool.clear() }
override def size = pool.size
override def iterator = new Iterator[(K,V)]() {
private val iter = pool.entrySet.iterator
def hasNext: Boolean = iter.hasNext
def next: (K, V) = {
val n = iter.next
(n.getKey, n.getValue)
}
}
}