util-core: Add Memoize.classValue

Problem

We'd like to use `java.lang.ClassValue`, but its Java API looks cumbersome in
Scala as it requires subclassing and overriding a protected method for every
class value.

Solution

Add a `Memoize.classValue` method for `Class[_] => T` functions that resembles
ordinary memoization in its API but uses `ClassValue` internally.

Differential Revision: https://phabricator.twitter.biz/D825673

Author: szegedi

CHANGELOG.rst

@@ -10,6 +10,8 @@ Unreleased
 New Features
 ~~~~~~~~~~~~
 
+* util-core: Added `Memoize.classValue` as a Scala-friendly API for `java.lang.ClassValue`. ``PHAB_ID=D825673``
+
 * util-jvm: Register JVM expression including memory pool usages (including code cache, compressed class space,
   eden space, sheap, metaspace, survivor space, and old gen) and open file descriptors count in StatsReceiver.
   ``PHAB_ID=D820472``

util-core/src/main/scala/com/twitter/util/Memoize.scala

@@ -164,4 +164,21 @@ object Memoize {
           case _ => missing(a)
         }
     }
+
+  /**
+   * Thread-safe memoization for a function taking Class objects,
+   * using [[java.lang.ClassValue]].
+   *
+   * @param f the function to memoize using a `ClassValue`
+   * @tparam T the return type of the function
+   * @return a memoized function over classes that stores memoized values
+   *         in a `ClassValue`
+   */
+  def classValue[T](f: Class[_] => T): Class[_] => T = {
+    val cv = new ClassValue[T] {
+      def computeValue(c: Class[_]): T = f(c)
+    }
+
+    cv.get
+  }
 }

util-core/src/test/scala/com/twitter/util/MemoizeTest.scala

@@ -2,7 +2,9 @@ package com.twitter.util
 
 import com.twitter.conversions.DurationOps._
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.concurrent.{TimeUnit, CountDownLatch => JavaCountDownLatch}
+import java.util.concurrent.atomic.AtomicReference
+import java.util.concurrent.{CountDownLatch => JavaCountDownLatch}
+import java.util.concurrent.TimeUnit
 import org.mockito.Mockito._
 import org.scalatest.funsuite.AnyFunSuite
 
@@ -185,4 +187,46 @@ class MemoizeTest extends AnyFunSuite {
     assert(memoizer.size == memoizer.snap.size)
     assert(memoizer.snap.size == 1)
   }
+
+  test("Memoize.classValue: returns the same value for the same input") {
+    val cv = Memoize.classValue { _ =>
+      new Object()
+    }
+    // just some anonymous classes
+    val clazz1 = new Object() {}.getClass
+    val clazz2 = new Object() {}.getClass
+    assert(clazz1 ne clazz2)
+
+    val o11 = cv(clazz1)
+    val o12 = cv(clazz1)
+    val o21 = cv(clazz2)
+    val o22 = cv(clazz2)
+    assert(o11 eq o12)
+    assert(o21 eq o22)
+    assert(o11 ne o21)
+  }
+
+  test("Memoize.classValue: racing threads observe the same value") {
+    val Concurrency = 16
+    val ready = new JavaCountDownLatch(Concurrency)
+    val cv = Memoize.classValue { _ =>
+      ready.await(1, TimeUnit.SECONDS)
+      new Object()
+    }
+
+    val clazz = new Object() {}.getClass
+
+    val values = new Array[AtomicReference[Object]](Concurrency)
+    val threads = (0 until Concurrency).map { i =>
+      new Thread(() => {
+        ready.countDown()
+        values(i) = new AtomicReference(cv(clazz))
+      })
+    }
+    threads.foreach { _.start() }
+    threads.foreach { _.join(1000L) }
+
+    val v0 = values(0).get()
+    assert(values.forall(_.get eq v0))
+  }
 }