From cff1d7f1f0a98d736e42520305354650ff4ff652 Mon Sep 17 00:00:00 2001
From: Sergey Makarov <sergeytop100@gmail.com>
Date: Fri, 6 Oct 2023 13:25:18 +0300
Subject: [PATCH] Make safeWithTimeoutContextFlagUnset test more explicit & fix
 it on CI (previously on CI executionThread1 != executionThread2)

---
 .../fingerprint/tools/threading/safe/Safe.kt  | 23 +++++++++---
 .../android/fingerprint/SafeTests.kt          | 35 +++++++++++--------
 2 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/fingerprint/src/main/java/com/fingerprintjs/android/fingerprint/tools/threading/safe/Safe.kt b/fingerprint/src/main/java/com/fingerprintjs/android/fingerprint/tools/threading/safe/Safe.kt
index e9cfe61..36250a2 100644
--- a/fingerprint/src/main/java/com/fingerprintjs/android/fingerprint/tools/threading/safe/Safe.kt
+++ b/fingerprint/src/main/java/com/fingerprintjs/android/fingerprint/tools/threading/safe/Safe.kt
@@ -8,13 +8,26 @@ import java.util.concurrent.Callable
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.TimeoutException
 import java.util.concurrent.atomic.AtomicReference
-import kotlin.concurrent.getOrSet
 
 internal object Safe {
     const val timeoutShort = 1_000L
 
+    private val runningInsideSafeWithTimeout = ThreadLocal<Boolean>()
+
+    @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
+    internal fun markInsideSafeWithTimeout() {
+        runningInsideSafeWithTimeout.set(true)
+    }
+
     @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
-    internal val runningInsideSafeWithTimeout = ThreadLocal<Boolean>()
+    internal fun clearInsideSafeWithTimeout() {
+        runningInsideSafeWithTimeout.remove()
+    }
+
+    @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
+    internal fun getInsideSafeWithTimeout(): Boolean {
+        return runningInsideSafeWithTimeout.get() ?: false
+    }
 
     @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
     internal fun logIllegalSafeWithTimeoutUsage() {
@@ -35,19 +48,19 @@ internal fun <T> safeWithTimeout(
     // we can't make a local variable volatile, hence using atomic reference here
     val executionThread = AtomicReference<Thread>(null)
 
-    if (Safe.runningInsideSafeWithTimeout.getOrSet { false }) {
+    if (Safe.getInsideSafeWithTimeout()) {
         Safe.logIllegalSafeWithTimeoutUsage()
     }
 
     val future = runCatching {
         sharedExecutor.submit(
             Callable {
-                Safe.runningInsideSafeWithTimeout.set(true)
+                Safe.markInsideSafeWithTimeout()
                 executionThread.set(Thread.currentThread())
                 try {
                     block()
                 } finally {
-                    Safe.runningInsideSafeWithTimeout.remove()
+                    Safe.clearInsideSafeWithTimeout()
                 }
             }
         )!!
diff --git a/fingerprint/src/test/java/com/fingerprintjs/android/fingerprint/SafeTests.kt b/fingerprint/src/test/java/com/fingerprintjs/android/fingerprint/SafeTests.kt
index 95654cd..0d1b564 100644
--- a/fingerprint/src/test/java/com/fingerprintjs/android/fingerprint/SafeTests.kt
+++ b/fingerprint/src/test/java/com/fingerprintjs/android/fingerprint/SafeTests.kt
@@ -9,13 +9,13 @@ import com.fingerprintjs.android.fingerprint.tools.threading.sharedExecutor
 import io.mockk.every
 import io.mockk.mockkObject
 import io.mockk.unmockkObject
+import io.mockk.verify
+import io.mockk.verifyOrder
 import junit.framework.TestCase
 import org.junit.After
 import org.junit.Test
-import java.util.concurrent.Callable
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.ExecutionException
-import kotlin.concurrent.getOrSet
 
 class SafeTests {
 
@@ -156,27 +156,32 @@ class SafeTests {
         safeWithTimeoutContextFlagUnset(whenBlockThrows = true)
 
     private fun safeWithTimeoutContextFlagUnset(whenBlockThrows: Boolean) {
-        var executionThread1: Long? = null
-        var executionThread2: Long? = null
+        mockkObject(Safe)
+        var clearThreadId: Long? = null
+        every { Safe.clearInsideSafeWithTimeout() } answers {
+            callOriginal().also { clearThreadId = Thread.currentThread().id }
+        }
+        var markThreadId: Long? = null
+        every { Safe.markInsideSafeWithTimeout() } answers {
+            callOriginal().also { markThreadId = Thread.currentThread().id }
+        }
 
         safeWithTimeout {
-            executionThread1 = Thread.currentThread().id
             if (whenBlockThrows)
                 throw Exception()
         }
 
-        val countDownLatch = CountDownLatch(1)
-        var insideSafe: Boolean? = null
-        sharedExecutor.submit(Callable {
-            executionThread2 = Thread.currentThread().id
-            insideSafe = Safe.runningInsideSafeWithTimeout.getOrSet { false }
-            countDownLatch.countDown()
-        })
-        countDownLatch.await()
+        verify(exactly = 1) {
+            Safe.markInsideSafeWithTimeout()
+            Safe.clearInsideSafeWithTimeout()
+        }
+        verifyOrder {
+            Safe.markInsideSafeWithTimeout()
+            Safe.clearInsideSafeWithTimeout()
+        }
 
+        TestCase.assertEquals(markThreadId, clearThreadId)
         unmockkObject(Safe)
-        TestCase.assertEquals(executionThread1, executionThread2)
-        TestCase.assertEquals(false, insideSafe)
     }
 
     @Test