MapAtomicCell

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

Start by reviewing the existing AtomicCell and MapRef APIs, then compare them with the proposed MapAtomicCell interface and ConcurrentHashMap implementation in the issue. Clarify the intended scope, including whether the sharded variant is required; done should mean an agreed API, implementation, and verification for concurrent per-key effectful updates.