From dad15ba15e00df4fef024dae3b101b7969775ed6 Mon Sep 17 00:00:00 2001
From: Przemko Robakowski <przemko.robakowski@elastic.co>
Date: Mon, 14 Oct 2019 18:33:46 +0200
Subject: [PATCH] LUCENE-9001: Fix race condition in SetOnce (#931)

---
 lucene/CHANGES.txt                            |  4 +--
 .../java/org/apache/lucene/util/SetOnce.java  | 36 +++++++++++++------
 .../org/apache/lucene/util/TestSetOnce.java   |  9 +++++
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 0a986abade2..4315a7d050e 100644
--- a/lucene/CHANGES.txt
+++ b/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
 ---------------------
diff --git a/lucene/core/src/java/org/apache/lucene/util/SetOnce.java b/lucene/core/src/java/org/apache/lucene/util/SetOnce.java
index 9be88ecfe4d..3c3f277ace1 100644
--- a/lucene/core/src/java/org/apache/lucene/util/SetOnce.java
+++ b/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;
   }
 }
diff --git a/lucene/core/src/test/org/apache/lucene/util/TestSetOnce.java b/lucene/core/src/test/org/apache/lucene/util/TestSetOnce.java
index 0d93b882d66..ae123214d68 100644
--- a/lucene/core/src/test/org/apache/lucene/util/TestSetOnce.java
+++ b/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 {