[fleet, util] replace atomic shims with stdlib functions

GitOrigin-RevId: aec14ee003b0b608f801feff7eb1d66e2abfa79d
This commit is contained in:
Alexander Zolotov
2025-10-07 15:28:33 +00:00
committed by intellij-monorepo-bot
parent 1c604ef58a
commit dbdff0785d
7 changed files with 21 additions and 64 deletions
@@ -4,12 +4,12 @@ package com.jetbrains.rhizomedb
import fleet.fastutil.longs.LongArrayList
import fleet.fastutil.longs.toArray
import fleet.util.computeShim
import fleet.util.updateAndGet
import kotlinx.collections.immutable.PersistentMap
import kotlinx.collections.immutable.PersistentSet
import kotlinx.collections.immutable.persistentHashMapOf
import kotlinx.collections.immutable.persistentHashSetOf
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.updateAndFetch
class QueryCache private constructor(private val cache: AtomicReference<QueryCacheData>) {
@@ -66,14 +66,14 @@ class QueryCache private constructor(private val cache: AtomicReference<QueryCac
companion object {
fun empty(): QueryCache =
QueryCache(kotlin.concurrent.atomics.AtomicReference(QueryCacheData(persistentHashMapOf(), persistentHashMapOf())))
QueryCache(AtomicReference(QueryCacheData(persistentHashMapOf(), persistentHashMapOf())))
}
fun <T> performQuery(query: CachedQuery<T>, compute: () -> CachedQueryResult<T>): CachedQueryResult<T> =
@Suppress("UNCHECKED_CAST")
(cache.load().find(query) as CachedQueryResult<T>?) ?: run {
val res = compute()
cache.updateAndGet { index ->
cache.updateAndFetch { index ->
val existing = index.find(query)
when {
existing == null -> index.insert(query, res)
@@ -3,9 +3,7 @@ package fleet.rpc.core
import fleet.util.UID
import fleet.util.async.coroutineNameAppended
import fleet.util.getAndUpdate
import fleet.util.logging.logger
import fleet.util.updateAndGet
import kotlinx.coroutines.*
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ReceiveChannel
@@ -13,6 +11,8 @@ import kotlinx.coroutines.channels.SendChannel
import kotlinx.coroutines.channels.consumeEach
import kotlinx.serialization.KSerializer
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.fetchAndUpdate
import kotlin.concurrent.atomics.updateAndFetch
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.resumeWithException
import kotlin.math.max
@@ -54,12 +54,12 @@ class Budget(initial: Int) {
private data class State(val cancellation: CancellationException?, val budget: Int, val continuation: CancellableContinuation<Unit>?)
private fun withdraw(): Boolean {
return state.getAndUpdate { it.copy(budget = max(0, it.budget - 1)) }.budget > 0
return state.fetchAndUpdate { it.copy(budget = max(0, it.budget - 1)) }.budget > 0
}
private suspend fun await() {
suspendCancellableCoroutine { continuation ->
val r = state.updateAndGet {
val r = state.updateAndFetch {
require(it.continuation == null) { "Budget is not intended to use by several producers" }
when {
it.budget > 0 -> it
@@ -82,14 +82,14 @@ class Budget(initial: Int) {
fun refill(quantity: Int) {
require(quantity > 0)
val r = state.getAndUpdate {
val r = state.fetchAndUpdate {
it.copy(budget = it.budget + quantity, continuation = null)
}
r.continuation?.resumeWith(Result.success(Unit))
}
internal fun cancel(cause: CancellationException) {
val was = state.getAndUpdate {
val was = state.fetchAndUpdate {
if (it.cancellation == null) {
it.copy(cancellation = cause, continuation = null)
}
@@ -1,9 +0,0 @@
// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package fleet.util
import fleet.util.multiplatform.linkToActual
import kotlin.concurrent.atomics.AtomicReference
fun <T> AtomicReference<T>.updateAndGet(f: (T) -> T): T = linkToActual()
fun <T> AtomicReference<T>.getAndUpdate(f: (T) -> T): T = linkToActual()
@@ -1,10 +1,10 @@
// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package fleet.util.async
import fleet.util.updateAndGet
import kotlinx.collections.immutable.persistentSetOf
import kotlinx.coroutines.*
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.updateAndFetch
import kotlin.coroutines.CoroutineContext
data class Handle<T>(val value: Deferred<ValueWithContext<T>>,
@@ -86,8 +86,12 @@ private suspend fun<T> handleScopeImpl(outerScope: CoroutineScope, body: suspend
override val coroutineContext: CoroutineContext get() = context
override fun <T> handle(launcher: Launcher<T>): Handle<T> {
val handle = outerScope.handle(launcher)
handles.updateAndGet { hs -> hs.add(handle) }
handle.job.invokeOnCompletion { handles.updateAndGet { hs -> hs.remove(handle) } }
handles.updateAndFetch { hs -> hs.add(handle) }
handle.job.invokeOnCompletion {
handles.updateAndFetch { hs ->
hs.remove(handle)
}
}
return handle
}
}.body()
@@ -1,9 +1,9 @@
// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package fleet.util.openmap
import fleet.util.updateAndGet
import kotlinx.collections.immutable.PersistentMap
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.updateAndFetch
internal class MutableBoundedOpenMapImpl<Domain, V : Any>(val map: AtomicReference<PersistentMap<Key<out V, in Domain>, V>>) : MutableBoundedOpenMap<Domain, V> {
override fun <T : Any> get(k: Key<T, in Domain>): T? {
@@ -13,11 +13,11 @@ internal class MutableBoundedOpenMapImpl<Domain, V : Any>(val map: AtomicReferen
override fun isEmpty() = map.load().size == 0
override fun <T : V> set(k: Key<T, Domain>, v: T) {
map.updateAndGet { map -> map.put(k, v) }
map.updateAndFetch { map -> map.put(k, v) }
}
override fun remove(k: Key<out V, Domain>) {
map.updateAndGet { map -> map.remove(k) }
map.updateAndFetch { map -> map.remove(k) }
}
override fun <T : V> assoc(k: Key<T, Domain>, v: T): BoundedOpenMap<Domain, V> {
@@ -33,14 +33,14 @@ internal class MutableBoundedOpenMapImpl<Domain, V : Any>(val map: AtomicReferen
}
override fun <T : V> update(key: Key<T, Domain>, f: (T?) -> T): T {
return map.updateAndGet { map ->
return map.updateAndFetch { map ->
val v1 = f(map.get(key) as T?)
map.put(key, v1)
}.get(key) as T
}
override fun <T : V> getOrInit(key: Key<T, Domain>, init: () -> T): T {
return map.updateAndGet { map ->
return map.updateAndFetch { map ->
val v = map.get(key)
if (v == null) {
map.put(key, init())
@@ -1,16 +0,0 @@
// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package fleet.util
import fleet.util.multiplatform.Actual
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.asJavaAtomic
@Actual
fun <T> AtomicReference<T>.updateAndGetJvm(f: (T) -> T): T {
return asJavaAtomic().updateAndGet(f)
}
@Actual
fun <T> AtomicReference<T>.getAndUpdateJvm(f: (T) -> T): T {
return asJavaAtomic().getAndUpdate(f)
}
@@ -1,22 +0,0 @@
// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
@file:OptIn(ExperimentalAtomicApi::class)
package fleet.util
import fleet.util.multiplatform.Actual
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.ExperimentalAtomicApi
@Actual
fun <T> AtomicReference<T>.updateAndGetWasmJs(f: (T) -> T): T {
val new = f(load())
store(new)
return new
}
@Actual
fun <T> AtomicReference<T>.getAndUpdateWasmJs(f: (T) -> T): T {
val old = load()
store(f(old))
return old
}