Merge pull request #88 from target/getIntoOnHeapMemoryBuffer

Add NativeMemoryMap.getIntoOnHeapMemoryBuffer

Author: aaronriekenberg

src/main/kotlin/com/target/nativememoryallocator/map/NativeMemoryMap.kt

@@ -1,6 +1,7 @@
 package com.target.nativememoryallocator.map
 
 import com.github.benmanes.caffeine.cache.stats.CacheStats
+import com.target.nativememoryallocator.buffer.NativeMemoryBuffer
 import com.target.nativememoryallocator.buffer.NativeMemoryBufferMetadata
 import com.target.nativememoryallocator.buffer.OnHeapMemoryBuffer
 
@@ -86,6 +87,13 @@ interface NativeMemoryMap<KEY_TYPE : Any, VALUE_TYPE : Any> : BaseNativeMemoryMa
      */
     fun get(key: KEY_TYPE): VALUE_TYPE?
 
+    /**
+     * Get a value from the map using the [key], copying the native memory contents into [onHeapMemoryBuffer].
+     *
+     * @return true if key was found and value copied into [onHeapMemoryBuffer], false otherwise.
+     */
+    fun getIntoOnHeapMemoryBuffer(key: KEY_TYPE, onHeapMemoryBuffer: OnHeapMemoryBuffer): Boolean
+
     /**
      * [Set] of [KEY_TYPE] for the map.
      *

src/main/kotlin/com/target/nativememoryallocator/map/impl/NativeMemoryMapImpl.kt

@@ -126,6 +126,20 @@ internal class NativeMemoryMapImpl<KEY_TYPE : Any, VALUE_TYPE : Any>(
         }
     }
 
+    override fun getIntoOnHeapMemoryBuffer(key: KEY_TYPE, onHeapMemoryBuffer: OnHeapMemoryBuffer): Boolean {
+        var found = false
+
+        cacheMap.computeIfPresent(key) { _, currentBuffer ->
+            found = true
+
+            currentBuffer.copyToOnHeapMemoryBuffer(onHeapMemoryBuffer = onHeapMemoryBuffer)
+
+            currentBuffer
+        }
+
+        return found
+    }
+
     override val keys: Set<KEY_TYPE>
         get() = cacheMap.keys
 

src/main/kotlin/com/target/nativememoryallocator/map/impl/OperationCountedNativeMemoryMapImpl.kt

@@ -1,5 +1,6 @@
 package com.target.nativememoryallocator.map.impl
 
+import com.target.nativememoryallocator.buffer.OnHeapMemoryBuffer
 import com.target.nativememoryallocator.map.NativeMemoryMap
 import com.target.nativememoryallocator.map.NativeMemoryMapOperationCounters
 import java.util.concurrent.atomic.AtomicLong
@@ -76,6 +77,23 @@ internal class OperationCountedNativeMemoryMapImpl<KEY_TYPE : Any, VALUE_TYPE :
         return getResult
     }
 
+    /**
+     * Delegate to [NativeMemoryMap.getIntoOnHeapMemoryBuffer] and then update [operationCounters].
+     */
+    override fun getIntoOnHeapMemoryBuffer(key: KEY_TYPE, onHeapMemoryBuffer: OnHeapMemoryBuffer): Boolean {
+        val getResult = nativeMemoryMap.getIntoOnHeapMemoryBuffer(key = key, onHeapMemoryBuffer = onHeapMemoryBuffer)
+
+        operationCounters.apply {
+            if (getResult) {
+                numGetsNonNullValue.incrementAndGet()
+            } else {
+                numGetsNullValue.incrementAndGet()
+            }
+        }
+
+        return getResult
+    }
+
     /**
      * Delegate to [NativeMemoryMap.put] and then update [operationCounters].
      */

src/test/kotlin/com/target/map/impl/NativeMemoryMapImplTest.kt

@@ -1,6 +1,7 @@
 package com.target.map.impl
 
 import com.target.nativememoryallocator.allocator.NativeMemoryAllocator
+import com.target.nativememoryallocator.allocator.NativeMemoryAllocatorBuilder
 import com.target.nativememoryallocator.buffer.NativeMemoryBuffer
 import com.target.nativememoryallocator.buffer.OnHeapMemoryBuffer
 import com.target.nativememoryallocator.buffer.OnHeapMemoryBufferFactory
@@ -279,6 +280,7 @@ class NativeMemoryMapImplTest {
         }
     }
 
+
     @Test
     fun `test put reuse buffer`() {
         val serializedValue1 = ByteArray(10)
@@ -441,4 +443,54 @@ class NativeMemoryMapImplTest {
             mockNativeMemoryAllocator.freeNativeMemoryBuffer(buffer = mockNativeMemoryBuffer)
         }
     }
+
+    @Test
+    fun `test put then getIntoOnHeapMemoryBuffer`() {
+
+        class NMAStringSerializer : NativeMemoryMapSerializer<String> {
+
+            override fun deserializeFromOnHeapMemoryBuffer(onHeapMemoryBuffer: OnHeapMemoryBuffer): String =
+                String(
+                    onHeapMemoryBuffer.array, 0, onHeapMemoryBuffer.getReadableBytes()
+                )
+
+            override fun serializeToByteArray(value: String): ByteArray = value.toByteArray()
+        }
+
+        clearAllMocks()
+
+        val nativeMemoryAllocator = NativeMemoryAllocatorBuilder(
+            pageSizeBytes = 1024,
+            nativeMemorySizeBytes = 16 * 1024 * 1024, // 16 MB
+        ).build()
+
+        val testValue = "hello world"
+
+        val nativeMemoryMap = NativeMemoryMapImpl(
+            valueSerializer = NMAStringSerializer(),
+            nativeMemoryAllocator = nativeMemoryAllocator,
+            useThreadLocalOnHeapReadBuffer = false,
+            threadLocalOnHeapReadBufferInitialCapacityBytes = (256 * 1024),
+            cacheMap = ConcurrentHashMap(),
+        )
+
+        val putResult = nativeMemoryMap.put(key = 1, value = testValue)
+        putResult shouldBe NativeMemoryMap.PutResult.ALLOCATED_NEW_BUFFER
+
+        val readOnHeapMemoryBuffer = OnHeapMemoryBufferFactory.newOnHeapMemoryBuffer(initialCapacityBytes = 1)
+
+        nativeMemoryMap.getIntoOnHeapMemoryBuffer(
+            key = 1,
+            onHeapMemoryBuffer = readOnHeapMemoryBuffer
+        ) shouldBe true
+
+        readOnHeapMemoryBuffer.getReadableBytes() shouldBe 11
+        String(readOnHeapMemoryBuffer.toTrimmedArray(), 0, readOnHeapMemoryBuffer.getReadableBytes()) shouldBe testValue
+
+        nativeMemoryMap.getIntoOnHeapMemoryBuffer(key = 2, onHeapMemoryBuffer = readOnHeapMemoryBuffer) shouldBe false
+
+        nativeMemoryMap.keys shouldBe setOf(1)
+        nativeMemoryMap.size shouldBe 1
+    }
+
 }