Repository: Kotlin/kotlinx-io
Pull request #352 (2024-07-04 13:55:28)

SegmentPool improvements
Currently, SegmentPool implementation on JVM works more like a cache than a pool: it has a relatively small fixed size, segments are taken from the pool on an effort basis, and it's easy to make a segment non-recyclable (meaning that it will never be returned to the pool).

This behavior is totally fine for many applications, but it's not for Ktor.
This PR addresses several issues:
- Support precise segment reference tracking to return shared segments back to the pool eventually. Currently, once a segment is shared, it is never returned to the pool because the segment state is tracked by a flag that could be set but could not be used. To solve that problem, on JVM, the exact number of shared copies is now tracked, and the value is decremented on attempts to recycle shared copies. Once the last copy is passed to the `Segment.recycle`, it'll be returned to the pool.
- Retry lost CAS attempts when taking or returning segments from the pool. The previous sentence is quite descriptive; I'll only add that it does not make much sense without the next change.
- Support segment pool size configuration. Currently, the segment pool consists of multiple buckets (the exact number depends on the system's logical CPU count), each with a max size of 64KiB. This PR adds a second level/tier to the pool with half as many buckets as an old first tier. The total size of the new tier is configurable via system properties and is `0` by default. All failed attempts to take or recycle segments in the first tier use the second tier as a fallback. Unlike the first tier, if the second tier's bucket is empty, `take` will scan all other buckets. The scan is also performed on `recycle` to a full bucket.

All these changes reduce the allocation count by combining more precise work with the pool and optionally extensible pool size.
Two-level design helps with avoiding performance penalties induced by the scan technique employed by the second tier when that tier is not used and, at the same time, allows using the same sharding technique as in the first tier enhanced with the bucket steeling techniques.

fzhinkin: CC @e5l @bjhham


Patch:
diff --git a/core/api/kotlinx-io-core.api b/core/api/kotlinx-io-core.api
index 60f0e5af..2662ef02 100644
--- a/core/api/kotlinx-io-core.api
+++ b/core/api/kotlinx-io-core.api
@@ -105,7 +105,7 @@ public abstract interface class kotlinx/io/RawSource : java/lang/AutoCloseable {
 }
 
 public final class kotlinx/io/Segment {
-	public synthetic fun <init> ([BIIZZLkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public synthetic fun <init> ([BIILkotlinx/io/SegmentCopyTracker;ZLkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final synthetic fun dataAsByteArray (Z)[B
 	public final synthetic fun getLimit ()I
 	public final synthetic fun getNext ()Lkotlinx/io/Segment;
diff --git a/core/common/src/Buffer.kt b/core/common/src/Buffer.kt
index c7e8153a..abd2eb14 100644
--- a/core/common/src/Buffer.kt
+++ b/core/common/src/Buffer.kt
@@ -488,7 +488,7 @@ public class Buffer : Source, Sink {
 
         while (remainingByteCount > 0L) {
             // Is a prefix of the source's head segment all that we need to move?
-            if (remainingByteCount < source.head!!.limit - source.head!!.pos) {
+            if (remainingByteCount < source.head!!.size) {
                 val tail = tail
                 if (tail != null && tail.owner &&
                     remainingByteCount + tail.limit - (if (tail.shared) 0 else tail.pos) <= Segment.SIZE
@@ -501,17 +501,13 @@ public class Buffer : Source, Sink {
                 } else {
                     // We're going to need another segment. Split the source's head
                     // segment in two, then move the first of those two to this buffer.
-                    val newHead = source.head!!.split(remainingByteCount.toInt())
-                    if (source.head == source.tail) {
-                        source.tail = newHead
-                    }
-                    source.head = newHead
+                    source.head = source.head!!.split(remainingByteCount.toInt())
                 }
             }
 
             // Remove the source's head segment and append it to our tail.
-            val segmentToMove = source.head
-            val movedByteCount = (segmentToMove!!.limit - segmentToMove.pos).toLong()
+            val segmentToMove = source.head!!
+            val movedByteCount = segmentToMove.size.toLong()
             source.head = segmentToMove.pop()
             if (source.head == null) {
                 source.tail = null
diff --git a/core/common/src/Segment.kt b/core/common/src/Segment.kt
index 79055d98..9d6eb06d 100644
--- a/core/common/src/Segment.kt
+++ b/core/common/src/Segment.kt
@@ -23,6 +23,47 @@ package kotlinx.io
 import kotlin.jvm.JvmField
 import kotlin.jvm.JvmSynthetic
 
+/**
+ * Tracks shared segment copies.
+ *
+ * A new [SegmentCopyTracker] instance should be not shared by default (i.e. `shared == false`).
+ * Any further [addCopy] calls should move the tracker to a shared state (i.e. `shared == true`).
+ * Once a shared segment copy is recycled, [removeCopy] should be called.
+ * Depending on implementation, calling [removeCopy] the same number of times as [addCopy] may
+ * or may not transition the tracked back to unshared stated.
+ *
+ * The class is not intended for public use and currently designed to fit the only use case - within JVM SegmentPool
+ * implementation.
+ */
+internal abstract class SegmentCopyTracker {
+    /**
+     * `true` if a tracker shared by multiple segment copies.
+     */
+    abstract val shared: Boolean
+
+    /**
+     * Track a new copy created by sharing an associated segment.
+     */
+    abstract fun addCopy()
+
+    /**
+     * Records reclamation of a shared segment copy associated with this tracker.
+     * If a tracker was in unshared state, this call should not affect an internal state.
+     *
+     * @return `true` if the segment was not shared *before* this called.
+     */
+    abstract fun removeCopy(): Boolean
+}
+
+/**
+ * Simple [SegmentCopyTracker] that always reports shared state.
+ */
+internal object AlwaysSharedCopyTracker : SegmentCopyTracker() {
+    override val shared: Boolean = true
+    override fun addCopy() = Unit
+    override fun removeCopy(): Boolean = true
+}
+
 /**
  * A segment of a buffer.
  *
@@ -59,8 +100,17 @@ public class Segment {
     internal var limit: Int = 0
 
     /** True if other segments or byte strings use the same byte array. */
-    @JvmField
-    internal var shared: Boolean = false
+    internal val shared: Boolean
+        get() = copyTracker?.shared ?: false
+
+    /**
+     * Tracks number shared copies
+     *
+     * Note that this reference is not `@Volatile` as segments are not thread-safe and it's an error
+     * to modify the same segment concurrently.
+     * At the same time, an object [copyTracker] refers to could be modified concurrently.
+     */
+    internal var copyTracker: SegmentCopyTracker? = null
 
     /** True if this segment owns the byte array and can append to it, extending `limit`. */
     @JvmField
@@ -81,14 +131,14 @@ public class Segment {
     private constructor() {
         this.data = ByteArray(SIZE)
         this.owner = true
-        this.shared = false
+        this.copyTracker = null
     }
 
-    private constructor(data: ByteArray, pos: Int, limit: Int, shared: Boolean, owner: Boolean) {
+    private constructor(data: ByteArray, pos: Int, limit: Int, shareToken: SegmentCopyTracker?, owner: Boolean) {
         this.data = data
         this.pos = pos
         this.limit = limit
-        this.shared = shared
+        this.copyTracker = shareToken
         this.owner = owner
     }
 
@@ -98,8 +148,10 @@ public class Segment {
      * prevents it from being pooled.
      */
     internal fun sharedCopy(): Segment {
-        shared = true
-        return Segment(data, pos, limit, true, false)
+        val t = copyTracker ?: SegmentPool.tracker().also {
+            copyTracker = it
+        }
+        return Segment(data, pos, limit, t.also { it.addCopy() }, false)
     }
 
     /**
@@ -284,8 +336,13 @@ public class Segment {
         internal fun new(): Segment = Segment()
 
         @JvmSynthetic
-        internal fun new(data: ByteArray, pos: Int, limit: Int, shared: Boolean, owner: Boolean): Segment
-            = Segment(data, pos, limit, shared, owner)
+        internal fun new(
+            data: ByteArray,
+            pos: Int,
+            limit: Int,
+            copyTracker: SegmentCopyTracker?,
+            owner: Boolean
+        ): Segment = Segment(data, pos, limit, copyTracker, owner)
     }
 }
 
diff --git a/core/common/src/SegmentPool.kt b/core/common/src/SegmentPool.kt
index 070cafbd..3bc511d6 100644
--- a/core/common/src/SegmentPool.kt
+++ b/core/common/src/SegmentPool.kt
@@ -38,4 +38,11 @@ internal expect object SegmentPool {
 
     /** Recycle a segment that the caller no longer needs. */
     fun recycle(segment: Segment)
+
+    /**
+     * Allocates a new copy tracker that'll be associated with a segment from this pool.
+     * For performance reasons, there's no tracker attached to a segment initially.
+     * Instead, it's allocated lazily on the first sharing attempt.
+     */
+    fun tracker(): SegmentCopyTracker
 }
diff --git a/core/common/src/unsafe/UnsafeBufferOperations.kt b/core/common/src/unsafe/UnsafeBufferOperations.kt
index a31f9ca1..2415714d 100644
--- a/core/common/src/unsafe/UnsafeBufferOperations.kt
+++ b/core/common/src/unsafe/UnsafeBufferOperations.kt
@@ -38,7 +38,8 @@ public object UnsafeBufferOperations {
     public fun moveToTail(buffer: Buffer, bytes: ByteArray, startIndex: Int = 0, endIndex: Int = bytes.size) {
         checkBounds(bytes.size, startIndex, endIndex)
         val segment = Segment.new(
-            bytes, startIndex, endIndex, shared = true /* to prevent recycling */,
+            bytes, startIndex, endIndex,
+            AlwaysSharedCopyTracker, /* to prevent recycling */
             owner = false /* can't append to it */
         )
         val tail = buffer.tail
diff --git a/core/js/src/SegmentPool.kt b/core/js/src/SegmentPool.kt
index 615a1497..c61fbeff 100644
--- a/core/js/src/SegmentPool.kt
+++ b/core/js/src/SegmentPool.kt
@@ -14,4 +14,6 @@ internal actual object SegmentPool {
 
     actual fun recycle(segment: Segment) {
     }
+
+    actual fun tracker(): SegmentCopyTracker = AlwaysSharedCopyTracker
 }
diff --git a/core/jvm/src/SegmentPool.kt b/core/jvm/src/SegmentPool.kt
index f70d22da..4a0614e6 100644
--- a/core/jvm/src/SegmentPool.kt
+++ b/core/jvm/src/SegmentPool.kt
@@ -25,26 +25,78 @@ import kotlinx.io.SegmentPool.LOCK
 import kotlinx.io.SegmentPool.MAX_SIZE
 import kotlinx.io.SegmentPool.recycle
 import kotlinx.io.SegmentPool.take
-import java.util.concurrent.atomic.AtomicReference
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater
+import java.util.concurrent.atomic.AtomicReferenceArray
+
+/**
+ * Precise [SegmentCopyTracker] implementation tracking a number of shared segment copies.
+ * Every [addCopy] call increments the counter, every [removeCopy] decrements it.
+ *
+ * After calling [removeCopy] the same number of time [addCopy] was called, tracker returns to the unshared state.
+ *
+ * The class is internal for testing only.
+ */
+internal class RefCountingCopyTracker : SegmentCopyTracker() {
+    companion object {
+        @JvmStatic
+        private val fieldUpdater = AtomicIntegerFieldUpdater.newUpdater(RefCountingCopyTracker::class.java, "copyCount")
+    }
+
+    @Volatile
+    private var copyCount: Int = 0
+
+    override val shared: Boolean
+        get() {
+            return copyCount > 0
+        }
+
+    override fun addCopy() {
+        fieldUpdater.incrementAndGet(this)
+    }
+
+    override fun removeCopy(): Boolean {
+        // The value could not be incremented from `0` under the race,
+        // so once it zero, it remains zero in the scope of this call.
+        if (copyCount == 0) return false
+
+        val updatedValue = fieldUpdater.decrementAndGet(this)
+        // If there are several copies, the last decrement will update copyCount from 0 to -1.
+        // That would be the last standing copy, and we can recycle it.
+        // If, however, the decremented value falls below -1, it's an error as there were more
+        // `removeCopy` than `addCopy` calls.
+        if (updatedValue >= 0) return true
+        check(updatedValue == -1) { "Shared copies count is negative: ${updatedValue + 1}" }
+        copyCount = 0
+        return false
+    }
+}
 
 /**
  * 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".
+ * does use a sentinel [LOCK] value to defend against races. To reduce the contention, the pool consists
+ * of several buckets (see [HASH_BUCKET_COUNT]), each holding a reference to its own segments stack.
+ * Every [take] or [recycle] choose one of the buckets depending on a [Thread.currentThread]'s [Thread.getId].
  *
  * 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 head 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.
+ * On conflict, operations are retried until they succeed.
  *
  * 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].
+ *
+ * [MAX_SIZE] is kept relatively small to avoid excessive memory consumption in case of a large [HASH_BUCKET_COUNT].
+ * For better handling of scenarios with high segments demand, an optional second-level pool could be enabled
+ * by setting up a value of `kotlinx.io.pool.size.bytes` system property.
+ *
+ * The second-level pool use half of the [HASH_BUCKET_COUNT] and if an initially selected bucket if empty on [take] or
+ * full or [recycle], all other buckets will be inspected before finally giving up (which means allocating a new segment
+ * on [take], or loosing a reference to a segment on [recycle]).
+ * That pool is used as a backup in case when [take] or [recycle] failed due to
+ * an empty or exhausted segments chain in a corresponding bucket (one of [HASH_BUCKET_COUNT] buckets).
  */
 internal actual object SegmentPool {
     /** The maximum number of bytes to pool per hash bucket. */
@@ -52,7 +104,8 @@ internal actual object SegmentPool {
     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.new(ByteArray(0), pos = 0, limit = 0, shared = false, owner = false)
+    private val LOCK =
+        Segment.new(ByteArray(0), pos = 0, limit = 0, copyTracker = null, owner = false)
 
     /**
      * The number of hash buckets. This number needs to balance keeping the pool small and contention
@@ -62,6 +115,22 @@ internal actual object SegmentPool {
     private val HASH_BUCKET_COUNT =
         Integer.highestOneBit(Runtime.getRuntime().availableProcessors() * 2 - 1)
 
+    private val HASH_BUCKET_COUNT_L2 = (HASH_BUCKET_COUNT / 2).coerceAtLeast(1)
+
+    // For now, keep things on Android as they were before, but on JVM - use second level cache.
+    // See https://developer.android.com/reference/java/lang/System#getProperties() for property name.
+    private val DEFAULT_SECOND_LEVEL_POOL_TOTAL_SIZE = when (System.getProperty("java.vm.name")) {
+        "Dalvik" -> "0"
+        else -> "4194304" // 4MB
+    }
+
+    private val SECOND_LEVEL_POOL_TOTAL_SIZE =
+        System.getProperty("kotlinx.io.pool.size.bytes", DEFAULT_SECOND_LEVEL_POOL_TOTAL_SIZE)
+            .toIntOrNull()?.coerceAtLeast(0) ?: 0
+
+    private val SECOND_LEVEL_POOL_BUCKET_SIZE =
+        (SECOND_LEVEL_POOL_TOTAL_SIZE / HASH_BUCKET_COUNT).coerceAtLeast(Segment.SIZE)
+
     /**
      * 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.
@@ -69,39 +138,85 @@ internal actual object SegmentPool {
      * 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<AtomicReference<Segment?>> = Array(HASH_BUCKET_COUNT) {
-        AtomicReference<Segment?>() // null value implies an empty bucket
-    }
+    private val hashBuckets: AtomicReferenceArray<Segment?> = AtomicReferenceArray(HASH_BUCKET_COUNT)
+    private val hashBucketsL2: AtomicReferenceArray<Segment?> = AtomicReferenceArray(HASH_BUCKET_COUNT_L2)
 
     actual val byteCount: Int
         get() {
-            val first = firstRef().get() ?: return 0
+            val first = hashBuckets[l1BucketId()] ?: return 0
             return first.limit
         }
 
     @JvmStatic
     actual fun take(): Segment {
-        val firstRef = firstRef()
+        val buckets = hashBuckets
+        val bucketId = l1BucketId()
 
-        val first = firstRef.getAndSet(LOCK)
-        when {
-            first === LOCK -> {
-                // We didn't acquire the lock. Don't take a pooled segment.
-                return Segment.new()
-            }
+        while (true) {
+            when (val first = buckets.getAndSet(bucketId, LOCK)) {
+                LOCK -> {
+                    // We didn't acquire the lock. Let's try again
+                    continue
+                }
+
+                null -> {
+                    // We acquired the lock but the pool was empty.
+                    // Unlock the bucket and either try to acquire a segment from the second level cache,
+                    // or, if the second level cache is disabled, allocate a brand-new segment.
+                    buckets.set(bucketId, null)
+
+                    if (SECOND_LEVEL_POOL_TOTAL_SIZE > 0) {
+                        return takeL2()
+                    }
 
-            first == null -> {
-                // We acquired the lock but the pool was empty. Unlock and return a new segment.
-                firstRef.set(null)
-                return Segment.new()
+                    return Segment.new()
+                }
+
+                else -> {
+                    // We acquired the lock and the pool was not empty. Pop the first element and return it.
+                    buckets.set(bucketId, first.next)
+                    first.next = null
+                    first.limit = 0
+                    return first
+                }
             }
+        }
+    }
+
+    @JvmStatic
+    private fun takeL2(): Segment {
+        val buckets = hashBuckets
+        var bucket = l2BucketId()
+        var attempts = 0
+        while (true) {
+            when (val first = buckets.getAndSet(bucket, LOCK)) {
+                LOCK -> {
+                    // We didn't acquire the lock, retry
+                    continue
+                }
+
+                null -> {
+                    // We acquired the lock but the pool was empty.
+                    // Unlock the current bucket and select a new one.
+                    // If all buckets were already scanned, allocate a new segment.
+                    buckets.set(bucket, null)
+
+                    if (attempts < HASH_BUCKET_COUNT_L2) {
+                        bucket = (bucket + 1) and (HASH_BUCKET_COUNT_L2 - 1)
+                        attempts++
+                        continue
+                    }
 
-            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
+                    return Segment.new()
+                }
+
+                else -> {
+                    // We acquired the lock and the pool was not empty. Pop the first element and return it.
+                    buckets.set(bucket, first.next)
+                    first.next = null
+                    first.limit = 0
+                    return first
+                }
             }
         }
     }
@@ -109,29 +224,78 @@ internal actual object SegmentPool {
     @JvmStatic
     actual fun recycle(segment: Segment) {
         require(segment.next == null && segment.prev == null)
-        if (segment.shared) return // This segment cannot be recycled.
+        if (segment.copyTracker?.removeCopy() == true) return // This segment cannot be recycled.
+
+        val buckets = hashBuckets
+        val bucketId = l1BucketId()
+
+        segment.pos = 0
+        segment.owner = true
+
+        while (true) {
+            val first = buckets[bucketId]
+            if (first === LOCK) continue // A take() is currently in progress.
+            val firstLimit = first?.limit ?: 0
+            if (firstLimit >= MAX_SIZE) {
+                // L1 pool is full.
+                if (SECOND_LEVEL_POOL_TOTAL_SIZE > 0) {
+                    recycleL2(segment)
+                }
+                return
+            }
 
-        val firstRef = firstRef()
+            segment.next = first
+            segment.limit = firstLimit + Segment.SIZE
 
-        val first = firstRef.get()
-        if (first === LOCK) return // A take() is currently in progress.
-        val firstLimit = first?.limit ?: 0
-        if (firstLimit >= MAX_SIZE) return // Pool is full.
+            if (buckets.compareAndSet(bucketId, first, segment)) {
+                return
+            }
+        }
+    }
 
-        segment.next = first
+    @JvmStatic
+    private fun recycleL2(segment: Segment) {
         segment.pos = 0
-        segment.limit = firstLimit + Segment.SIZE
+        segment.owner = true
+
+        var bucket = l2BucketId()
+        val buckets = hashBucketsL2
+        var attempts = 0
+
+        while (true) {
+            val first = buckets[bucket]
+            if (first === LOCK) continue // A take() is currently in progress.
+            val firstLimit = first?.limit ?: 0
+            if (firstLimit + Segment.SIZE > SECOND_LEVEL_POOL_BUCKET_SIZE) {
+                // The current bucket is full, try to find another one and return the segment there.
+                if (attempts < HASH_BUCKET_COUNT_L2) {
+                    attempts++
+                    bucket = (bucket + 1) and (HASH_BUCKET_COUNT_L2 - 1)
+                    continue
+                }
+                // L2 pool is full.
+                return
+            }
 
-        // If we lost a race with another operation, don't recycle this segment.
-        if (!firstRef.compareAndSet(first, segment)) {
-            segment.next = null // Don't leak a reference in the pool either!
+            segment.next = first
+            segment.limit = firstLimit + Segment.SIZE
+
+            if (buckets.compareAndSet(bucket, first, segment)) {
+                return
+            }
         }
     }
 
-    private fun firstRef(): AtomicReference<Segment?> {
-        // Get a value in [0..HASH_BUCKET_COUNT) based on the current thread.
+    @JvmStatic
+    actual fun tracker(): SegmentCopyTracker = RefCountingCopyTracker()
+
+    private fun l1BucketId() = bucketId (HASH_BUCKET_COUNT - 1L)
+
+    private fun l2BucketId() = bucketId (HASH_BUCKET_COUNT_L2 - 1L)
+
+    private fun bucketId(mask: Long): Int {
+        // Get a value in [0..HASH_BUCKET_COUNT_L2) based on the current thread.
         @Suppress("DEPRECATION") // TODO: switch to threadId after JDK19
-        val hashBucket = (Thread.currentThread().id and (HASH_BUCKET_COUNT - 1L)).toInt()
-        return hashBuckets[hashBucket]
+        return (Thread.currentThread().id and mask).toInt()
     }
 }
diff --git a/core/native/src/SegmentPool.kt b/core/native/src/SegmentPool.kt
index 709e6f89..e2732153 100644
--- a/core/native/src/SegmentPool.kt
+++ b/core/native/src/SegmentPool.kt
@@ -29,4 +29,6 @@ internal actual object SegmentPool {
 
     actual fun recycle(segment: Segment) {
     }
+
+    actual fun tracker(): SegmentCopyTracker = AlwaysSharedCopyTracker
 }
diff --git a/core/wasm/src/SegmentPool.kt b/core/wasm/src/SegmentPool.kt
index 615a1497..c61fbeff 100644
--- a/core/wasm/src/SegmentPool.kt
+++ b/core/wasm/src/SegmentPool.kt
@@ -14,4 +14,6 @@ internal actual object SegmentPool {
 
     actual fun recycle(segment: Segment) {
     }
+
+    actual fun tracker(): SegmentCopyTracker = AlwaysSharedCopyTracker
 }


Test patch:
diff --git a/core/common/test/AlwaysSharedCopyTrackerTest.kt b/core/common/test/AlwaysSharedCopyTrackerTest.kt
new file mode 100644
index 00000000..62418c8c
--- /dev/null
+++ b/core/common/test/AlwaysSharedCopyTrackerTest.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+
+package kotlinx.io
+
+import kotlin.test.Test
+import kotlin.test.assertTrue
+
+class AlwaysSharedCopyTrackerTest {
+    @Test
+    fun stateTransition() {
+        val tracker = AlwaysSharedCopyTracker
+        assertTrue(tracker.shared)
+
+        assertTrue(tracker.removeCopy())
+        assertTrue(tracker.shared)
+
+        tracker.addCopy()
+        assertTrue(tracker.shared)
+
+        assertTrue(tracker.removeCopy())
+        assertTrue(tracker.shared)
+    }
+}
diff --git a/core/common/test/CommonBufferTest.kt b/core/common/test/CommonBufferTest.kt
index d450557a..d009c5f6 100644
--- a/core/common/test/CommonBufferTest.kt
+++ b/core/common/test/CommonBufferTest.kt
@@ -601,4 +601,21 @@ class CommonBufferTest {
         buffer.clear()
         assertEquals(ByteString(), buffer.snapshot())
     }
+
+    @Test
+    fun splitHead() {
+        val dst = Buffer()
+        val src = Buffer().also { it.write(ByteArray(SEGMENT_SIZE)) }
+
+        dst.write(src, src.size / 2)
+        assertEquals(src.size, dst.size)
+
+        assertEquals(src.head, src.tail)
+        assertEquals(null, src.head?.prev)
+        assertEquals(null, src.tail?.next)
+
+        assertEquals(dst.head, dst.tail)
+        assertEquals(null, dst.head?.prev)
+        assertEquals(null, dst.tail?.next)
+    }
 }
diff --git a/core/jvm/test/PoolingTest.kt b/core/jvm/test/PoolingTest.kt
new file mode 100644
index 00000000..95c5c36c
--- /dev/null
+++ b/core/jvm/test/PoolingTest.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+
+package kotlinx.io
+
+import kotlin.test.Test
+import kotlin.test.assertTrue
+
+class PoolingTest {
+    @Test
+    fun segmentSharing() {
+        val buffer = Buffer()
+
+        buffer.writeByte(1)
+        var poolSize = SegmentPool.byteCount
+        buffer.clear()
+        // clear should return a segment to a pool, so the pool size should grow
+        assertTrue(poolSize < SegmentPool.byteCount)
+
+        buffer.writeByte(1)
+        poolSize = SegmentPool.byteCount
+        val copy = buffer.copy()
+        buffer.clear()
+        copy.clear()
+        assertTrue(poolSize < SegmentPool.byteCount)
+
+        buffer.writeByte(1)
+        poolSize = SegmentPool.byteCount
+        val peek = buffer.peek().buffered()
+        peek.readByte()
+        buffer.clear()
+        assertTrue(poolSize < SegmentPool.byteCount)
+
+        buffer.writeByte(1)
+        poolSize = SegmentPool.byteCount
+        val otherBuffer = Buffer()
+        otherBuffer.write(buffer, buffer.size)
+        otherBuffer.clear()
+        assertTrue(poolSize < SegmentPool.byteCount)
+    }
+}
diff --git a/core/jvm/test/RefCounteringCopyTrackerTest.kt b/core/jvm/test/RefCounteringCopyTrackerTest.kt
new file mode 100644
index 00000000..49192ea3
--- /dev/null
+++ b/core/jvm/test/RefCounteringCopyTrackerTest.kt
@@ -0,0 +1,35 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+
+package kotlinx.io
+
+import kotlin.test.Test
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+class RefCounteringCopyTrackerTest {
+    @Test
+    fun stateTransition() {
+        val tracker = RefCountingCopyTracker()
+        assertFalse(tracker.shared)
+
+        assertFalse(tracker.removeCopy())
+        assertFalse(tracker.shared)
+
+        tracker.addCopy()
+        assertTrue(tracker.shared)
+        assertTrue(tracker.removeCopy())
+        assertFalse(tracker.shared)
+
+        tracker.addCopy()
+        assertTrue(tracker.shared)
+        tracker.addCopy()
+        assertTrue(tracker.shared)
+        assertTrue(tracker.removeCopy())
+        assertTrue(tracker.shared)
+        assertTrue(tracker.removeCopy())
+        assertFalse(tracker.shared)
+    }
+}