fix(level-hash/persistent): implement expansion logic for persistent level hash

2026-10-16 07:49:44 +08:00 · 2024-08-01 14:21:20 +05:30
parent dfb4f78b2f
commit 77bc0e2fdd
7 changed files with 249 additions and 68 deletions
@@ -126,17 +126,11 @@ interface LevelHash<K : Any, V : Any?> : AutoCloseable {
  fun findSlot(key: K): LevelSlot<K, V>?

  /**
-   * Expand the level hash by the given ADDITIONAL level size.
-   *
-   * @param addtionalLevelSize The additional level size of the level hash. If this is
-   * less than the current level size, an exception is thrown. If this is equal
-   * to the current level size, this function has no effect.
+   * Expand the level hash, doubling its capacity.
   *
   * @return Whether the level hash was resized.
-   * @throws IllegalArgumentException If the level size if less than the current
-   * level size.
   */
-  fun expand(addtionalLevelSize: Int): Boolean
+  fun expand(): Boolean

  /**
   * Get the load factor of the level hash.
@@ -201,7 +195,7 @@ interface LevelHash<K : Any, V : Any?> : AutoCloseable {
   * An [InsertionError] which is thrown when the level hash is full and unable
   * to insert new values.
   */
-  class OverflowException(capacity: Int, occupied: Int) : InsertionError("Level hash has reached is capacity. cap=$capacity occ=$occupied")
+  class OverflowException(capacity: Int, occupied: Int) : InsertionError("Level hash needs to be expanded. cap=$capacity occ=$occupied")

  /**
   * Thrown when the level hash cannot be resized.
@@ -123,16 +123,16 @@ internal abstract class AbstractLevelHash<K : Any, V : Any?> internal constructo

  override fun insert(key: K, value: V): Boolean {

-    if (this.levelItemCounts[0] == this.topLevelBucketCount * this.bucketSize) {
-      if (autoExpand) {
-        check(expand(1)) {
-          "Failed to expand the level hash"
-        }
-      } else {
-        throw OverflowException(this.totalSlotCount, this.levelItemCounts.sum())
+    if (loadFactor() >= 0.92 && autoExpand) {
+      check(expand()) {
+        "Failed to expand the level hash"
      }
    }

+    if (loadFactor() == 1f) {
+      throw OverflowException(totalSlotCount, levelItemCounts.sum())
+    }
+
    val insertFn =
      fun(levelIdx: Int, bucketIdx: Int, slotIdx: Int): Boolean {
        val slot = getSlot(levelIdx, bucketIdx, slotIdx)
@@ -245,15 +245,14 @@ internal abstract class AbstractLevelHash<K : Any, V : Any?> internal constructo
    return null
  }

-  override fun expand(addtionalLevelSize: Int): Boolean {
+  override fun expand(): Boolean {
    check(this.resizeState == RESIZE_STATE_NOT_RESIZING) {
      "Level hash is already being resized. Maybe on a different thread?"
    }
-    require(addtionalLevelSize > 0) { "levelSize must be > 0" }

    try {
      this.resizeState = RESIZE_STATE_EXPANDING
-      return doExpand(this.levelSize + addtionalLevelSize)
+      return doExpand(this.levelSize + 1)
    } finally {
      this.resizeState = RESIZE_STATE_NOT_RESIZING
    }
@@ -17,6 +17,7 @@

 package com.itsaky.androidide.levelhash.internal

+import androidx.annotation.VisibleForTesting
 import com.itsaky.androidide.levelhash.DataExternalizer
 import com.itsaky.androidide.levelhash.HashT
 import com.itsaky.androidide.levelhash.LevelHashFn
@@ -42,13 +43,14 @@ internal class PersistentLevelHash<K : Any, V : Any?>(
  uniqueKeys = uniqueKeys, autoExpand = autoExpand, levelHashFn = levelHashFn,
  seeds = seeds) {

-  private val io = PersistentLevelHashIO(indexFile, levelSize, bucketSize, keyExternalizer,
+  @VisibleForTesting
+  internal val io = PersistentLevelHashIO(indexFile, levelSize, bucketSize, keyExternalizer,
    valueExternalizer)

  override var levelSize: Int
-    get() = io.metaIo.levelSize
+    get() = io.levelSize
    set(value) {
-      io.metaIo.levelSize = value
+      io.levelSize = value
    }

  override fun getSlot(levelIdx: Int, bucketIdx: Int, slotIdx: Int
@@ -64,11 +66,21 @@ internal class PersistentLevelHash<K : Any, V : Any?>(
    io.close()
  }

+  override fun prepareExpansion(bucketCount: Int) {
+    super.prepareExpansion(bucketCount)
+    io.prepareInterimLevel(bucketCount)
+  }
+
  override fun moveForExpansion(slot: LevelSlot<K, V>, bucketIdx: Int,
                                slotIdx: Int
  ): Boolean {
    slot as PersistentLevelSlot<K, V>
-    return false
+    return io.moveToInterim(slot.levelIdx, slot.bucketIdx, slot.slotIdx, bucketIdx, slotIdx)
+  }
+
+  override fun onExpand(newLevelSize: Int, interimItemCount: Int) {
+    super.onExpand(newLevelSize, interimItemCount)
+    io.finalizeExpansion()
  }

  private class PersistentLevelSlot<K : Any, V : Any?>(
@@ -17,6 +17,7 @@

 package com.itsaky.androidide.levelhash.internal

+import androidx.annotation.VisibleForTesting
 import com.itsaky.androidide.levelhash.DataExternalizer
 import com.itsaky.androidide.levelhash.LevelHash.ResizeFailure
 import com.itsaky.androidide.levelhash.seekInt
@@ -32,6 +33,16 @@ import kotlin.math.pow

 private val logger = LoggerFactory.getLogger(PersistentLevelHashIO::class.java)

+// TODO: Update valuesFirstEntry and valuesNxtEntry whenever an entry is added or deleted
+// TODO: Entries in the values file are always appended at the end. When an entry
+//    is to be deleted, we simply use fallocate with FALLOC_FL_PUNCH_HOLE to deallocate
+//    the disk space occupied by that entry. However, this may result in the region
+//    at the start of the values file become empty. In that case, we might consider
+//    moving the entries at the start of the file. This may (or may not) result in
+//    improved performance when seeking through the file.
+// TODO: The above is also applicable for the keymap file when the level hash is
+//    expanded and new regions for the levels are allocated.
+
 /**
 * This class is responsible for performing the IO operations for the level hash
 * backed by the disk.
@@ -162,9 +173,11 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  // as the values file is randomly accessed, and by multiple functions, or
  // even multiple threads, the current position cannot be determined

+  @VisibleForTesting
  internal val metaIo =
    PersistentMetaIO(File(indexFile.parentFile, "${indexFile.name}._meta"),
      levelSize, bucketSize)
+
  private val keymapFile = File(indexFile.parentFile, "${indexFile.name}._i")
  private val raKeymapFile by lazy { RandomAccessFile(keymapFile, "rw") }
  private val raValuesFile by lazy { RandomAccessFile(indexFile, "rw") }
@@ -172,17 +185,18 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  private val keymapIo = MappedRandomAccessIO()
  private var valIo = MappedRandomAccessIO()

+  private var interimLvlAddr: Long? = null
+
  private val topLevelBucketCount: Long
-    get() = 2.0.pow(metaIo.levelSize).toLong()
+    get() = 2.0.pow(metaIo.levelSize)
+      .toLong()

  private val topLevelSizeBytes: Long
    get() = topLevelBucketCount * metaIo.bucketSize * KEYMAP_ENTRY_SIZE_BYTES

-  private val totalBucketCount: Long
-    get() = topLevelBucketCount + (topLevelBucketCount shr 1)
-
-  private val totalSlotCount: Long
-    get() = totalBucketCount * metaIo.bucketSize
+  internal var levelSize: Int
+    get() = metaIo.levelSize
+    set(value) { metaIo.levelSize = value }

  init {
    BinaryFileUtils.initSparseFile(file = keymapFile,
@@ -196,8 +210,8 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  }

  private fun realValueOffset(offset: Long) = VALUES_HEADER_SIZE_BYTES + offset
-
  private fun realKeymapOffset(offset: Long) = KEYMAP_HEADER_SIZE_BYTES + offset
+  private fun keymapBufOffset(offset: Long) = offset - KEYMAP_HEADER_SIZE_BYTES

  private fun valuesDeallocate(offset: Long, len: Long) =
    FileUtils.deallocate(raValuesFile, realValueOffset(offset), len)
@@ -212,13 +226,13 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
   *
   * @param newSize The new size of the values file, in bytes.
   */
-  private fun valuesResize(newSize: Long)  {
+  private fun valuesResize(newSize: Long) {
    metaIo.valuesFileSize = newSize
    raValuesFile.setLength(newSize)
    valuesRemap(newSize)
  }

-  private fun keymapResize(newSize: Long)  {
+  private fun keymapResize(newSize: Long) {
    if (raKeymapFile.length() != newSize) {
      raKeymapFile.setLength(newSize)
    }
@@ -236,7 +250,7 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  ) {
    val buffer = file.channel.map(FileChannel.MapMode.READ_WRITE, offset, size)
    io.close()
-    io.reset(buffer, 0L, size)
+    io.reset(buffer)
  }

  private fun slotAddress(
@@ -244,13 +258,21 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
    bucketIdx: Int,
    slotIdx: Int,
  ): Long {
-    val levelPos = when (levelIdx) {
+    val levelAddr = when (levelIdx) {
      0 -> metaIo.l0Addr
      1 -> metaIo.l1Addr
      else -> throw IllegalArgumentException("Invalid level index: $levelIdx")
    }

-    return levelPos + // start position of level
+    return slotAddrForLvlAddr(levelAddr, bucketIdx, slotIdx)
+  }
+
+  private fun slotAddrForLvlAddr(
+    lvlAddr: Long,
+    bucketIdx: Int,
+    slotIdx: Int,
+  ): Long {
+    return lvlAddr + // start position of level
      (KEYMAP_ENTRY_SIZE_BYTES * metaIo.bucketSize * bucketIdx) + // bucket position
      (KEYMAP_ENTRY_SIZE_BYTES * slotIdx) // slot position in bucket
  }
@@ -391,8 +413,6 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
    return value
  }

-
-
  fun writeEntry(
    levelNum: Int,
    bucketIdx: Int,
@@ -518,11 +538,17 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  }

  fun clear() {
-    var keymapSize = topLevelSizeBytes
-    keymapSize += topLevelSizeBytes shr 1
+    metaIo.valuesFirstEntry = 0L
+    metaIo.valuesNextEntry = 0L
+    metaIo.l0Addr = 0
+    metaIo.l1Addr = 2.0.pow(metaIo.levelSize).toLong() * metaIo.bucketSize * KEYMAP_ENTRY_SIZE_BYTES
+    val kmSize = metaIo.l1Addr + (metaIo.l1Addr shr 1)

-    keymapDeallocate(KEYMAP_HEADER_SIZE_BYTES, keymapSize)
-    valuesDeallocate(VALUES_HEADER_SIZE_BYTES, metaIo.valuesFileSize)
+    keymapResize(realKeymapOffset(kmSize))
+    keymapDeallocate(0, kmSize)
+
+    valuesResize(realValueOffset(VALUES_SEGMENT_SIZE_BYTES))
+    valuesDeallocate(0, VALUES_SEGMENT_SIZE_BYTES)
  }

  private fun <T> writeSlotKeyOrVal(obj: T?, externalizer: DataExternalizer<T>
@@ -550,6 +576,72 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
    }
  }

+  internal fun prepareInterimLevel(bucketCount: Int) {
+    check(interimLvlAddr == null) {
+      "Interim level address must be null. Is the level hash being resized on multiple threads?"
+    }
+
+    val interimSize = bucketCount * metaIo.bucketSize * KEYMAP_ENTRY_SIZE_BYTES
+    val start = metaIo.kmStartAddr
+    if (realKeymapOffset(interimSize) < start) {
+      // we can reuse the empty space at the beginning of the file
+      interimLvlAddr = 0L
+    } else {
+      // ensure the keymap can accommodate the interim level
+      val length = raKeymapFile.length()
+      keymapResize(length + interimSize)
+      interimLvlAddr = keymapBufOffset(length)
+    }
+  }
+
+  internal fun moveToInterim(
+    levelNum: Int,
+    bucketIdx: Int,
+    slotIdx: Int,
+    interimBucketIdx: Int,
+    interimSlotIdx: Int,
+  ): Boolean {
+    // current (source) slot
+    val sSlotAddr = slotAddress(levelNum, bucketIdx, slotIdx)
+
+    // destination slot
+    val dSlotAddr =
+      slotAddrForLvlAddr(interimLvlAddr!!, interimBucketIdx, interimSlotIdx)
+
+    keymapIo.position(dSlotAddr)
+    val dValAddr = keymapIo.readLong()
+    if (dValAddr > 0) {
+      // interim (destination) slot is occupied
+      return false
+    }
+
+    // interim slot is not occupied
+    // as the values in the keymap are just pointers to the actual value entries
+    // in the values file, we just need to update destination slot to point
+    // where the source slot points
+
+    // 1. read the address where the source slot points
+    keymapIo.position(sSlotAddr)
+    val eValAddr = keymapIo.readLong()
+
+    // 2. move the destination slot and write the address of the source slot's value
+    keymapIo.position(dSlotAddr)
+    keymapIo.writeLong(eValAddr)
+
+    // 3. deallocate the space occupied by the source slot
+    keymapDeallocate(sSlotAddr, KEYMAP_HEADER_SIZE_BYTES)
+
+    return true
+  }
+
+  internal fun finalizeExpansion() {
+    // current top level becomes the bottom level
+    // and interim level becomes the new top level
+    metaIo.l1Addr = metaIo.l0Addr
+    metaIo.l0Addr = interimLvlAddr!!
+    interimLvlAddr = null
+  }
+
  override fun close() {
    closeAndLogErrs(metaIo, keymapIo, valIo, raKeymapFile, raValuesFile)
  }
@@ -17,6 +17,7 @@

 package com.itsaky.androidide.levelhash.internal

+import androidx.annotation.VisibleForTesting
 import androidx.collection.MutableLongLongMap
 import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.KEYMAP_ENTRY_SIZE_BYTES
 import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.LEVEL_KEYMAP_VERSION
@@ -78,14 +79,15 @@ internal class PersistentMetaIO(private val metaFile: File,
    get() = readLong(KM__L1_ADDR__OFFSET)
    set(value) { writeLong(KM__L1_ADDR__OFFSET, value) }

+  val kmStartAddr: Long
+    get() = min(l0Addr, l1Addr)
+
  val keymapSize: Long
    get() {
      // for this to work, l0 and l1 must be in a contiguosly allocated space
      // in the keymap file
-      val l0Addr = this.l0Addr
-      val l1Addr = this.l1Addr
      val l0SizeBytes = 2.0.pow(levelSize).toLong() * bucketSize * KEYMAP_ENTRY_SIZE_BYTES
-      var size = min(l0Addr, l1Addr)
+      var size = kmStartAddr
      size += l0SizeBytes
      size += l0SizeBytes shr 1
      return size
@@ -105,12 +107,6 @@ internal class PersistentMetaIO(private val metaFile: File,
    if (valuesFileSize == 0L) {
      valuesFileSize = VALUES_SEGMENT_SIZE_BYTES
    }
-    if (valuesFirstEntry == 0L) {
-      valuesFirstEntry = 0
-    }
-    if (valuesNextEntry == 0L) {
-      valuesNextEntry = 0
-    }
    if (this.levelSize != levelSize) {
      this.levelSize = levelSize
    }
@@ -168,13 +164,18 @@ internal class PersistentMetaIO(private val metaFile: File,

  override fun close() {
    try {
-      cache.clear()
+      clearCache()
      raFile.close()
    } catch (e: Exception) {
      log.error("Failed to close meta file", e)
    }
  }

+  @VisibleForTesting
+  internal fun clearCache() {
+    this.cache.clear()
+  }
+
  private inline fun <reified T : Number> readOrDefault(default: T,
                                                        read: () -> T
  ): T {
@@ -31,6 +31,7 @@ import java.io.DataInputStream
 import java.io.DataOutput
 import java.io.EOFException
 import java.nio.MappedByteBuffer
+import kotlin.math.max
 import kotlin.math.min

 /**
@@ -43,7 +44,6 @@ internal class MappedRandomAccessIO : DataInput, DataOutput, RandomAccessIO, Aut

  private var positionOffset = 0L
  private var position = 0L
-  private var size = 0L

  private lateinit var buffer: MappedByteBuffer

@@ -57,10 +57,9 @@ internal class MappedRandomAccessIO : DataInput, DataOutput, RandomAccessIO, Aut
    return this.buffer.apply { position(position.toInt()) }
  }

-  fun reset(buffer: MappedByteBuffer, position: Long, length: Long, positionOffset: Long = 0L) {
+  fun reset(buffer: MappedByteBuffer, position: Long = 0L, positionOffset: Long = 0L) {
    this.buffer = buffer
    this.position = position
-    this.size = length
    this.positionOffset = positionOffset
  }

@@ -76,16 +75,25 @@ internal class MappedRandomAccessIO : DataInput, DataOutput, RandomAccessIO, Aut
    }
  }

+  override fun seekRelative(count: Long) {
+    if (count == 0L) {
+      return
+    }
+    var pos = position() + count
+    pos = max(0L, min(size().toLong(), pos))
+    position(pos)
+  }
+
  override fun position(position: Long) {
    require(position >= 0L) { "position must be >= 0, position=$position" }
-    require(position < buffer.limit()) {
-      "position must be < buffer.limit(), position=$position buffer.limit=${buffer.limit()}"
+    require(position <= size()) {
+      "position must be < buffer.limit(), position=$position buffer.limit=${size()}"
    }
-    tryPosition(position)
+    check(tryPosition(position))
  }

  override fun tryPosition(position: Long): Boolean {
-    if (position < 0L || position >= buffer.limit()) {
+    if (position < 0L || position > size()) {
      return false
    }
    val currentPos = position()
@@ -100,7 +108,7 @@ internal class MappedRandomAccessIO : DataInput, DataOutput, RandomAccessIO, Aut
    return this.position
  }

-  fun size() = this.size
+  fun size() = buffer.limit()

  override fun readFully(bytes: ByteArray) {
    if (bytes.isEmpty()) return
@@ -130,7 +138,7 @@ internal class MappedRandomAccessIO : DataInput, DataOutput, RandomAccessIO, Aut
  }

  override fun readByte(): Byte {
-    return buf(position).get().also { seekByte() }
+    return buf().get().also { seekByte() }
  }

  override fun readUnsignedByte(): Int {
@@ -19,6 +19,8 @@ package com.itsaky.androidide.levelhash

 import com.google.common.truth.Truth.assertThat
 import com.itsaky.androidide.levelhash.LevelHash.OverflowException
+import com.itsaky.androidide.levelhash.internal.PersistentLevelHash
+import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.KEYMAP_ENTRY_SIZE_BYTES
 import com.itsaky.androidide.levelhash.util.DataExternalizers
 import com.itsaky.androidide.levelhash.util.nullable
 import com.itsaky.androidide.testing.common.LinuxOnlyTestRule
@@ -75,7 +77,7 @@ class PersistentLevelHashTest {
  }

  @Test
-  fun testInsertion() {
+  fun `test insertion`() {
    persistentStringHash("insert").use { hash ->
      val key = "key"
      val value = "value"
@@ -88,7 +90,7 @@ class PersistentLevelHashTest {
  }

  @Test(expected = OverflowException::class)
-  fun testOverflow() {
+  fun `test overflow`() {
    persistentStringHash("overflow") { autoExpand(false) }.use { hash ->
      for (i in 0..<hash.totalSlotCount) {
        val key = "key${i}"
@@ -102,7 +104,11 @@ class PersistentLevelHashTest {

  @Test
  fun `test values file expands and remaps in case of buffer overflow`() {
-    persistentStringHash("huge") { autoExpand(false).levelSize(15).bucketSize(10) }.use { hash ->
+    persistentStringHash("huge") {
+      autoExpand(true)
+      levelSize(15)
+      bucketSize(10)
+    }.use { hash ->
      for (i in 0..<hash.totalSlotCount) {
        val key = "key${i}"
        val value = "value${i}"
@@ -115,7 +121,7 @@ class PersistentLevelHashTest {
  }

  @Test
-  fun testRemoval() {
+  fun `test removal`() {
    persistentStringHash("remove").use { hash ->
      assertThat(hash.insert("key", "value")).isTrue()
      assertThat(hash.findSlot("key")).isNotNull()
@@ -203,22 +209,91 @@ class PersistentLevelHashTest {
    persistentStringHash(
      fileName = "expand",
      valueExternalizer = DataExternalizers.STRING.nullable()
-    ).use { hash ->
-      for (i in 0..<hash.totalSlotCount) {
+    ) {
+      levelSize(5)
+      bucketSize(10)
+      autoExpand(false)
+    }.use { hash ->
+      val slots = hash.totalSlotCount - hash.bucketSize
+      for (i in 0..<slots) {
        val key = "key${i}"
        val value = "value${i}"
        assertThat(hash.insert(key, value)).isTrue()
      }

      // expand the level hash
-      hash.expand(1)
+      hash.expand()

      // then check that all existing entries are still present
-      for (i in 0..<hash.totalSlotCount) {
+      for (i in 0..<slots) {
        val key = "key${i}"
        val value = "value${i}"
        assertThat(hash[key]).isEqualTo(value)
      }
    }
  }
+
+  @Test
+  fun `test expansion with huge slot count`() {
+    persistentStringHash(
+      fileName = "huge-expand",
+      valueExternalizer = DataExternalizers.STRING.nullable()
+    ){
+      levelSize(15)
+      bucketSize(10)
+      autoExpand(false)
+    }.use { hash ->
+      var i = 0
+      while(hash.insert("key${i}", "value${i}")) {
+        i++
+      }
+
+      hash.expand()
+
+      for (j in 0..<i) {
+        val key = "key${j}"
+        val value = "value${j}"
+        assertThat(hash[key]).isEqualTo(value)
+      }
+    }
+  }
+
+  @Test
+  fun `test meta after expansion`() {
+    persistentStringHash(
+      fileName = "expand",
+      valueExternalizer = DataExternalizers.STRING.nullable()
+    ) {
+      levelSize(5)
+      bucketSize(10)
+      autoExpand(false)
+    }.use { hash ->
+      hash as PersistentLevelHash<String, String?>
+
+      val l0Size =
+        hash.topLevelBucketCount * hash.io.metaIo.bucketSize * KEYMAP_ENTRY_SIZE_BYTES
+
+      assertThat(hash.io.metaIo.levelSize).isEqualTo(5)
+      assertThat(hash.io.metaIo.bucketSize).isEqualTo(10)
+      assertThat(hash.io.metaIo.l0Addr).isEqualTo(0)
+      assertThat(hash.io.metaIo.l1Addr).isEqualTo(l0Size)
+
+      hash.expand()
+
+      // clear cache so we read values directly from the file
+      hash.io.metaIo.clearCache()
+
+      assertThat(hash.io.metaIo.levelSize).isEqualTo(6)
+      assertThat(hash.io.metaIo.bucketSize).isEqualTo(10)
+
+      // top level (l0) is moved after the bottom level (l1)
+      // the bottom level used to span from `l0Size` to `l0Size + l0Size / 2`
+      // and the interim level was added AFTER the previous bottom level
+      assertThat(hash.io.metaIo.l0Addr).isEqualTo(l0Size + (l0Size shr 1))
+
+      // previous top level, which used to be located at 0,
+      // becomes the new bottom level
+      assertThat(hash.io.metaIo.l1Addr).isEqualTo(0L)
+    }
+  }
 }