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))
      }
  }
}

[reardonj]

Personally, I'd want to look for a more robust cache solution instead of something like this (ie. LRU or entry timeout stuff).

On your implementation, why Async? It looks like all you're using is delay which is in Sync.

---

Aside: I have to wonder if there is a way to abstract these Map* things, it's so much boilerplate 🙍‍♂️