Synchronization between threads

[p-lenart]

I wanted to create a very small PoC that showed case offloading CPU-intensive calls to the background thread. I cannot confirm it, but I have a suspicion that Swift's stdlib implementation for WASI is not actually thread-safe, and I cannot synchronize code in an expected way. I attached the code at the bottom. I can see various crashes around accessing invalid memory, locking, etc. Here are some examples:

Thread 1Thread 2
Thread 1Thread 2
Error: error while executing at wasm backtrace:
    0: 0xef40ca - MT.wasm!__wasi_fd_write
    1: 0xf00db4 - MT.wasm!writev
    2: 0xf00eeb - MT.wasm!__stdio_write
    ...
   12:  0x3da41 - MT.wasm!$s2MT13spawnBlockingyxxyYbcYaKs8SendableRzlFyyYbcfU_TA
   13:  0x400ca - MT.wasm!$s2MT18_wasi_thread_startyys5Int32V_SvSgtF
   14:  0x4000b - MT.wasm!wasi_thread_start

Thread 1
Error: error while executing at wasm backtrace:
    0:  0x414a6 - MT.wasm!$sypWOc
    1: 0x1dde66 - MT.wasm!$ss6_print_9separator10terminator2toySayypG_S2Sxzts16TextOutputStreamRzlFs7_StdoutV_Tg5
    2: 0x1deb20 - MT.wasm!$ss5print_9separator10terminatoryypd_S2StFTm
    ...
    7:  0x400ca - MT.wasm!$s2MT18_wasi_thread_startyys5Int32V_SvSgtF
    8:  0x4000b - MT.wasm!wasi_thread_start

Here is my code:

import WASILibc
import Foundation

// Simple runtime for spawning blocking operation on separate thread.

enum ThreadError: Error {
    case threadSpawnFailed
}

class BoxedValue<T: Sendable>: @unchecked Sendable {
    private let lock = NSLock()
    private var value: T?

    func put(_ newValue: T) {
        lock.lock()
        value = newValue
        lock.unlock()
    }

    func take() -> T? {
        lock.lock()
        let currentValue = value
        value = nil
        lock.unlock()
        return currentValue
    }

    deinit {
        print("BoxedValue deinitialized")
    }
}

class ThreadContext: @unchecked Sendable {
    let operation: @Sendable () -> Void
    init(operation: @Sendable @escaping ()-> Void ) {
        self.operation = operation
    }

    deinit {
        print("ThreadContext deinitialized")
    }
}

public func spawnBlocking<T: Sendable>(_ operation: @Sendable @escaping () -> T) async throws -> T {
    let returnValue = BoxedValue<T>()
    let context = ThreadContext(operation: {
        let result = operation()
        returnValue.put(result)
    })
    let unmanagedContext = Unmanaged<ThreadContext>.passRetained(context)
    let result = WASILibc.__wasi_thread_spawn(unmanagedContext.toOpaque())
    guard result >= 0 else {
        throw ThreadError.threadSpawnFailed
    }
    // Tried implementation based on unchecked continuation and async stream, but we
    // cannot call theses APIs from the WASM thread, unfortunately.
    while true {
        if let value = returnValue.take() {
            return value
        }
        try await Task.sleep(nanoseconds: 10_000_000)
    }
}

@_expose(wasm, "wasi_thread_start")
@_cdecl("wasi_thread_start")
public func _wasi_thread_start(_ threadId: Int32, _ arg: UnsafeMutableRawPointer?) {
    guard let arg = arg else {
        fatalError("Thread started with null argument")
    }
    let unmanagedContext = Unmanaged<ThreadContext>.fromOpaque(arg)
    unmanagedContext.takeRetainedValue().operation()
}


// Example

print("Hello from main thread")

let task1 = Task {
    try await spawnBlocking {
        print("Starting counter 1 thread")
        while true {
            print("Thread 1")
            usleep(100_000)
        }
        print("Finished counter 1 thread")
    }
}

let task2 = Task {
    try await spawnBlocking {
        print("Starting counter 2 thread")
        while true {
            print("Thread 2")
            usleep(100_000)
        }
        print("Finished counter 2 thread")

    }
}

try await task1.value
try await task2.value

Swift SDK version that I use: 6.1-RELEASE-wasm32-unknown-wasip1-threads
Compiled with: swift build --swift-sdk wasm32-unknown-wasip1-threads
Run with: wasmtime --wasi threads .build/debug/MT.wasm

[p-lenart]

Actually, I wasn't aware that we had access to the pthread implementation under the hood and using it seems to solve most of the issues that I encountered.

import wasi_pthread

// Simple runtime for spawning blocking operation on separate thread.

enum ThreadError: Error {
    case threadSpawnFailed
}

class ThreadContext: @unchecked Sendable {
    let operation: @Sendable () -> Void
    init(operation: @Sendable @escaping ()-> Void ) {
        self.operation = operation
    }

    deinit {
        print("ThreadContext deinitialized")
    }
}

public func spawnBlocking<T: Sendable>(_ operation: @Sendable @escaping () -> T) async throws -> T {
    // NOTE: This ensures that Swift runtime will pickup value from the thread. For
    // some reason main thread is not getting notified when continuation is resumed.
    let task = Task { while true { try await Task.sleep(nanoseconds: 10_000_000) } }
    defer { task.cancel() }

    return try await withCheckedThrowingContinuation { continuation in
        let context = ThreadContext(operation: {
            let result = operation()
            continuation.resume(returning: result)
        })
        let unmanagedContext = Unmanaged<ThreadContext>.passRetained(context)
        var thread = pthread_t(bitPattern: 0)
        let result = wasi_pthread.pthread_create(
            &thread,
            nil,
            { data in
                guard let data = data else {
                    return nil
                }
                let unmanagedContext = Unmanaged<ThreadContext>.fromOpaque(data)
                unmanagedContext.takeRetainedValue().operation()
                return nil
            },
            unmanagedContext.toOpaque()
        )
        guard result == 0 else {
            continuation.resume(throwing: ThreadError.threadSpawnFailed)
            return
        }
    }
}

// Spawn main task...
let mainTask = Task {
    print("Hello from main thread")
    while true {
        try await Task.sleep(nanoseconds: 150_000_000)
        print("Waiting for threads to finish...")
    }
}
defer {
    mainTask.cancel()
    print("Program finished")
}

// Spawn two blocking tasks...
try await withThrowingTaskGroup { group in
    defer { group.cancelAll() }

    group.addTask {
        try await spawnBlocking {
            print("Starting 1 thread")
            for _ in 0...10 {
                usleep(UInt32(100_000))
                print("Doing work on 1 thread...")
            }
            print("Finished 1 thread")
            return 1
        }
    }

    group.addTask {
        try await spawnBlocking {
            print("Starting 2 thread")
            for _ in 0...10 {
                usleep(UInt32(100_000))
                print("Doing work on 2 thread...")
            }
            print("Finished 2 thread")
            return 2
        }
    }

    for try await value in group {
        print("Got value: \(value)")
    }
}

My original question still holds. Is it safe to assume that Swift's standard library is thread-safe? I'm thinking primarily about allocation, locking, etc.

[kateinoigakukun]

Swift stdlib and runtime code are already thread-safe as long as you use pthread_create. (implementing wasi_thread_start by yourself is totally unsafe, we must use the one from wasi-libc). One missing thread unsafety is Foundation, I just started working on it: swiftlang/swift-foundation#1490

[p-lenart]

Thank you, that's what I wanted to know. I'll close this thread as my original question was answered. I really appreciate your effort in this area. Thanks again!