HashStore.java

package datadog.trace.api;

import datadog.apicontract.Restricted;
import datadog.trace.api.TagMap.BucketGroup;
import datadog.trace.api.TagMap.Entry;
import datadog.trace.api.TagMap.EntryReader;
import datadog.trace.api.function.TriConsumer;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

/**
 * The custom-tag half of a {@link TagMap}: custom tags by name, in their own bucket array ({@link
 * HashSupport}), reading through to the frozen parent's custom store. Name-keyed methods take only
 * custom tag names (asserted): a known tag is resolved to its id first, and goes to the {@link
 * IdStore}.
 *
 * <p>Created lazily by its owner -- a {@link TagMap}, or a span in the tracer core -- on the first
 * custom write ({@link #create}); until then the owner borrows its parent's custom store, if any,
 * read-only. No locking: guarded by the owner's lock.
 */
@Restricted(allowedIn = {"datadog.trace.core"})
public final class HashStore {
  final Object[] buckets;

  /** Local entries. */
  int count;

  /** The frozen parent's custom store, or null. */
  final HashStore parent;

  /**
   * Parent keys removed locally (read-through tombstones): a tombstone stops read-through for its
   * key. Allocated on the first such removal, so while it is {@code null} the hot paths skip the
   * tombstone checks.
   */
  private Set<String> removedFromParent;

  /**
   * {@code parent}'s custom store, for an owner to borrow read-only until its first custom write,
   * or null.
   */
  public static HashStore inheritedFrom(TagMap parent) {
    return parent == null ? null : parent.nameStore;
  }

  /**
   * A fresh custom store for an owner's first custom write, reading through to {@code borrowed},
   * the store it borrowed until now (may be null).
   */
  public static HashStore create(HashStore borrowed) {
    return new HashStore(borrowed);
  }

  HashStore(HashStore parent) {
    this.buckets = HashSupport.newBuckets();
    this.parent = parent;
  }

  /** Exact count, including the visible part of the read-through chain. */
  int size() {
    return this.parent == null ? this.count : this.count + this.visibleParentCount();
  }

  /** Ancestor entries not shadowed or tombstoned by a nearer level, nearest-level-wins. */
  private int visibleParentCount() {
    // a plain loop, not forEachParent: counting through a lambda would allocate a counter
    int count = 0;
    for (HashStore ancestor = this.parent; ancestor != null; ancestor = ancestor.parent) {
      Object[] parentBuckets = ancestor.buckets;
      for (int i = 0; i < parentBuckets.length; ++i) {
        Object bucket = parentBuckets[i];
        if (bucket instanceof Entry) {
          Entry entry = (Entry) bucket;
          if (this.parentEntryVisible(entry, ancestor)) {
            count++;
          }
        } else if (bucket instanceof BucketGroup) {
          for (BucketGroup group = (BucketGroup) bucket; group != null; group = group.prev) {
            for (int j = 0; j < BucketGroup.LEN; ++j) {
              Entry entry = group._entryAt(j);
              if (entry != null && this.parentEntryVisible(entry, ancestor)) {
                count++;
              }
            }
          }
        }
      }
    }
    return count;
  }

  boolean isEmpty() {
    if (this.count != 0) {
      return false;
    }
    HashStore parent = this.parent;
    if (parent == null) {
      return true;
    }
    if (this.removedFromParent == null) {
      return parent.isEmpty(); // nothing hidden: the parent's tags are all visible
    }
    return this.visibleParentCount() == 0;
  }

  /** Cheap: no level of the chain holds a local entry (ignores shadowing and tombstones). */
  boolean isDefinitelyEmpty() {
    for (HashStore level = this; level != null; level = level.parent) {
      if (level.count != 0) {
        return false;
      }
    }
    return true;
  }

  /** Upper bound: the local counts of every level. */
  int estimateSize() {
    int total = 0;
    for (HashStore level = this; level != null; level = level.parent) {
      total += level.count;
    }
    return total;
  }

  /** The custom tag {@code tag}'s entry, reading through to the parent chain. */
  public Entry get(String tag) {
    assert KnownTagCodec.keyOf(tag) == 0 : "'" + tag + "' is a known tag; look it up by its id";
    return this.get(tag, Entry.bucketHash(Entry.customHash(tag)));
  }

  private Entry get(String tag, int hash) {
    Entry local = HashSupport.find(this.buckets, hash, tag);
    if (local != null) {
      return local;
    }
    HashStore parent = this.parent;
    if (parent == null) {
      return null;
    }
    if (this.removedFromParent != null && this.removedFromParent.contains(tag)) {
      return null; // tombstoned: removed locally, do not read through
    }
    return parent.get(tag, hash);
  }

  /*
   * Sets the custom tag {@code tag}, which must not be a known tag's name.
   */

  public void set(String tag, Object value) {
    if (value == null) {
      this.remove(tag); // a null value removes the tag
      return;
    }
    this.put(Entry.newAnyEntry(tag, value));
  }

  public void set(String tag, CharSequence value) {
    if (value == null) {
      this.remove(tag); // a null value removes the tag
      return;
    }
    this.put(Entry.newObjectEntry(tag, value));
  }

  public void set(String tag, boolean value) {
    this.put(Entry.newBooleanEntry(tag, value));
  }

  public void set(String tag, int value) {
    this.put(Entry.newIntEntry(tag, value));
  }

  public void set(String tag, long value) {
    this.put(Entry.newLongEntry(tag, value));
  }

  public void set(String tag, float value) {
    this.put(Entry.newFloatEntry(tag, value));
  }

  public void set(String tag, double value) {
    this.put(Entry.newDoubleEntry(tag, value));
  }

  /**
   * Inserts or replaces a custom tag's entry, returning the local entry it replaced (or null). Does
   * not consult the read-through parent.
   */
  public Entry put(Entry newEntry) {
    // Re-setting a key clears any read-through tombstone for it (the new value overrides the
    // removal). Gated on the lazy field, so this is a no-op for the common no-tombstone case.
    if (this.removedFromParent != null) {
      this.removedFromParent.remove(newEntry.tag);
    }
    Entry replaced = HashSupport.put(this.buckets, newEntry);
    if (replaced == null) {
      this.count++;
    }
    return replaced;
  }

  /** {@link #put}, returning the prior visible entry -- local, else read through. */
  public Entry getAndPut(Entry newEntry) {
    // Capture whether the key was tombstoned BEFORE put clears it: a tombstoned key had no visible
    // prior value (it was removed), so getAndPut must report null rather than the parent's.
    boolean wasTombstoned =
        this.removedFromParent != null && this.removedFromParent.contains(newEntry.tag);

    Entry priorLocal = this.put(newEntry);
    if (priorLocal != null) {
      return priorLocal;
    }
    if (wasTombstoned || this.parent == null) {
      return null;
    }
    return this.parent.get(newEntry.tag, newEntry.hash());
  }

  /**
   * Removes the custom tag {@code tag}, returning the prior visible entry. Removing a key the
   * parent chain still exposes also installs a tombstone, so it stops reading through.
   */
  public Entry remove(String tag) {
    assert KnownTagCodec.keyOf(tag) == 0 : "'" + tag + "' is a known tag; remove it by its id";
    int hash = Entry.bucketHash(Entry.customHash(tag));
    Entry localRemoved = HashSupport.remove(this.buckets, hash, tag);
    if (localRemoved != null) {
      this.count--;
    }

    HashStore parent = this.parent;
    if (parent != null) {
      boolean alreadyTombstoned =
          this.removedFromParent != null && this.removedFromParent.contains(tag);
      if (!alreadyTombstoned) {
        Entry parentEntry = parent.get(tag, hash);
        if (parentEntry != null) {
          if (this.removedFromParent == null) {
            // Small initial capacity: this set is rare and almost always holds only a handful of
            // tombstoned keys, so the default 16-bucket HashSet table would be wasteful.
            this.removedFromParent = new HashSet<>(4);
          }
          this.removedFromParent.add(tag);
          return localRemoved != null ? localRemoved : parentEntry;
        }
      }
    }
    return localRemoved;
  }

  /**
   * Whether {@code parentEntry}, from the ancestor level {@code fromAncestor}, is visible through
   * this leaf: not shadowed and not tombstoned by any nearer level.
   */
  boolean parentEntryVisible(Entry parentEntry, HashStore fromAncestor) {
    int hash = parentEntry.hash();
    String tag = parentEntry.tag;
    for (HashStore nearer = this; nearer != fromAncestor; nearer = nearer.parent) {
      if (nearer.removedFromParent != null && nearer.removedFromParent.contains(tag)) {
        return false;
      }
      if (HashSupport.find(nearer.buckets, hash, tag) != null) {
        return false;
      }
    }
    return true;
  }

  /** Visits the local entries, then the visible inherited ones. */
  void forEach(Consumer<? super EntryReader> consumer) {
    HashSupport.forEach(this.buckets, consumer);
    if (this.parent != null) {
      this.forEachParent(consumer, null, (c, unused, entry) -> c.accept(entry));
    }
  }

  <T> void forEach(T thisObj, BiConsumer<T, ? super EntryReader> consumer) {
    HashSupport.forEach(this.buckets, thisObj, consumer);
    if (this.parent != null) {
      this.forEachParent(thisObj, consumer, (t, c, entry) -> c.accept(t, entry));
    }
  }

  <T, U> void forEach(T thisObj, U otherObj, TriConsumer<T, U, ? super EntryReader> consumer) {
    HashSupport.forEach(this.buckets, thisObj, otherObj, consumer);
    if (this.parent != null) {
      this.forEachParent(thisObj, otherObj, consumer);
    }
  }

  private <T, U> void forEachParent(
      T thisObj, U otherObj, TriConsumer<T, U, ? super EntryReader> consumer) {
    for (HashStore ancestor = this.parent; ancestor != null; ancestor = ancestor.parent) {
      Object[] parentBuckets = ancestor.buckets;
      for (int i = 0; i < parentBuckets.length; ++i) {
        Object bucket = parentBuckets[i];
        if (bucket instanceof Entry) {
          Entry entry = (Entry) bucket;
          if (this.parentEntryVisible(entry, ancestor)) {
            consumer.accept(thisObj, otherObj, entry);
          }
        } else if (bucket instanceof BucketGroup) {
          for (BucketGroup group = (BucketGroup) bucket; group != null; group = group.prev) {
            for (int j = 0; j < BucketGroup.LEN; ++j) {
              Entry entry = group._entryAt(j);
              if (entry != null && this.parentEntryVisible(entry, ancestor)) {
                consumer.accept(thisObj, otherObj, entry);
              }
            }
          }
        }
      }
    }
  }

  /**
   * An independently mutable copy: the same (frozen, shared) parent, this store's own entries
   * (shared Entry objects) and its tombstones.
   */
  HashStore copy() {
    HashStore copy = new HashStore(this.parent);
    HashSupport.copyInto(copy.buckets, this.buckets);
    copy.count = this.count;
    if (this.removedFromParent != null) {
      copy.removedFromParent = new HashSet<>(this.removedFromParent);
    }
    return copy;
  }

  /** Copies the entries of the parentless store {@code that}, sharing Entry objects. */
  void putAllLocal(HashStore that) {
    if (that.count == 0) {
      return;
    }
    if (this.count == 0) {
      HashSupport.copyInto(this.buckets, that.buckets);
      this.count = that.count;
    } else {
      this.count += HashSupport.merge(this.buckets, that.buckets);
    }
  }

  void fillMapLocal(Map<? super String, Object> map) {
    HashSupport.fillMap(this.buckets, map);
  }

  void fillStringMapLocal(Map<? super String, ? super String> map) {
    HashSupport.fillStringMap(this.buckets, map);
  }

  /** Tests only: throws if the store's structure is inconsistent. */
  void checkIntegrity() {
    HashSupport.checkBuckets(this.buckets);
    if (this.count != this.computeSize()) {
      throw new IllegalStateException("incorrect size");
    }
  }

  /** Tests only: the local count, recomputed from the structure. */
  int computeSize() {
    return HashSupport.count(this.buckets);
  }

  /** For deep debugging. */
  String toInternalString() {
    return HashSupport.toInternalString(this.buckets);
  }
}