add put() and getIfPresent() to DDCache

internal-api/src/main/java/datadog/trace/api/cache/CHMCache.java

@@ -11,6 +11,24 @@ public CHMCache(final int initialCapacity) {
     this.chm = new ConcurrentHashMap<>(initialCapacity);
   }
 
+  @Override
+  public V put(K key, V value) {
+    if (null == key) {
+      return null;
+    }
+
+    return chm.put(key, value);
+  }
+
+  @Override
+  public V getIfPresent(K key) {
+    if (null == key) {
+      return null;
+    }
+
+    return chm.get(key);
+  }
+
   @Override
   public V computeIfAbsent(K key, Function<K, ? extends V> func) {
     if (null == key) {

internal-api/src/main/java/datadog/trace/api/cache/DDCache.java

@@ -3,6 +3,12 @@
 import datadog.trace.api.Function;
 
 public interface DDCache<K, V> {
+  /** Assigns key to value. Returns the previous value or null if no previous value existed */
+  V put(final K key, V value);
 
+  /** Returns the value assigned to this key or null if no such value exists */
+  V getIfPresent(final K key);
+
+  /** Computes a value for key if key has no assignment in this cache */
   V computeIfAbsent(final K key, Function<K, ? extends V> func);
 }

internal-api/src/main/java/datadog/trace/api/cache/FixedSizeCache.java

@@ -99,6 +99,60 @@ public V computeIfAbsent(K key, Function<K, ? extends V> creator) {
     return value;
   }
 
+  @Override
+  public V getIfPresent(K key) {
+    if (key == null) {
+      return null;
+    }
+
+    int h = key.hashCode();
+    for (int i = 1; i <= 3; i++) {
+      int pos = h & mask;
+      Node<K, V> current = elements[pos];
+      if (current != null && key.equals(current.key)) {
+        // we found a cached key, so use that value
+        return current.value;
+      }
+      // slot was occupied by someone else, so try another slot
+      h = rehash(h);
+    }
+    return null;
+  }
+
+  @Override
+  public V put(K key, V value) {
+    if (key == null) {
+      return null;
+    }
+
+    int h = key.hashCode();
+
+    // insert at the first position unless we find another empty slot
+    int insertPos = h & mask;
+    boolean foundEmptyPosition = false;
+
+    // try to find a slot or a match 3 times
+    for (int i = 1; i <= 3; i++) {
+      int pos = h & mask;
+      Node<K, V> current = elements[pos];
+      if (current == null && !foundEmptyPosition) {
+        insertPos = pos;
+        foundEmptyPosition = true;
+      } else if (current != null && key.equals(current.key)) {
+        // we found a cached key, so overwrite that position
+        elements[pos] = new Node<>(key, value);
+
+        // current is no longer in the array
+        return current.value;
+      }
+      // slot was occupied by someone else, so try another slot
+      h = rehash(h);
+    }
+
+    elements[insertPos] = new Node<>(key, value);
+    return null;
+  }
+
   private V createAndStoreValue(K key, Function<K, ? extends V> creator, int pos) {
     V value = creator.apply(key);
     Node<K, V> node = new Node<>(key, value);

internal-api/src/test/groovy/datadog/trace/api/cache/FixedSizeCacheTest.groovy

@@ -10,7 +10,7 @@ import java.util.concurrent.ThreadLocalRandom
 import java.util.concurrent.atomic.AtomicInteger
 
 class FixedSizeCacheTest extends DDSpecification {
-  def "fixed size should store and retrieve values"() {
+  def "fixed size should computeIfAbsent and retrieve values"() {
     setup:
     def fsCache = DDCaches.newFixedSizeCache(15)
     def creationCount = new AtomicInteger(0)
@@ -38,6 +38,75 @@ class FixedSizeCacheTest extends DDSpecification {
     null                      | null           | 3     // do nothing
   }
 
+  def "fixed size should get values if present"() {
+    setup:
+    def fsCache = DDCaches.newFixedSizeCache(15)
+    def creationCount = new AtomicInteger(0)
+    def tvc = new TVC(creationCount)
+    def tk1 = new TKey(1, 1, "one")
+    def tk6 = new TKey(6, 6, "six")
+    def tk10 = new TKey(10, 10, "ten")
+    // insert some values that happen to be the chain of hashes 1 -> 6 -> 10
+    fsCache.computeIfAbsent(tk1, tvc)
+    fsCache.computeIfAbsent(tk6, tvc)
+    fsCache.computeIfAbsent(tk10, tvc)
+
+    expect:
+    fsCache.getIfPresent(tk) == value
+    creationCount.get() == 3
+
+    where:
+    tk                        | value
+    new TKey(1, 1, "one")     | "one_value"
+    new TKey(6, 6, "six")     | "six_value"
+    new TKey(10, 10, "ten")   | "ten_value"
+    new TKey(1, 11, "eleven") | null
+    null                      | null
+  }
+
+  def "fixed size put values"() {
+    setup:
+    def fsCache = DDCaches.newFixedSizeCache(15)
+    def creationCount = new AtomicInteger(0)
+    def tvc = new TVC(creationCount)
+
+    expect:
+    fsCache.put(null, "value") == null
+    fsCache.getIfPresent(null) == null
+
+    when:
+    def tk1 = new TKey(1, 1, "one")
+    fsCache.put(tk1, "put_one")
+
+    then:
+    fsCache.getIfPresent(tk1) == "put_one"
+
+    when:
+    def tk6 = new TKey(6, 6, "six")
+    fsCache.computeIfAbsent(tk6, tvc)
+
+    then:
+    fsCache.getIfPresent(tk6) == "six_value"
+
+    when: "overwrite value in second slot"
+    def oldValue = fsCache.put(tk6, "put_six")
+
+    then:
+    oldValue == "six_value"
+    fsCache.getIfPresent(tk6) == "put_six"
+    fsCache.computeIfAbsent(tk6, tvc) == "put_six"
+
+    when: "put value in first slot when all slots used"
+    def tk10 = new TKey(10, 10, "ten")
+    fsCache.computeIfAbsent(tk10, tvc)
+    def tk11 = new TKey(1, 11, "eleven")
+    def slotValue = fsCache.put(tk11, "put_eleven")
+
+    then:
+    slotValue == null //slot was occupied but equality failed
+    fsCache.getIfPresent(tk11) == "put_eleven"
+  }
+
   def "chm cache should store and retrieve values"() {
     setup:
     def fsCache = DDCaches.newUnboundedCache(15)
@@ -65,6 +134,64 @@ class FixedSizeCacheTest extends DDSpecification {
     null                      | null           | 3
   }
 
+  def "chm cache should get values if present"() {
+    def fsCache = DDCaches.newUnboundedCache(15)
+    def creationCount = new AtomicInteger(0)
+    def tvc = new TVC(creationCount)
+    def tk1 = new TKey(1, 1, "one")
+    def tk6 = new TKey(6, 6, "six")
+    def tk10 = new TKey(10, 10, "ten")
+    // insert some values that happen to be the chain of hashes 1 -> 6 -> 10
+    fsCache.computeIfAbsent(tk1, tvc)
+    fsCache.computeIfAbsent(tk6, tvc)
+    fsCache.computeIfAbsent(tk10, tvc)
+
+    expect:
+    fsCache.getIfPresent(tk) == value
+    creationCount.get() == 3
+
+    where:
+    tk                        | value
+    new TKey(1, 1, "one")     | "one_value"
+    new TKey(6, 6, "six")     | "six_value"
+    new TKey(10, 10, "ten")   | "ten_value"
+    new TKey(1, 11, "eleven") | null
+    null                      | null
+  }
+
+  def "chm cahce put values"() {
+    setup:
+    def fsCache = DDCaches.newUnboundedCache(15)
+    def creationCount = new AtomicInteger(0)
+    def tvc = new TVC(creationCount)
+
+    expect:
+    fsCache.put(null, "value") == null
+    fsCache.getIfPresent(null) == null
+
+    when:
+    def tk1 = new TKey(1, 1, "one")
+    fsCache.put(tk1, "put_one")
+
+    then:
+    fsCache.getIfPresent(tk1) == "put_one"
+
+    when:
+    def tk6 = new TKey(6, 6, "six")
+    fsCache.computeIfAbsent(tk6, tvc)
+
+    then:
+    fsCache.getIfPresent(tk6) == "six_value"
+
+    when:
+    def oldValue = fsCache.put(tk6, "put_six")
+
+    then:
+    oldValue == "six_value"
+    fsCache.getIfPresent(tk6) == "put_six"
+    fsCache.computeIfAbsent(tk6, tvc) == "put_six"
+  }
+
   def "should handle concurrent usage"() {
     setup:
     def numThreads = 5