simplify — use getAndSet

platform/platform-tests/testSrc/org/jetbrains/concurrency/AsyncPromiseTest.kt

@@ -35,6 +35,31 @@ class AsyncPromiseTest {
     doHandlerTest(true)
   }
 
+  @Test
+  fun state() {
+    val promise = AsyncPromise<String>()
+    val count = AtomicInteger()
+
+    val r = {
+      promise.done { count.incrementAndGet() }
+    }
+
+    val s = {
+      promise.setResult("test")
+    }
+
+    val numThreads = 30
+    assertConcurrent(*Array(numThreads, {
+      if (it and 1 === 0) r else s
+    }))
+
+    assertThat(count.get()).isEqualTo(numThreads / 2)
+    assertThat(promise.get()).isEqualTo("test")
+
+    r()
+    assertThat(count.get()).isEqualTo((numThreads / 2) + 1)
+  }
+
   fun doHandlerTest(reject: Boolean) {
     val promise = AsyncPromise<String>()
     val count = AtomicInteger()
@@ -68,7 +93,7 @@ class AsyncPromiseTest {
   }
 }
 
-fun assertConcurrent(vararg runnables: () -> Any, maxTimeoutSeconds: Int = 5) {
+fun assertConcurrent(vararg runnables: () -> Any?, maxTimeoutSeconds: Int = 5) {
   val numThreads = runnables.size
   val exceptions = ContainerUtil.createLockFreeCopyOnWriteList<Throwable>()
   val threadPool = Executors.newFixedThreadPool(numThreads)

platform/projectModel-api/src/org/jetbrains/concurrency/AsyncPromise.kt

@@ -171,7 +171,7 @@ open class AsyncPromise<T> : Promise<T>, Getter<T> {
 
     this.result = result
 
-    val done = getAndClearHandler(doneRef)
+    val done = doneRef.getAndSet(null)
     rejectedRef.set(null)
 
     if (done != null && !isObsolete(done)) {
@@ -192,7 +192,7 @@ open class AsyncPromise<T> : Promise<T>, Getter<T> {
 
     result = error
 
-    val rejected = getAndClearHandler(rejectedRef)
+    val rejected = rejectedRef.getAndSet(null)
     doneRef.set(null)
 
     if (rejected == null) {
@@ -204,15 +204,6 @@ open class AsyncPromise<T> : Promise<T>, Getter<T> {
     return true
   }
 
-  private fun <T> getAndClearHandler(ref: AtomicReference<Consumer<in T>?>): Consumer<in T>? {
-    var handler: Consumer<in T>?
-    do {
-      handler = ref.get()
-    }
-    while (!ref.compareAndSet(handler, null))
-    return handler
-  }
-
   override fun processed(processed: Consumer<in T>): Promise<T> {
     done(processed)
     rejected { processed.consume(null) }
@@ -271,7 +262,7 @@ open class AsyncPromise<T> : Promise<T>, Getter<T> {
     }
 
     if (state == targetState) {
-      getAndClearHandler(ref)?.let {
+      ref.getAndSet(null)?.let {
         @Suppress("UNCHECKED_CAST")
         it.consume(result as T?)
       }