Scala TrieMap, in depth: the Ctrie, lock-free updates and constant-time snapshots

The problem TrieMap solves

Start with what goes wrong without it. A plain mutable.HashMap shared between threads can lose updates or corrupt its internal table. Wrapping it in synchronized fixes correctness but serialises every reader behind every writer. Java's ConcurrentHashMap is far better: reads take no lock and writes lock only one bin. But its iterators are weakly consistent. They never fail, yet they may or may not reflect writes that happen during iteration, so two passes over the map, say one to count and one to sum, can describe two different maps.

TrieMap answers a different question: what if readers need a stable view of the entire map while writers continue? Its snapshot makes an exact point-in-time copy available immediately, and the copying cost is paid lazily, only for the parts of the trie that are later modified. Iteration, size and bulk reads all run against such a snapshot.

Inside the Ctrie

A hash trie is a tree indexed by the bits of the key's hash. At the root, the lowest five bits pick one of 32 possible branches; at the next level, the next five bits; and so on. Since 32 slots per node would waste memory when most are empty, each branch node stores a 32-bit bitmap plus a dense array holding only the branches that exist. To find slot i, check bit i in the bitmap; the array position is the number of set bits below it, one popcount instruction.

The Ctrie makes this concurrent by making every node immutable except one kind:

With a well-mixed hash, a map of n entries is about log32(n) levels deep: four levels hold a million keys. A read walks down that path with volatile reads and no CAS, so lookups are fast and never wait. Each update allocates a new CNode (up to 32 references) and usually a new SNode, which is the main cost compared with an in-place table.

Insertion with a single CAS

Insertion shows the core trick. The thread walks down to the INode whose CNode should contain the key, builds a modified copy of that CNode, and tries to swing the INode's reference from the old CNode to the new one with a single compare-and-swap. If another thread changed that INode first, the CAS fails and the thread retries from the root. In pseudocode:

insert(key, value):
  hc = hash(key)
  loop:                                     # restart point after a failed CAS
    inode = root; level = 0
    while true:
      main = inode.main                     # volatile read
      match main:
        CNode(bitmap, array):
          idx  = (hc >>> level) & 0x1f       # five bits for this level
          flag = 1 << idx
          pos  = popcount(bitmap & (flag - 1))
          if bitmap & flag == 0:            # empty slot: add a new SNode
            if CAS(inode.main, main, main.inserted(pos, flag, SNode(key, value, hc))): return
            else: continue loop
          match array[pos]:
            INode child:  inode = child; level += 5     # descend
            SNode s if s.key == key:                    # replace the value
              if CAS(inode.main, main, main.updated(pos, SNode(key, value, hc))): return
              else: continue loop
            SNode s:                                    # different key, same 5 bits: push down a level
              sub = INode(CNode.dual(s, SNode(key, value, hc), level + 5))
              if CAS(inode.main, main, main.updated(pos, sub)): return
              else: continue loop
        TNode:  clean(parent); continue loop             # help finish a pending removal
        LNode:  CAS(inode.main, main, main.inserted(key, value)) or continue loop

Every change is published by one CAS on one INode, so readers see either the old CNode or the new one, never a partial state. Removal is the mirror image, with one extra step: if a CNode shrinks to a single SNode, it is replaced by a TNode so the parent can compress the path, and any thread that encounters a TNode helps finish that compression. That helping is what makes the structure lock-free rather than merely non-blocking for readers.

Constant-time snapshots

Snapshots are the reason TrieMap exists. Each INode carries a generation tag. Taking a snapshot replaces the root with a new INode of a new generation, using a double-compare single-swap (RDCSS) on the root so that it cannot race with a concurrent root update. Now two roots, old and new, share the entire trie below them. Nothing has been copied.

The copying happens lazily. When any writer descends through an INode whose generation is older than the root it started from, it first replaces that INode with a copy tagged with the current generation, then continues. The CAS used throughout is a generation-aware variant (GCAS) that commits only if the root's generation has not changed since the operation began; if a snapshot slipped in, the operation aborts and retries against the new generation. The effect is copy-on-write at the granularity of trie paths: a snapshot costs O(1) to take, and later writes pay a few extra node copies the first time they touch each old path.

Two flavours exist. snapshot() returns a new, fully mutable TrieMap that diverges from the original. readOnlySnapshot() returns a read-only view, which is cheaper because only the original needs to copy paths on later writes. TrieMap's own iterator and size use a read-only snapshot internally, so iteration is point-in-time consistent. The catch is that size must count the snapshot, which is linear in the map size rather than a stored counter; avoid it on hot paths.

The API and its atomicity contract

TrieMap implements scala.collection.concurrent.Map, whose atomic operations are the ones to build on:

That last row matters. Unlike ConcurrentHashMap.computeIfAbsent, which runs the function at most once per key by holding the bin lock while it runs, TrieMap never blocks, so racing threads may each evaluate op. That is fine for cheap, pure computations and wrong for expensive or side-effecting ones. Compound operations written as separate calls, such as if (!m.contains(k)) m.put(k, v), are not atomic and are the most common TrieMap bug in code review. Read-modify-write with m(k) = m(k) + 1 loses increments under contention; use updateWith or a replace loop instead.

Worked examples: a registry and a compute-once cache

Two patterns cover most real uses. The first is a live registry that is updated constantly and reported on periodically. Using one snapshot per report guarantees that the count and the totals describe the same moment:

import scala.collection.concurrent.TrieMap

final case class Session(user: String, startedAt: Long, bytes: Long)

object Sessions {
  private val live = TrieMap.empty[String, Session]

  def open(id: String, user: String): Unit =
    live.putIfAbsent(id, Session(user, System.currentTimeMillis(), 0L))

  def addBytes(id: String, n: Long): Unit =
    live.updateWith(id)(_.map(s => s.copy(bytes = s.bytes + n)))   // atomic read-modify-write

  def close(id: String): Unit = live.remove(id)

  def report(): (Int, Long, Map[String, Long]) = {
    val snap  = live.readOnlySnapshot()        // O(1); writers keep going
    val total = snap.valuesIterator.map(_.bytes).sum
    val byUser = snap.values.groupMapReduce(_.user)(_.bytes)(_ + _)
    (snap.size, total, byUser)                 // all three describe the same instant
  }
}

The second is a compute-once cache for expensive work. Because getOrElseUpdate may evaluate its argument more than once, store a cheap placeholder, a Promise's future, and let only the thread that won the putIfAbsent start the computation:

import scala.collection.concurrent.TrieMap
import scala.concurrent.{ExecutionContext, Future, Promise}

final class OnceCache[K, V](compute: K => Future[V])(implicit ec: ExecutionContext) {
  private val entries = TrieMap.empty[K, Future[V]]

  def get(key: K): Future[V] =
    entries.get(key) match {
      case Some(f) => f                                    // fast path: no allocation
      case None =>
        val p = Promise[V]()
        entries.putIfAbsent(key, p.future) match {
          case Some(existing) => existing                  // another thread won the race
          case None =>
            p.completeWith(compute(key))                   // only the winner computes
            p.future.failed.foreach(_ => entries.remove(key, p.future))  // do not cache failures
            p.future
        }
    }
}

Losing threads allocate one unused promise and nothing more. The conditional remove(key, p.future) clears a failed entry without deleting a newer successful one. Add a size bound or expiry before using this for unbounded key spaces; TrieMap has no eviction. The general memoisation trade-offs are covered in memoization in Scala.

TrieMap versus ConcurrentHashMap

Choose TrieMap when you need consistent snapshots or iteration of a map that is written concurrently, or want strictly non-blocking behaviour. Choose ConcurrentHashMap (via asScala if you like) for write-heavy maps where allocation and GC pressure dominate, or when compute-once semantics under a lock are what you want. Benchmark with your own key distribution; Scala collections performance explains how to measure fairly, and lock-free data structures covers the CAS foundations underneath both.

Failure modes and trade-offs

The underlying trade-off is allocation for consistency. Every TrieMap write creates garbage that ConcurrentHashMap would not, and in return you get non-blocking progress and snapshots that no lock-based map can offer cheaply.

What to do next

• Find shared mutable maps in your codebase and classify each by need: snapshots or consistent iteration point to TrieMap, write-heavy caches to ConcurrentHashMap.
• Search TrieMap call sites for check-then-act sequences and replace them with putIfAbsent, replace or updateWith.
• Audit every getOrElseUpdate whose argument is expensive or has side effects; switch to the promise pattern.
• Move size and full iteration out of request paths into periodic snapshot-based reporting.
• Benchmark both maps with your real key distribution and thread count, watching allocation rate as well as throughput.
• Read the Ctrie paper's sections on GCAS and RDCSS if you maintain lock-free code yourself.