Skip to content

Commit

Permalink
Use Handle for passing foreign objects to Rust
Browse files Browse the repository at this point in the history
Consolidated the various handle maps into a single implementation for
each language.  This handle map works basically the same as all the
others, but it's API is based on the `HandleAlloc` trait.  Handles have
a couple properties:

* All foreign handles are odd, which allows us to distinguish between
  Rust and foreign handles.
* For handles store a map ID that can detect when a handle is used with
  the wrong map.

Made all languages always use the handle maps for passing objects.  No
more trying to leak pointers from to foreign objects.

Started updating the ForeignExecutor code to use handles, but this is
still a WIP while the ForeignExecutor type is in it's limbo state.
  • Loading branch information
bendk committed Nov 20, 2023
1 parent 9f3f914 commit bc5e827
Show file tree
Hide file tree
Showing 52 changed files with 388 additions and 533 deletions.
6 changes: 5 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,11 @@
- The `rust_future_continuation_callback_set` FFI function was removed. `rust_future_poll` now
inputs the callback pointer. External bindings authors will need to update their code.
- The object handle FFI has changed. External bindings generators will need to update their code to
use the new handle system.
use the new handle system:
* A single `FfiType::Handle` is used for all object handles.
* `FfiType::Handle` is always a 64-bit int.
* Foreign handles must always set the lowest bit of that int.
- The `NoPointer` singleton was renamed to `NoHandle`

### What's new?

Expand Down
4 changes: 2 additions & 2 deletions fixtures/coverall/tests/bindings/test_coverall.kts
Original file line number Diff line number Diff line change
Expand Up @@ -446,11 +446,11 @@ Coveralls("test_bytes").use { coveralls ->

// Test fakes using open classes

class FakePatch(private val color: Color): Patch(NoPointer) {
class FakePatch(private val color: Color): Patch(NoHandle) {
override fun `getColor`(): Color = color
}

class FakeCoveralls(private val name: String) : Coveralls(NoPointer) {
class FakeCoveralls(private val name: String) : Coveralls(NoHandle) {
private val repairs = mutableListOf<Repair>()

override fun `addPatch`(patch: Patch) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,3 @@ runBlocking {
assert(result.delayMs <= 200U)
tester.close()
}

// Test that we cleanup when dropping a ForeignExecutor handles
assert(FfiConverterForeignExecutor.handleCount() == 0)
val tester = ForeignExecutorTester(coroutineScope)
val tester2 = ForeignExecutorTester.newFromSequence(listOf(coroutineScope))
tester.close()
tester2.close()
assert(FfiConverterForeignExecutor.handleCount() == 0)
2 changes: 0 additions & 2 deletions uniffi_bindgen/src/bindings/kotlin/gen_kotlin/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -333,12 +333,10 @@ impl KotlinCodeOracle {
}
FfiType::ForeignBytes => "ForeignBytes.ByValue".to_string(),
FfiType::ForeignCallback => "ForeignCallback".to_string(),
FfiType::ForeignExecutorHandle => "USize".to_string(),
FfiType::ForeignExecutorCallback => "UniFfiForeignExecutorCallback".to_string(),
FfiType::RustFutureContinuationCallback => {
"UniFffiRustFutureContinuationCallbackType".to_string()
}
FfiType::RustFutureContinuationData => "USize".to_string(),
}
}

Expand Down
10 changes: 5 additions & 5 deletions uniffi_bindgen/src/bindings/kotlin/templates/Async.kt
Original file line number Diff line number Diff line change
Expand Up @@ -3,18 +3,18 @@
internal const val UNIFFI_RUST_FUTURE_POLL_READY = 0.toShort()
internal const val UNIFFI_RUST_FUTURE_POLL_MAYBE_READY = 1.toShort()

internal val uniffiContinuationHandleMap = UniFfiHandleMap<CancellableContinuation<Short>>()
internal val uniffiContinuationHandleMap = UniffiHandleMap<CancellableContinuation<Short>>()

// FFI type for Rust future continuations
internal object uniffiRustFutureContinuationCallback: UniFffiRustFutureContinuationCallbackType {
override fun callback(continuationHandle: USize, pollResult: Short) {
uniffiContinuationHandleMap.remove(continuationHandle)?.resume(pollResult)
override fun callback(continuationHandle: UniffiHandle, pollResult: Short) {
uniffiContinuationHandleMap.consumeHandle(continuationHandle).resume(pollResult)
}
}

internal suspend fun<T, F, E: Exception> uniffiRustCallAsync(
rustFuture: UniffiHandle,
pollFunc: (UniffiHandle, UniFffiRustFutureContinuationCallbackType, USize) -> Unit,
pollFunc: (UniffiHandle, UniFffiRustFutureContinuationCallbackType, UniffiHandle) -> Unit,
completeFunc: (UniffiHandle, RustCallStatus) -> F,
freeFunc: (UniffiHandle) -> Unit,
liftFunc: (F) -> T,
Expand All @@ -26,7 +26,7 @@ internal suspend fun<T, F, E: Exception> uniffiRustCallAsync(
pollFunc(
rustFuture,
uniffiRustFutureContinuationCallback,
uniffiContinuationHandleMap.insert(continuation)
uniffiContinuationHandleMap.newHandle(continuation)
)
}
} while (pollResult != UNIFFI_RUST_FUTURE_POLL_READY);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,11 @@
// Implement the foreign callback handler for {{ interface_name }}
internal class {{ callback_handler_class }} : ForeignCallback {
@Suppress("TooGenericExceptionCaught")
override fun invoke(handle: Handle, method: Int, argsData: Pointer, argsLen: Int, outBuf: RustBufferByReference): Int {
override fun invoke(handle: UniffiHandle, method: Int, argsData: Pointer, argsLen: Int, outBuf: RustBufferByReference): Int {
val cb = {{ ffi_converter_name }}.handleMap.get(handle)
return when (method) {
IDX_CALLBACK_FREE -> {
{{ ffi_converter_name }}.handleMap.remove(handle)
{{ ffi_converter_name }}.handleMap.consumeHandle(handle)

// Successful return
// See docs of ForeignCallback in `uniffi_core/src/ffi/foreigncallbacks.rs`
Expand Down
Original file line number Diff line number Diff line change
@@ -1,39 +1,8 @@
{{- self.add_import("java.util.concurrent.atomic.AtomicLong") }}
{{- self.add_import("java.util.concurrent.locks.ReentrantLock") }}
{{- self.add_import("kotlin.concurrent.withLock") }}

internal typealias Handle = Long
internal class ConcurrentHandleMap<T>(
private val leftMap: MutableMap<Handle, T> = mutableMapOf(),
) {
private val lock = java.util.concurrent.locks.ReentrantLock()
private val currentHandle = AtomicLong(0L)
private val stride = 1L

fun insert(obj: T): Handle =
lock.withLock {
currentHandle.getAndAdd(stride)
.also { handle ->
leftMap[handle] = obj
}
}

fun get(handle: Handle) = lock.withLock {
leftMap[handle] ?: throw InternalException("No callback in handlemap; this is a Uniffi bug")
}

fun delete(handle: Handle) {
this.remove(handle)
}

fun remove(handle: Handle): T? =
lock.withLock {
leftMap.remove(handle)
}
}

interface ForeignCallback : com.sun.jna.Callback {
public fun invoke(handle: Handle, method: Int, argsData: Pointer, argsLen: Int, outBuf: RustBufferByReference): Int
public fun invoke(handle: UniffiHandle, method: Int, argsData: Pointer, argsLen: Int, outBuf: RustBufferByReference): Int
}

// Magic number for the Rust proxy to call using the same mechanism as every other method,
Expand All @@ -44,20 +13,16 @@ internal const val UNIFFI_CALLBACK_SUCCESS = 0
internal const val UNIFFI_CALLBACK_ERROR = 1
internal const val UNIFFI_CALLBACK_UNEXPECTED_ERROR = 2

public abstract class FfiConverterCallbackInterface<CallbackInterface>: FfiConverter<CallbackInterface, Handle> {
internal val handleMap = ConcurrentHandleMap<CallbackInterface>()

internal fun drop(handle: Handle) {
handleMap.remove(handle)
}
public abstract class FfiConverterCallbackInterface<CallbackInterface>: FfiConverter<CallbackInterface, UniffiHandle> {
internal val handleMap = UniffiHandleMap<CallbackInterface>()

override fun lift(value: Handle): CallbackInterface {
override fun lift(value: UniffiHandle): CallbackInterface {
return handleMap.get(value)
}

override fun read(buf: ByteBuffer) = lift(buf.getLong())

override fun lower(value: CallbackInterface) = handleMap.insert(value)
override fun lower(value: CallbackInterface) = handleMap.newHandle(value)

override fun allocationSize(value: CallbackInterface) = 8

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ internal interface UniFfiRustTaskCallback : com.sun.jna.Callback {
}

internal object UniFfiForeignExecutorCallback : com.sun.jna.Callback {
fun callback(handle: USize, delayMs: Int, rustTask: UniFfiRustTaskCallback?, rustTaskData: Pointer?) : Byte {
fun callback(handle: UniffiHandle, delayMs: Int, rustTask: UniFfiRustTaskCallback?, rustTaskData: Pointer?) : Byte {
if (rustTask == null) {
FfiConverterForeignExecutor.drop(handle)
return UNIFFI_FOREIGN_EXECUTOR_CALLBACK_SUCCESS
Expand All @@ -42,11 +42,11 @@ internal object UniFfiForeignExecutorCallback : com.sun.jna.Callback {
}
}

public object FfiConverterForeignExecutor: FfiConverter<CoroutineScope, USize> {
internal val handleMap = UniFfiHandleMap<CoroutineScope>()
public object FfiConverterForeignExecutor: FfiConverter<CoroutineScope, UniffiHandle> {
internal val handleMap = UniffiHandleMap<CoroutineScope>()

internal fun drop(handle: USize) {
handleMap.remove(handle)
internal fun drop(handle: UniffiHandle) {
handleMap.consumeHandle(handle)
}

internal fun register(lib: _UniFFILib) {
Expand All @@ -58,26 +58,21 @@ public object FfiConverterForeignExecutor: FfiConverter<CoroutineScope, USize> {
{% endmatch %}
}
// Number of live handles, exposed so we can test the memory management
public fun handleCount() : Int {
return handleMap.size
}
override fun allocationSize(value: CoroutineScope) = USize.size
override fun allocationSize(value: CoroutineScope) = 8
override fun lift(value: USize): CoroutineScope {
return handleMap.get(value) ?: throw RuntimeException("unknown handle in FfiConverterForeignExecutor.lift")
override fun lift(value: UniffiHandle): CoroutineScope {
return handleMap.get(value)
}
override fun read(buf: ByteBuffer): CoroutineScope {
return lift(USize.readFromBuffer(buf))
return lift(buf.getLong())
}
override fun lower(value: CoroutineScope): USize {
return handleMap.insert(value)
override fun lower(value: CoroutineScope): UniffiHandle {
return handleMap.newHandle(value)
}
override fun write(value: CoroutineScope, buf: ByteBuffer) {
lower(value).writeToBuffer(buf)
buf.putLong(lower(value))
}
}
53 changes: 53 additions & 0 deletions uniffi_bindgen/src/bindings/kotlin/templates/HandleMap.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
internal class UniffiHandleMap<T> {
private val lock = ReentrantReadWriteLock()
private var mapId: Long = UniffiHandleMap.nextMapId()
private val map: MutableMap<Long, T> = mutableMapOf()
// Note: Foreign handles are always odd
private var keyCounter = 1L

private fun nextKey(): Long = keyCounter.also {
keyCounter = (keyCounter + 2L).and(0xFFFF_FFFF_FFFFL)
}

private fun makeHandle(key: Long): UniffiHandle = key.or(mapId)

private fun key(handle: UniffiHandle): Long {
if (handle.and(0x7FFF_0000_0000_0000L) != mapId) {
throw InternalException("Handle map ID mismatch")
}
return handle.and(0xFFFF_FFFF_FFFFL)
}

fun newHandle(obj: T): UniffiHandle = lock.writeLock().withLock {
val key = nextKey()
map[key] = obj
makeHandle(key)
}

fun get(handle: UniffiHandle) = lock.readLock().withLock {
map[key(handle)] ?: throw InternalException("Missing key in handlemap: was the handle used after being freed?")
}

fun cloneHandle(handle: UniffiHandle): UniffiHandle = lock.writeLock().withLock {
val obj = map[key(handle)] ?: throw InternalException("Missing key in handlemap: was the handle used after being freed?")
val clone = nextKey()
map[clone] = obj
makeHandle(clone)
}

fun consumeHandle(handle: UniffiHandle): T = lock.writeLock().withLock {
map.remove(key(handle)) ?: throw InternalException("Missing key in handlemap: was the handle used after being freed?")
}

companion object {
// Generate map IDs that are likely to be unique
private var mapIdCounter: Long = {{ ci.namespace_hash() }}.and(0x7FFF)

// Map ID, shifted into the top 16 bits
internal fun nextMapId(): Long = mapIdCounter.shl(48).also {
// On Kotlin, map ids are only 15 bits to get around signed/unsigned issues
mapIdCounter = ((mapIdCounter + 1).and(0x7FFF))
}
}
}

84 changes: 1 addition & 83 deletions uniffi_bindgen/src/bindings/kotlin/templates/Helpers.kt
Original file line number Diff line number Diff line change
Expand Up @@ -78,89 +78,7 @@ private inline fun <U> rustCall(callback: (RustCallStatus) -> U): U {
return rustCallWithError(NullCallStatusErrorHandler, callback);
}

// IntegerType that matches Rust's `usize` / C's `size_t`
public class USize(value: Long = 0) : IntegerType(Native.SIZE_T_SIZE, value, true) {
// This is needed to fill in the gaps of IntegerType's implementation of Number for Kotlin.
override fun toByte() = toInt().toByte()
// Needed until https://youtrack.jetbrains.com/issue/KT-47902 is fixed.
@Deprecated("`toInt().toChar()` is deprecated")
override fun toChar() = toInt().toChar()
override fun toShort() = toInt().toShort()

fun writeToBuffer(buf: ByteBuffer) {
// Make sure we always write usize integers using native byte-order, since they may be
// casted to pointer values
buf.order(ByteOrder.nativeOrder())
try {
when (Native.SIZE_T_SIZE) {
4 -> buf.putInt(toInt())
8 -> buf.putLong(toLong())
else -> throw RuntimeException("Invalid SIZE_T_SIZE: ${Native.SIZE_T_SIZE}")
}
} finally {
buf.order(ByteOrder.BIG_ENDIAN)
}
}

companion object {
val size: Int
get() = Native.SIZE_T_SIZE

fun readFromBuffer(buf: ByteBuffer) : USize {
// Make sure we always read usize integers using native byte-order, since they may be
// casted from pointer values
buf.order(ByteOrder.nativeOrder())
try {
return when (Native.SIZE_T_SIZE) {
4 -> USize(buf.getInt().toLong())
8 -> USize(buf.getLong())
else -> throw RuntimeException("Invalid SIZE_T_SIZE: ${Native.SIZE_T_SIZE}")
}
} finally {
buf.order(ByteOrder.BIG_ENDIAN)
}
}
}
}


// Map handles to objects
//
// This is used when the Rust code expects an opaque pointer to represent some foreign object.
// Normally we would pass a pointer to the object, but JNA doesn't support getting a pointer from an
// object reference , nor does it support leaking a reference to Rust.
//
// Instead, this class maps USize values to objects so that we can pass a pointer-sized type to
// Rust when it needs an opaque pointer.
//
// TODO: refactor callbacks to use this class
internal class UniFfiHandleMap<T: Any> {
private val map = ConcurrentHashMap<USize, T>()
// Use AtomicInteger for our counter, since we may be on a 32-bit system. 4 billion possible
// values seems like enough. If somehow we generate 4 billion handles, then this will wrap
// around back to zero and we can assume the first handle generated will have been dropped by
// then.
private val counter = java.util.concurrent.atomic.AtomicInteger(0)

val size: Int
get() = map.size

fun insert(obj: T): USize {
val handle = USize(counter.getAndAdd(1).toLong())
map.put(handle, obj)
return handle
}

fun get(handle: USize): T? {
return map.get(handle)
}

fun remove(handle: USize): T? {
return map.remove(handle)
}
}

// FFI type for Rust future continuations
internal interface UniFffiRustFutureContinuationCallbackType : com.sun.jna.Callback {
fun callback(continuationHandle: USize, pollResult: Short);
fun callback(continuationHandle: Long, pollResult: Short);
}
13 changes: 12 additions & 1 deletion uniffi_bindgen/src/bindings/kotlin/templates/ObjectRuntime.kt
Original file line number Diff line number Diff line change
@@ -1,5 +1,16 @@
{{- self.add_import("java.util.concurrent.atomic.AtomicLong") }}
{{- self.add_import("java.util.concurrent.atomic.AtomicBoolean") }}

// Wraps `UniffiHandle` to pass to object constructors.
//
// This class exists because `UniffiHandle` is a typealias to `Long`. If the object constructor
// inputs `UniffiHandle` directly and the user defines a primary constructor than inputs a single
// `Long` or `ULong` input, then we get JVM signature conflicts. To avoid this, we pass this type
// in instead.
//
// Let's try to remove this when we update the code based on ADR-0008.
data class UniffiHandleWrapper(val handle: UniffiHandle)

// The base class for all UniFFI Object types.
//
// This class provides core operations for working with the Rust handle to the live Rust struct on
Expand Down Expand Up @@ -100,7 +111,7 @@ abstract class FFIObject: Disposable, AutoCloseable {
this.handle = null
}

protected val handle: UniffiHandle?
internal val handle: UniffiHandle?

private val wasDestroyed = AtomicBoolean(false)
private val callCounter = AtomicLong(1)
Expand Down
Loading

0 comments on commit bc5e827

Please sign in to comment.