Immutable cache implementation

cache-sample-2.scala

import scala.annotation.tailrec
import scala.collection.SortedSet
import scala.concurrent.duration.{Duration, _}

import Cache._
import Cache.ExpiryInfo.ordering


/**
 * An immutable cache implementation.
 */
case class Cache[+T] private (
  values: Map[String, Value[T]],
  expiryCalendar: SortedSet[ExpiryInfo],
  cleanupInterval: FiniteDuration,
  currentTimestamp: Long,
  lastCleanupTS: Long
) extends (String => T) {
  /**
   * Returns the value associated with `key`, if the value hasn't expired yet.
   *
   * @throws NoSuchElementException in case the key is missing or is expired.
   */
  def apply(key: String): T = {
    values.get(key) match {
      case Some(Value(value, expiresAt)) =>
        if (expiresAt > currentTimestamp)
          value
        else
          throw new NoSuchElementException(s"$key is expired in the cache")

      case None =>
        throw new NoSuchElementException(s"$key is missing from the cache")
    }
  }

  /**
   * Fetches a key from the cache. Returns `None` in case the key does not
   * exists or is expired.
   *
   * @param key is the key to fetch for its associated value
   */
  def get(key: String): Option[T] =
    values.get(key).collect {
      case Value(value, expiresAt)
        if expiresAt > currentTimestamp => value
    }

  /**
   * Returns true if the given key exists in the cache and is not expired 
   * (relative to the present moment indicated by `now`), or false otherwise.
   */
  def exists(key: String): Boolean =
    values.get(key) match {
      case None => false
      case Some(value) =>
        value.notExpiredAt(currentTimestamp)
    }

  /**
   * Returns a cache that is free of expired junk.
   *
   * @param forced if true, then forces cleanup, even if we are not at the right time.
   */
  def cleaned(forced: Boolean = false): Cache[T] = {
    if (lastCleanupTS + cleanupInterval.toMillis <= currentTimestamp)
      removeExpired(copy(lastCleanupTS = currentTimestamp), currentTimestamp)
    else
      this
  }

  /**
   * For removing a key, returning a tuple representing our removed
   * value + a new cache instance with the key removed.
   */
  def remove(key: String): (Option[T], Cache[T]) = {
    val cache = cleaned(forced = false)
    cache.values.get(key) match {
      case None => (None, cache)

      case Some(value) =>
        val newExpiries = cache.expiryCalendar - ExpiryInfo(key, value.expiresAt)
        val newValues = cache.values - key
        val result = if (value.notExpiredAt(currentTimestamp)) Some(value.value) else None
        (result, copy(values = newValues, expiryCalendar = newExpiries))
    }
  }

  /**
   * Returns a cache with the given `key` updated by `value`.
   */
  def update[U >: T](key: String, value: U, expiry: Duration): Cache[U] = {
    require(expiry > Duration.Zero, "The given expiry time must be greater than zero")

    val expiresAt = currentTimestamp + (
      if (expiry.isFinite())
        expiry.toMillis
      else
        365.days.toMillis
      )

    val (_, cache) = remove(key)
    val newValues = cache.values.updated(key, Value(value, expiresAt))
    val newCalendar = cache.expiryCalendar + ExpiryInfo(key, expiresAt)
    copy(values = newValues, expiryCalendar = newCalendar)
  }

  /**
   * Returns a cache with the given `key` associated with the given `value`,
   * but only if the `key` is not in cache already or is expired.
   */
  def add[U >: T](key: String, value: U, expiry: Duration): Cache[U] = {
    require(expiry > Duration.Zero, "The given expiry time must be greater than zero")
    if (exists(key)) this
    else
      update(key, value, expiry)
  }

  /**
   * Returns the cache with a new current time applied and thus with
   * expired keys.
   */
  def atMoment(timestamp: Long): Cache[T] = {
    copy(currentTimestamp = timestamp)
  }

  /**
   * Advances the current moment in time by the specified timespan,
   * expiring keys that have to expire at
   * `currentTimestamp` + `timespan`.
   */
  def advanceInTime(timespan: FiniteDuration): Cache[T] = {
    copy(currentTimestamp = currentTimestamp + timespan.toMillis)
      .cleaned(forced = false)
  }

  /**
   * Traverses all the elements in the cache that aren't expired.
   */
  def foreach[B](f: ((String, T)) => B): Unit = {
    for ((key, Value(value, expiresAt)) <- values)
      if (expiresAt > currentTimestamp) f(key, value)
  }

  /**
   * Builds a new cache with the given mapping function applied to 
   * each value that isn't expired.
   */
  def mapValues[U](f: T => U): Cache[U] = {
    val cache = cleaned(forced = true)
    cache.copy(values = 
      cache.values.map { case (key, boxedValue) =>
        (key, boxedValue.copy(value = f(boxedValue.value)))
      })
  }

  /**
   * Builds a new cache with the given mapping function applied to 
   * each value that isn't expired.
   */
  def map[U](f: ((String, T)) => (String, U)): Cache[U] = {
    val source = cleaned(forced = true)
    var newValues = Map.empty[String, Value[U]]
    var newCalendar = SortedSet.empty[ExpiryInfo]
    
    for ((key, boxedVal) <- source.values) {
      val (newKey, newValue) = f((key, boxedVal.value))
      newValues = newValues.updated(newKey, boxedVal.copy(newValue))
      newCalendar = newCalendar + ExpiryInfo(newKey, boxedVal.expiresAt)
    }

    source.copy(newValues, newCalendar)
  }

  /**
   * Returns a union of the source with the given `other` cache,
   * the result containing the non-expired key-value pairs in both.
   */
  def union[U >: T](other: Cache[U]): Cache[U] = {
    val currentTimestamp = math.max(this.currentTimestamp, other.currentTimestamp)
    val sourceCleaned = this.atMoment(currentTimestamp).cleaned(forced = true)
    val otherCleaned = other.atMoment(currentTimestamp).cleaned(forced = true)

    Cache(
      values = sourceCleaned.values ++ otherCleaned.values,
      expiryCalendar = sourceCleaned.expiryCalendar ++ otherCleaned.expiryCalendar,
      cleanupInterval = cleanupInterval,
      currentTimestamp = currentTimestamp,
      lastCleanupTS = currentTimestamp
    )
  }

  /**
   * Alias for [[Cache.union]].
   */
  def ++[U >: T](other: Cache[U]): Cache[U] = {
    union(other)
  }

  /**
   * The flatMap operation takes a function that returns a cache for each
   * element of the source and returns the union of all resulting caches.
   */
  def flatMap[U](f: ((String, T)) => Cache[U]): Cache[U] = {
    var union = Cache.empty[U](cleanupInterval, currentTimestamp)
    for ((key, boxedValue) <- cleaned(forced = true).values) {
      val newCache = f((key, boxedValue.value))
      union = union ++ newCache
    }
    
    union
  }

  /**
   * Filters this cache by the given predicate, returning a new cache for
   * which the given predicate holds for all elements.
   */
  def filter[U](f: ((String, T)) => Boolean): Cache[T] = {
    val source = cleaned(forced = true)
    var newValues = source.values
    var newCalendar = SortedSet.empty[ExpiryInfo]

    for ((key, boxedVal) <- source.values)
      if (!f((key, boxedVal.value)))
        newValues = newValues - key
      else
        newCalendar = newCalendar + ExpiryInfo(key, boxedVal.expiresAt)

    source.copy(newValues, newCalendar)
  }

  /**
   * Filters this cache by the given predicate, returning a new cache for
   * which the given predicate does not hold for any of the elements.
   */
  def filterNot[U](f: ((String, T)) => Boolean): Cache[T] = {
    val source = cleaned(forced = true)
    var newValues = source.values
    var newCalendar = SortedSet.empty[ExpiryInfo]

    for ((key, boxedVal) <- source.values) {
      if (f((key, boxedVal.value)))
        newValues = newValues - key
      else
        newCalendar = newCalendar + ExpiryInfo(key, boxedVal.expiresAt)
    }

    source.copy(newValues, newCalendar)
  }

  /**
   * Returns a map with all the key-value pairs that haven't expired yet.
   */
  def toMap: Map[String, T] = {
    var map = Map.empty[String, T]
    for ((key, boxed @ Value(value, _)) <- values)
      if (boxed.notExpiredAt(currentTimestamp)) {
        map = map.updated(key, value)
      }

    map
  }

  /**
   * Returns an iterable with all the key-value pairs that haven't expired yet.
   */
  def toIterable: Iterable[(String, T)] =
    new Iterable[(String, T)] {
      def iterator = new Iterator[(String, T)] {
        private[this] val cursor = values.iterator
          .filter { case (_, boxed) => boxed.notExpiredAt(currentTimestamp) }

        def hasNext =
          cursor.hasNext

        def next() = {
          val (key, Value(value, _)) = cursor.next()
          (key, value)
        }
      }
    }

  override lazy val toString = {
    val iterator = toIterable.iterator
    val first100 = iterator.take(20).map(x => s"${x._1} -> ${x._2}").mkString(", ")
    val elems = first100 + (if (iterator.hasNext) ", ..." else "")
    "Cache(" + elems + ")"
  }
}


object Cache {
  /**
   * Builds a cache out of the specified sequence of key and value pairs, where
   * every item has the expiration specified by `expiry`.
   */
  def apply[T](expiry: Duration)(seq: (String, T)*)(implicit config: InitConfig): Cache[T] = {
    var cache = Cache.empty[T]
    for ((key, value) <- seq)
      cache = cache.update(key, value, expiry)
    cache
  }

  /**
   * Returns an empty [[Cache]].
   */
  def empty[T](implicit config: InitConfig): Cache[T] = {
    val now = config.currentMoment
    Cache(
      values = Map.empty[String, Value[T]],
      expiryCalendar = SortedSet.empty[ExpiryInfo],
      cleanupInterval = config.cleanupInterval,
      lastCleanupTS = now,
      currentTimestamp = now
    )
  }

  /**
   * Returns an empty [[Cache]].
   *
   * @param cleanupInterval is the interval of time at which the [[Cache.cleaned]]
   *                        operation should be triggered on writes.
   *
   * @param now indicates the present moment.
   */
  def empty[T](cleanupInterval: FiniteDuration, now: Long = System.currentTimeMillis()): Cache[T] = Cache(
    values = Map.empty[String, Value[T]],
    expiryCalendar = SortedSet.empty[ExpiryInfo],
    cleanupInterval = cleanupInterval,
    lastCleanupTS = 0,
    currentTimestamp = now
  )

  /**
   * Taken as an implicit when building a new cache, to spare one for manually
   * specifying the current moment and the cleanup interval every time you do it.
   */
  trait InitConfig {
    def cleanupInterval: FiniteDuration
    def currentMoment: Long
  }

  object InitConfig {
    /**
     * Default init configuration.
     */
    implicit val default = new InitConfig {
      def cleanupInterval = 1.second
      def currentMoment = System.currentTimeMillis()
    }
  }

  /**
   * Tuple representing a value with an expiration date attached.
   */
  case class Value[+T](value: T, expiresAt: Long) {
    def isExpiredAt(timestamp: Long): Boolean = 
      !notExpiredAt(timestamp)
    
    def notExpiredAt(timestamp: Long): Boolean =
      expiresAt > timestamp
  }

  /**
   * A tuple representing the timestamp at which keys should expire,
   * given in UTC timestamps.
   */
  case class ExpiryInfo(key: String, expiresAt: Long)

  object ExpiryInfo {
    implicit val ordering: Ordering[ExpiryInfo] =
      new Ordering[ExpiryInfo] {
        def compare(x: ExpiryInfo, y: ExpiryInfo) = {
          if (x.expiresAt < y.expiresAt) -1
          else if (x.expiresAt > y.expiresAt) 1
          else x.key.compareTo(y.key)
        }
      }
  }

  @tailrec
  private def removeExpired[T](cache: Cache[T], now: Long): Cache[T] = {
    val calendar = cache.expiryCalendar

    if (calendar.isEmpty) cache else {
      val expiry = calendar.firstKey
      if (expiry.expiresAt > now)
        cache
      else
        removeExpired(now = now, cache = cache.copy(
          values = cache.values - expiry.key,
          expiryCalendar = calendar.tail
        ))
    }
  }
}