From 5e41e675ebc18dbfa41fd5c779fe2675075487e8 Mon Sep 17 00:00:00 2001 From: MohammedKHC Date: Tue, 10 Feb 2026 03:04:27 +0200 Subject: [PATCH] Implement SegmentPool for all platforms. Closes: #1629 --- .../commonMain/kotlin/okio/CommonPlatform.kt | 3 + .../src/commonMain/kotlin/okio/SegmentPool.kt | 113 ++++++++++++++- okio/src/jsMain/kotlin/okio/JsPlatform.kt | 17 ++- okio/src/jvmMain/kotlin/okio/-JvmPlatform.kt | 3 + okio/src/jvmMain/kotlin/okio/SegmentPool.kt | 129 ------------------ .../kotlin/okio/NativePlatform.kt} | 19 +-- okio/src/wasmMain/kotlin/okio/WasmPlatform.kt | 6 + 7 files changed, 147 insertions(+), 143 deletions(-) delete mode 100644 okio/src/jvmMain/kotlin/okio/SegmentPool.kt rename okio/src/{nonJvmMain/kotlin/okio/SegmentPool.kt => nativeMain/kotlin/okio/NativePlatform.kt} (56%) diff --git a/okio/src/commonMain/kotlin/okio/CommonPlatform.kt b/okio/src/commonMain/kotlin/okio/CommonPlatform.kt index e073d0b342..010a49183c 100644 --- a/okio/src/commonMain/kotlin/okio/CommonPlatform.kt +++ b/okio/src/commonMain/kotlin/okio/CommonPlatform.kt @@ -32,6 +32,9 @@ expect inline fun Lock.withLock(action: () -> T): T internal expect fun newLock(): Lock +internal expect fun getCurrentThreadId(): Long +internal expect fun getAvailableProcessors(): Int + expect open class IOException(message: String?, cause: Throwable?) : Exception { constructor(message: String?) constructor() diff --git a/okio/src/commonMain/kotlin/okio/SegmentPool.kt b/okio/src/commonMain/kotlin/okio/SegmentPool.kt index f21c7a29fb..ab0a3986b2 100644 --- a/okio/src/commonMain/kotlin/okio/SegmentPool.kt +++ b/okio/src/commonMain/kotlin/okio/SegmentPool.kt @@ -15,22 +15,127 @@ */ package okio +import kotlin.concurrent.atomics.AtomicReference +import kotlin.concurrent.atomics.ExperimentalAtomicApi +import kotlin.jvm.JvmStatic +import okio.SegmentPool.LOCK +import okio.SegmentPool.recycle +import okio.SegmentPool.take + /** * A collection of unused segments, necessary to avoid GC churn and zero-fill. * This pool is a thread-safe static singleton. + * + * This class pools segments in a lock-free singly-linked stack. Though this code is lock-free it + * does use a sentinel [LOCK] value to defend against races. Conflicted operations are not retried, + * so there is no chance of blocking despite the term "lock". + * + * On [take], a caller swaps the stack's next pointer with the [LOCK] sentinel. If the stack was + * not already locked, the caller replaces the head node with its successor. + * + * On [recycle], a caller swaps the stack's next pointer with the [LOCK] sentinel. If the stack was + * not already locked, the caller replaces the head node with a new node whose successor is the + * replaced head. + * + * On conflict, operations succeed, but segments are not pushed into the stack. For example, a + * [take] that loses a race allocates a new segment regardless of the pool size. A [recycle] call + * that loses a race will not increase the size of the pool. Under significant contention, this pool + * will have fewer hits and the VM will do more GC and zero filling of arrays. + * + * This tracks the number of bytes in each linked list in its [Segment.limit] property. Each element + * has a limit that's one segment size greater than its successor element. The maximum size of the + * pool is a product of [MAX_SIZE] and [HASH_BUCKET_COUNT]. */ -internal expect object SegmentPool { - val MAX_SIZE: Int +@OptIn(ExperimentalAtomicApi::class) +internal object SegmentPool { + /** The maximum number of bytes to pool per hash bucket. */ + // TODO: Is 64 KiB a good maximum size? Do we ever have that many idle segments? + const val MAX_SIZE = 64 * 1024 // 64 KiB. + + /** A sentinel segment to indicate that the linked list is currently being modified. */ + private val LOCK = Segment(ByteArray(0), pos = 0, limit = 0, shared = false, owner = false) + + /** + * The number of hash buckets. This number needs to balance keeping the pool small and contention + * low. We use the number of processors rounded up to the nearest power of two. For example a + * machine with 6 cores will have 8 hash buckets. + */ + private val HASH_BUCKET_COUNT = + (getAvailableProcessors() * 2 - 1).takeHighestOneBit() + + /** + * Hash buckets each contain a singly-linked list of segments. The index/key is a hash function of + * thread ID because it may reduce contention or increase locality. + * + * We don't use [ThreadLocal] because we don't know how many threads the host process has and we + * don't want to leak memory for the duration of a thread's life. + */ + private val hashBuckets = Array(HASH_BUCKET_COUNT) { + AtomicReference(null) // null value implies an empty bucket + } /** * For testing only. Returns a snapshot of the number of bytes currently in the pool. If the pool * is segmented such as by thread, this returns the byte count accessible to the calling thread. */ val byteCount: Int + get() { + val first = firstRef().load() ?: return 0 + return first.limit + } /** Return a segment for the caller's use. */ - fun take(): Segment + @JvmStatic + fun take(): Segment { + val firstRef = firstRef() + + val first = firstRef.exchange(LOCK) + when { + first === LOCK -> { + // We didn't acquire the lock. Don't take a pooled segment. + return Segment() + } + first == null -> { + // We acquired the lock but the pool was empty. Unlock and return a new segment. + firstRef.store(null) + return Segment() + } + else -> { + // We acquired the lock and the pool was not empty. Pop the first element and return it. + firstRef.store(first.next) + first.next = null + first.limit = 0 + return first + } + } + } /** Recycle a segment that the caller no longer needs. */ - fun recycle(segment: Segment) + @JvmStatic + fun recycle(segment: Segment) { + require(segment.next == null && segment.prev == null) + if (segment.shared) return // This segment cannot be recycled. + + val firstRef = firstRef() + + val first = firstRef.exchange(LOCK) + if (first === LOCK) return // A take() or recycle() is currently in progress. + val firstLimit = first?.limit ?: 0 + if (firstLimit >= MAX_SIZE) { + firstRef.store(first) // Pool is full. + return + } + + segment.next = first + segment.pos = 0 + segment.limit = firstLimit + Segment.SIZE + + firstRef.store(segment) + } + + private fun firstRef(): AtomicReference { + // Get a value in [0..HASH_BUCKET_COUNT] based on the current thread. + val hashBucket = (getCurrentThreadId() and (HASH_BUCKET_COUNT - 1L)).toInt() + return hashBuckets[hashBucket] + } } diff --git a/okio/src/jsMain/kotlin/okio/JsPlatform.kt b/okio/src/jsMain/kotlin/okio/JsPlatform.kt index 04350d157a..924f07f821 100644 --- a/okio/src/jsMain/kotlin/okio/JsPlatform.kt +++ b/okio/src/jsMain/kotlin/okio/JsPlatform.kt @@ -16,7 +16,7 @@ package okio /* - * This file exposes Node.js `os` and `path` APIs to Kotlin/JS, using reasonable default behaviors + * This file exposes Node.js `os`, `path` and `worker_threads` APIs to Kotlin/JS, using reasonable default behaviors * if those symbols aren't available. * * This was originally implemented using a Kotlin/JS [JsModule], but that broke browser builds @@ -50,3 +50,18 @@ private val path: dynamic internal val tmpdir: String get() = os?.tmpdir() as? String ?: "/tmp" + +private val worker_threads: dynamic + get() { + return try { + js("require('worker_threads')") + } catch (t: Throwable) { + null + } + } + +internal actual fun getCurrentThreadId() = + (worker_threads?.threadId as? Int)?.toLong() ?: 0L + +internal actual fun getAvailableProcessors() = + os?.availableParallelism() as? Int ?: js("navigator.hardwareConcurrency") as? Int ?: 1 diff --git a/okio/src/jvmMain/kotlin/okio/-JvmPlatform.kt b/okio/src/jvmMain/kotlin/okio/-JvmPlatform.kt index a95b242538..ed6c02dd18 100644 --- a/okio/src/jvmMain/kotlin/okio/-JvmPlatform.kt +++ b/okio/src/jvmMain/kotlin/okio/-JvmPlatform.kt @@ -30,6 +30,9 @@ actual typealias ArrayIndexOutOfBoundsException = java.lang.ArrayIndexOutOfBound actual typealias Lock = ReentrantLock +internal actual fun getCurrentThreadId() = Thread.currentThread().id +internal actual fun getAvailableProcessors() = Runtime.getRuntime().availableProcessors() + internal actual fun newLock(): Lock = ReentrantLock() actual inline fun Lock.withLock(action: () -> T): T { diff --git a/okio/src/jvmMain/kotlin/okio/SegmentPool.kt b/okio/src/jvmMain/kotlin/okio/SegmentPool.kt deleted file mode 100644 index 4a05478c6b..0000000000 --- a/okio/src/jvmMain/kotlin/okio/SegmentPool.kt +++ /dev/null @@ -1,129 +0,0 @@ -/* - * Copyright (C) 2014 Square, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package okio - -import java.util.concurrent.atomic.AtomicReference -import okio.SegmentPool.LOCK -import okio.SegmentPool.recycle -import okio.SegmentPool.take - -/** - * This class pools segments in a lock-free singly-linked stack. Though this code is lock-free it - * does use a sentinel [LOCK] value to defend against races. Conflicted operations are not retried, - * so there is no chance of blocking despite the term "lock". - * - * On [take], a caller swaps the stack's next pointer with the [LOCK] sentinel. If the stack was - * not already locked, the caller replaces the head node with its successor. - * - * On [recycle], a caller swaps the stack's next pointer with the [LOCK] sentinel. If the stack was - * not already locked, the caller replaces the head node with a new node whose successor is the - * replaced head. - * - * On conflict, operations succeed, but segments are not pushed into the stack. For example, a - * [take] that loses a race allocates a new segment regardless of the pool size. A [recycle] call - * that loses a race will not increase the size of the pool. Under significant contention, this pool - * will have fewer hits and the VM will do more GC and zero filling of arrays. - * - * This tracks the number of bytes in each linked list in its [Segment.limit] property. Each element - * has a limit that's one segment size greater than its successor element. The maximum size of the - * pool is a product of [MAX_SIZE] and [HASH_BUCKET_COUNT]. - */ -internal actual object SegmentPool { - /** The maximum number of bytes to pool per hash bucket. */ - // TODO: Is 64 KiB a good maximum size? Do we ever have that many idle segments? - actual val MAX_SIZE = 64 * 1024 // 64 KiB. - - /** A sentinel segment to indicate that the linked list is currently being modified. */ - private val LOCK = Segment(ByteArray(0), pos = 0, limit = 0, shared = false, owner = false) - - /** - * The number of hash buckets. This number needs to balance keeping the pool small and contention - * low. We use the number of processors rounded up to the nearest power of two. For example a - * machine with 6 cores will have 8 hash buckets. - */ - private val HASH_BUCKET_COUNT = - Integer.highestOneBit(Runtime.getRuntime().availableProcessors() * 2 - 1) - - /** - * Hash buckets each contain a singly-linked list of segments. The index/key is a hash function of - * thread ID because it may reduce contention or increase locality. - * - * We don't use [ThreadLocal] because we don't know how many threads the host process has and we - * don't want to leak memory for the duration of a thread's life. - */ - private val hashBuckets: Array> = Array(HASH_BUCKET_COUNT) { - AtomicReference() // null value implies an empty bucket - } - - actual val byteCount: Int - get() { - val first = firstRef().get() ?: return 0 - return first.limit - } - - @JvmStatic - actual fun take(): Segment { - val firstRef = firstRef() - - val first = firstRef.getAndSet(LOCK) - when { - first === LOCK -> { - // We didn't acquire the lock. Don't take a pooled segment. - return Segment() - } - first == null -> { - // We acquired the lock but the pool was empty. Unlock and return a new segment. - firstRef.set(null) - return Segment() - } - else -> { - // We acquired the lock and the pool was not empty. Pop the first element and return it. - firstRef.set(first.next) - first.next = null - first.limit = 0 - return first - } - } - } - - @JvmStatic - actual fun recycle(segment: Segment) { - require(segment.next == null && segment.prev == null) - if (segment.shared) return // This segment cannot be recycled. - - val firstRef = firstRef() - - val first = firstRef.getAndSet(LOCK) - if (first === LOCK) return // A take() or recycle() is currently in progress. - val firstLimit = first?.limit ?: 0 - if (firstLimit >= MAX_SIZE) { - firstRef.set(first) // Pool is full. - return - } - - segment.next = first - segment.pos = 0 - segment.limit = firstLimit + Segment.SIZE - - firstRef.set(segment) - } - - private fun firstRef(): AtomicReference { - // Get a value in [0..HASH_BUCKET_COUNT) based on the current thread. - val hashBucket = (Thread.currentThread().id and (HASH_BUCKET_COUNT - 1L)).toInt() - return hashBuckets[hashBucket] - } -} diff --git a/okio/src/nonJvmMain/kotlin/okio/SegmentPool.kt b/okio/src/nativeMain/kotlin/okio/NativePlatform.kt similarity index 56% rename from okio/src/nonJvmMain/kotlin/okio/SegmentPool.kt rename to okio/src/nativeMain/kotlin/okio/NativePlatform.kt index 50e709913b..8ebf98624b 100644 --- a/okio/src/nonJvmMain/kotlin/okio/SegmentPool.kt +++ b/okio/src/nativeMain/kotlin/okio/NativePlatform.kt @@ -1,5 +1,5 @@ /* - * Copyright (C) 2014 Square, Inc. + * Copyright (C) 2026 Square, Inc. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -15,13 +15,14 @@ */ package okio -internal actual object SegmentPool { - actual val MAX_SIZE: Int = 0 +import kotlin.experimental.ExperimentalNativeApi +import kotlin.native.concurrent.ObsoleteWorkersApi +import kotlin.native.concurrent.Worker - actual val byteCount: Int = 0 +@OptIn(ObsoleteWorkersApi::class, ExperimentalStdlibApi::class) +internal actual fun getCurrentThreadId() = + Worker.current.platformThreadId.toLong() - actual fun take(): Segment = Segment() - - actual fun recycle(segment: Segment) { - } -} +@OptIn(ExperimentalNativeApi::class) +internal actual fun getAvailableProcessors() = + Platform.getAvailableProcessors() diff --git a/okio/src/wasmMain/kotlin/okio/WasmPlatform.kt b/okio/src/wasmMain/kotlin/okio/WasmPlatform.kt index 4a874ec826..5a8d094448 100644 --- a/okio/src/wasmMain/kotlin/okio/WasmPlatform.kt +++ b/okio/src/wasmMain/kotlin/okio/WasmPlatform.kt @@ -17,3 +17,9 @@ package okio internal actual val PLATFORM_DIRECTORY_SEPARATOR: String get() = "/" + +// FIXME Should wasmJs be shared with js? +// Maybe create a new webMain source set? +// And what about wasmWasi? I don't think that wasi provides any APIS to get these.. +internal actual fun getCurrentThreadId() = 0L +internal actual fun getAvailableProcessors() = 1