Optimize KeyedLock and related concurrency primitives (#96372)

This class was quite hot in recent benchmarks of shared-cached based searches and we can make instantiating the releasable locks a little cheaper. Also, those same benchmarks showed a lot of visible time spent on dealing with ref counts. I removed one layer of indirection in atomic use from both the release-once and the abstract ref count which should save a little in CPU caches as well.

Author: original-brownbear

libs/core/src/main/java/org/elasticsearch/core/Releasables.java

@@ -122,21 +122,7 @@ public String toString() {
  * Wraps a {@link Releasable} such that its {@link Releasable#close()} method can be called multiple times without double-releasing.
  */
  public static Releasable releaseOnce(final Releasable releasable) {
- final var ref = new AtomicReference<>(releasable);
- return new Releasable() {
- @Override
- public void close() {
- final var acquired = ref.getAndSet(null);
- if (acquired != null) {
- acquired.close();
- }
- }
-
- @Override
- public String toString() {
- return "releaseOnce[" + ref.get() + "]";
- }
- };
+ return new ReleaseOnce(releasable);
  }
 
  public static Releasable assertOnce(final Releasable delegate) {
@@ -165,4 +151,23 @@ public String toString() {
  return delegate;
  }
  }
+
+ private static class ReleaseOnce extends AtomicReference<Releasable> implements Releasable {
+ ReleaseOnce(Releasable releasable) {
+ super(releasable);
+ }
+
+ @Override
+ public void close() {
+ final var acquired = getAndSet(null);
+ if (acquired != null) {
+ acquired.close();
+ }
+ }
+
+ @Override
+ public String toString() {
+ return "releaseOnce[" + get() + "]";
+ }
+ }
 }

server/src/main/java/org/elasticsearch/common/util/concurrent/KeyedLock.java

@@ -40,7 +40,6 @@ public Releasable acquire(T key) {
  return newLock;
  }
  } else {
- assert perNodeLock != null;
  int i = perNodeLock.count.get();
  if (i > 0 && perNodeLock.count.compareAndSet(i, i + 1)) {
  perNodeLock.lock();
@@ -104,10 +103,9 @@ private void release(T key, KeyLock lock) {
  assert decrementAndGet >= 0 : decrementAndGet + " must be >= 0 but wasn't";
  }
 
- private final class ReleasableLock implements Releasable {
+ private final class ReleasableLock extends AtomicBoolean implements Releasable {
  final T key;
  final KeyLock lock;
- final AtomicBoolean closed = new AtomicBoolean();
 
  private ReleasableLock(T key, KeyLock lock) {
  this.key = key;
@@ -116,7 +114,7 @@ private ReleasableLock(T key, KeyLock lock) {
 
  @Override
  public void close() {
- if (closed.compareAndSet(false, true)) {
+ if (compareAndSet(false, true)) {
  release(key, lock);
  }
  }