LUCENE-9001: Fix race condition in SetOnce (#931)

lucene/CHANGES.txt

@@ -25,8 +25,8 @@ Optimizations
   (Ignacio Vera, Adrien Grand)
 
 Bug Fixes
----------------------
-(No changes)
+
+* LUCENE-9001: Fix race condition in SetOnce. (Przemko Robakowski)
 
 Other
 ---------------------

lucene/core/src/java/org/apache/lucene/util/SetOnce.java

@@ -16,7 +16,7 @@
  */
 package org.apache.lucene.util;
 
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 
 /**
@@ -36,16 +36,24 @@ public final class SetOnce<T> implements Cloneable {
       super("The object cannot be set twice!");
     }
   }
-  
-  private volatile T obj = null;
-  private final AtomicBoolean set;
+
+  /** Holding object and marking that it was already set */
+  private static final class Wrapper<T> {
+    private T object;
+
+    private Wrapper(T object) {
+      this.object = object;
+    }
+  }
+
+  private final AtomicReference<Wrapper<T>> set;
   
   /**
    * A default constructor which does not set the internal object, and allows
    * setting it by calling {@link #set(Object)}.
    */
   public SetOnce() {
-    set = new AtomicBoolean(false);
+    set = new AtomicReference<>();
   }
 
   /**
@@ -57,21 +65,27 @@ public final class SetOnce<T> implements Cloneable {
    * @see #set(Object)
    */
   public SetOnce(T obj) {
-    this.obj = obj;
-    set = new AtomicBoolean(true);
+    set = new AtomicReference<>(new Wrapper<>(obj));
   }
   
   /** Sets the given object. If the object has already been set, an exception is thrown. */
   public final void set(T obj) {
-    if (set.compareAndSet(false, true)) {
-      this.obj = obj;
-    } else {
+    if (!trySet(obj)) {
       throw new AlreadySetException();
     }
   }
+
+  /** Sets the given object if none was set before.
+   *
+   * @return true if object was set successfully, false otherwise
+   * */
+  public final boolean trySet(T obj) {
+    return set.compareAndSet(null, new Wrapper<>(obj));
+  }
   
   /** Returns the object set by {@link #set(Object)}. */
   public final T get() {
-    return obj;
+    Wrapper<T> wrapper = set.get();
+    return wrapper == null ? null : wrapper.object;
   }
 }

lucene/core/src/test/org/apache/lucene/util/TestSetOnce.java

@@ -69,6 +69,15 @@ public class TestSetOnce extends LuceneTestCase {
     assertEquals(5, set.get().intValue());
     set.set(7);
   }
+
+  @Test
+  public void testTrySet() {
+    SetOnce<Integer> set = new SetOnce<>();
+    assertTrue(set.trySet(5));
+    assertEquals(5, set.get().intValue());
+    assertFalse(set.trySet(7));
+    assertEquals(5, set.get().intValue());
+  }
   
   @Test
   public void testSetMultiThreaded() throws Exception {