MapAtomicCell

[vangogh500]

Been a big fan of Refs and MapRefs provided by cats-effect but always found myself wanting a way to effect-fully modify the value since it is a common production use case e.g. if there is a cache miss go and fetch from x.

Given now we have AtomicCells which allows us to do this I was surprised at the missing MapAtomicCell which would allow for concurrent access to the values but also allow for effectful modifications.

Here is an implementation with example usage I wrote for myself. Was wondering if others would find it useful as well :)

Example:

// Cache that fetches users on miss — concurrent callers for the same key
// serialize (default runs at-most-once), different keys don't block each other.
for {
  cache <- MapAtomicCell[IO, String, User].map(_.withDefaultF(fetchUser))
  user  <- cache("user-123").get
} yield user

// Manual get-or-update without withDefault
for {
  map  <- MapAtomicCell[IO, String, Config]
  conf <- map("service-a").evalModify {
    case Some(c) => IO.pure((Some(c), c))
    case None    => loadConfig("service-a").map(c => (Some(c), c))
  }
} yield conf 

// Per-key atomic updates — e.g. rate limiting
for {
  counters <- MapAtomicCell[IO, String, Int].map(_.withDefaultValue(0))
  count    <- counters("endpoint-a").evalUpdateAndGet(n => IO.pure(n + 1))
} yield count

Implementation (note only did the ConcurrentMap version for now, doing the sharded version like MapRef should be trivial):

import cats.implicits._
import cats.effect.Async
import cats.effect.std.{AtomicCell}

import java.util.concurrent.ConcurrentHashMap

trait MapAtomicCell[F[_], K, V] {
  def apply(key: K): AtomicCell[F, V]

  def withDefault[A](default: K => A)(implicit ev: V =:= Option[A], F: Async[F]): MapAtomicCell[F, K, A] =
    withDefaultF(k => F.pure(default(k)))

  def withDefaultF[A](default: K => F[A])(implicit ev: V =:= Option[A], F: Async[F]): MapAtomicCell[F, K, A] =
    MapAtomicCell.withDefault(this.asInstanceOf[MapAtomicCell[F, K, Option[A]]], default)

  def withDefaultValue[A](value: => A)(implicit ev: V =:= Option[A], F: Async[F]): MapAtomicCell[F, K, A] =
    withDefaultF(_ => F.pure(value))

  def withDefaultValueF[A](value: F[A])(implicit ev: V =:= Option[A], F: Async[F]): MapAtomicCell[F, K, A] =
    withDefaultF(_ => value)
}

object MapAtomicCell {
  def apply[F[_]: Async, K, V]: F[MapAtomicCell[F, K, Option[V]]] =
    ofConcurrentHashMap[F, K, V]

  def ofConcurrentHashMap[F[_]: Async, K, V](
    initialCapacity: Int = 16,
    loadFactor: Float = 0.75f,
    concurrencyLevel: Int = 16
  ): F[MapAtomicCell[F, K, Option[V]]] =
    Async[F]
      .delay(new ConcurrentHashMap[K, AtomicCell[F, Option[V]]](initialCapacity, loadFactor, concurrencyLevel))
      .map(fromConcurrentHashMap[F, K, V])

  def fromConcurrentHashMap[F[_]: Async, K, V](map: ConcurrentHashMap[K, AtomicCell[F, Option[V]]]): MapAtomicCell[F, K, Option[V]] =
    new ConcurrentHashMapImpl[F, K, V](map)

  def withDefault[F[_]: Async, K, V](underlying: MapAtomicCell[F, K, Option[V]], default: K => F[V]): MapAtomicCell[F, K, V] =
    new DefaultImpl[F, K, V](underlying, default)

  private class DefaultImpl[F[_]: Async, K, V](
    underlying: MapAtomicCell[F, K, Option[V]],
    default: K => F[V]
  ) extends MapAtomicCell[F, K, V] {

    private def eval[B](key: K)(f: V => F[(V, B)]): F[B] =
      underlying(key).evalModify {
        case Some(v) => f(v).map { case (next, b) => (next.some, b) }
        case None    => default(key).flatMap(f).map { case (next, b) => (next.some, b) }
      }

    def apply(key: K): AtomicCell[F, V] =
      new AtomicCell[F, V] {
        def get: F[V] = eval(key)(v => (v, v).pure[F])
        def set(a: V): F[Unit] = underlying(key).set(a.some)
        def modify[B](f: V => (V, B)): F[B] = eval(key)(v => f(v).pure[F])
        def evalModify[B](f: V => F[(V, B)]): F[B] = eval(key)(f)
        def evalUpdate(f: V => F[V]): F[Unit] = eval(key)(v => f(v).map(next => (next, ())))
        def evalGetAndUpdate(f: V => F[V]): F[V] = eval(key)(v => f(v).map(next => (next, v)))
        def evalUpdateAndGet(f: V => F[V]): F[V] = eval(key)(v => f(v).map(next => (next, next)))
      }
  }

  private class ConcurrentHashMapImpl[F[_]: Async, K, V](
    chm: ConcurrentHashMap[K, AtomicCell[F, Option[V]]]
  ) extends MapAtomicCell[F, K, Option[V]] {

    private def cellFor(key: K): F[AtomicCell[F, Option[V]]] =
      Async[F].delay(Option(chm.get(key))).flatMap {
        case Some(cell) => cell.pure[F]
        case None =>
          AtomicCell[F].of(none[V]).flatMap { newCell =>
            Async[F].delay(chm.putIfAbsent(key, newCell)).map {
              case null     => newCell
              case existing => existing
            }
          }
      }

    def apply(key: K): AtomicCell[F, Option[V]] =
      new AtomicCell[F, Option[V]] {
        def get: F[Option[V]] = cellFor(key).flatMap(_.get)
        def set(a: Option[V]): F[Unit] = cellFor(key).flatMap(_.set(a))
        def modify[B](f: Option[V] => (Option[V], B)): F[B] = cellFor(key).flatMap(_.modify(f))
        def evalModify[B](f: Option[V] => F[(Option[V], B)]): F[B] = cellFor(key).flatMap(_.evalModify(f))
        def evalUpdate(f: Option[V] => F[Option[V]]): F[Unit] = cellFor(key).flatMap(_.evalUpdate(f))
        def evalGetAndUpdate(f: Option[V] => F[Option[V]]): F[Option[V]] = cellFor(key).flatMap(_.evalGetAndUpdate(f))
        def evalUpdateAndGet(f: Option[V] => F[Option[V]]): F[Option[V]] = cellFor(key).flatMap(_.evalUpdateAndGet(f))
      }
  }
}