microsoft / vscode-wasm

A WASI implementation that uses VS Code's extension host as the implementing API
MIT License
376 stars 27 forks source link

Command callback not receiving WASM returned data when in Worker #201

Open AFCMS opened 2 days ago

AFCMS commented 2 days ago

I am exploring the use of a Rust file treatment library straight inside VSCode in WebAssembly.

I successfully built my fist POC, using the wasm32-unknown-unknown target + wit2ts. It exports a single function which takes file data (Uint8Array) as a parameter and return the modified file in the same format (it's a trivial function on my side since the library have already a Uint8Array->Uint8Array function). So the extension command handler read the file (with vscode.workspace.fs.readFile), call the WASM function and write it back. (not sure if it's the best thing to do, but reading/writing file from WASI directly for just a single function which access one file at a time seemed a bit too much for me)

This setup worked really well, but according to the second blog post long running tasks should be put in a worker. In my case the treatment can go from half a second to more than 10 seconds, so it should ideally be in a worker I guess.

I followed the exemple to make the wasm part run in a worker, and I went into a strange problem.

WIT file:

package afcms:test-extension;

world test-extension {
    import log: func(msg: string);
    export add: func(left: u64, right: u64) -> u64;
    export optimise: func(data: list<u8>) -> list<u8>;
}

Rust code (minus boilerplate):

fn add(left: u64, right: u64) -> u64 {
    log("Calculating...");
    return left + right;
}

fn optimise(in_data: Vec<u8>) -> Vec<u8> {
    log("Starting optimising...");
    let new_data = do_stuff(&in_data);
    log("Finished optimising");
    return new_data.unwrap();
}

Extension code:

async function loadWasmModule(context: vscode.ExtensionContext) {
    const filename = vscode.Uri.joinPath(
        context.extensionUri,
        "target",
        "wasm32-unknown-unknown",
        "debug", // TODO: Change to release
        "test-extension.wasm"
    );

    // The channel for printing the log.
    const log = vscode.window.createOutputChannel("Test Extension - Worker", { log: true });
    context.subscriptions.push(log);

    const bits = await vscode.workspace.fs.readFile(filename);
    const module = await WebAssembly.compile(bits);

    const worker = new Worker(vscode.Uri.joinPath(context.extensionUri, "./dist/worker.js").fsPath);

    const service: testExtension.Imports.Promisified = {
        log: (msg: string) => {
            log.info(msg);
        },
    };

    // Bind the TypeScript Api
    const api = await testExtension._.bind(service, module, worker);

    return api;
}

export async function activate(context: vscode.ExtensionContext) {
    const api = await loadWasmModule(context);

    const disposable = vscode.commands.registerCommand("testExtension.helloWorld", async () => {
        const n = await api.add(BigInt(2), BigInt(2));
        vscode.window.showInformationMessage("Hello World! 2+2=" + n.toString());
    });

    const commandOptimise = vscode.commands.registerCommand("testExtension.optimise", async (param: vscode.Uri) => {
        vscode.window.showInformationMessage("Called: " + param.fsPath);
        const in_data = await vscode.workspace.fs.readFile(param);
        const out_data = await api.optimise(in_data);
        vscode.window.showInformationMessage("Optimised: " + param.fsPath);
        vscode.workspace.fs.writeFile(param, out_data);
    });

    context.subscriptions.push(disposable, commandOptimise);
}

export function deactivate() {}

I didn't change anything on the Rust/WIT side and the testExtension.helloWorld command works fine in both cases but now when calling the testExtension.optimise command, the first info message is displayed, the two log messages from the worker are displayed in the right output channel, but the second info message is never displayed and neither is the file modified on disk.

What I don't understand is why it was working without the worker, and now it seems the rest of the command callback is just never called, despite the wasm function apparently succeeding, without any extension error message or anything in logs.

I am still a beginner in WASM, maybe it's just a stupid error but I really can't figure where to search at this point.

dbaeumer commented 1 day ago

@AFCMS could you share the repository with me. Makes it a lot easier to understand what is going on.

AFCMS commented 1 day ago

Didn't have a repository yet, created one.

https://github.com/AFCMS/vscode-oxipng

It aims to provide the oxipng right inside VSCode.

There are two VSCode commands, the hello world one (which works) and the optimise one (which no longer works). For the optimise one, it's accessible in the file explorer and it should work on PNG files (small files can be used for testing since bigger ones tends to take a long time to process).