What context.{get,set} slots are used when an async task returns a string?

[alexcrichton]

I've been thinking about bytecodealliance/wasmtime#13890 and how best to fix it recently. I've come across behavior though where I'm not entirely sure what to do, hence the question here. The basic problem of this Wasmtime issue is that when a thread exits there remains lingering references to the thread's ID and data structure which causes problems when they're later accessed. My original goal was to continue to "promptly deallocate" thread structures when a thread exits (as opposed to keeping them alive until subtasks have returned or something like that), but this led to a bit of a deeper question.

Specifically there's currently the scenario where a host async function returns in a "delayed" fashion. For example it's invoked, not ready yet, and then a subtask is created within the guest. Eventually when this host subtask actually generates its result it lowers the result into guest memory and then signals an event saying that the subtask is ready. The process of lowering, however, can invoke guest functions like realloc and need context.{get,set} intrinsics to work correctly. At this time I don't actually know what the right values are to use here and I think Wasmtime is relatively buggy.

A rough example of what I'm thinking of is this test:

Details

(component
  (core module $libc
    (import "" "context.set" (func $context.get (result i32)))
    (memory (export "m") 1)
    (func (export "realloc") (param i32 i32 i32 i32) (result i32)
      (i32.store
        (i32.const 100)
        (call $context.get))
      i32.const 0
    )
  )
  (component $A
    (core func $context.get (canon context.get i32 0))
    (core instance $libc (instantiate $libc
      (with "" (instance
        (export "context.set" (func $context.get))
      ))
    ))

    (core module $m
      (import "libc" "m" (memory 1))
      (import "" "task.return" (func $task.return (param i32 i32)))

      (global $go (mut i32) (i32.const 0))

      (func (export "run") (result i32)
        i32.const 1 ;; CALLBACK_CODE_YIELD
      )

      (func (export "run-cb") (param i32 i32 i32) (result i32)
        (if (result i32) (global.get $go)
          (then
            (call $task.return (i32.const 0) (i32.const 0))
            i32.const 0 ;; CALLBACK_CODE_EXIT
          )
          (else (i32.const 1)) ;; CALLBACK_CODE_YIELD
        )
      )

      (func (export "go") (global.set $go (i32.const 1)))
    )

    (core func $task.return (canon task.return (result string) (memory $libc "m")))
    (core instance $m (instantiate $m
      (with "libc" (instance $libc))
      (with "" (instance
        (export "task.return" (func $task.return))
      ))
    ))

    (func (export "run") async (result string)
      (canon lift (core func $m "run") async (memory $libc "m")
          (callback (func $m "run-cb")))
    )
    (func (export "go") (canon lift (core func $m "go") ))
  )

  (component $B
    (import "a" (instance $a
      (export "run" (func async (result string)))
      (export "go" (func))
    ))

    (core func $context.get (canon context.get i32 0))
    (core instance $libc (instantiate $libc
      (with "" (instance
        (export "context.set" (func $context.get))
      ))
    ))

    (core module $m
      (import "" "run" (func $run (param i32) (result i32)))
      (import "" "go" (func $go))
      (import "" "thread.yield" (func $thread.yield (result i32)))
      (import "" "context.set" (func $context.set (param i32)))
      (import "libc" "m" (memory 1))

      (global $subtask (mut i32) (i32.const 0))

      (func (export "start")
        (local $r i32)

        (call $context.set (i32.const 1))


        (local.set $r (call $run (i32.const 0)))
        (if (i32.ne
              (i32.and (local.get $r) (i32.const 0xf))
              (i32.const 1)) ;; STARTED
          (then (unreachable)))
        (global.set $subtask (i32.shr_u (local.get $r) (i32.const 4)))
      )

      (func (export "finish") (result i32)
        (call $context.set (i32.const 2))

        call $go
        (drop (call $thread.yield))
        (i32.load (i32.const 100))
      )
    )

    (core func $run
      (canon lower
        (func $a "run")
        async
        (memory $libc "m")
        (realloc (func $libc "realloc"))
      )
    )
    (core func $go (canon lower (func $a "go")))
    (core func $thread.yield (canon thread.yield))
    (core func $context.set (canon context.set i32 0))

    (core instance $i (instantiate $m
      (with "libc" (instance $libc))
      (with "" (instance
        (export "run" (func $run))
        (export "go" (func $go))
        (export "thread.yield" (func $thread.yield))
        (export "context.set" (func $context.set))
      ))
    ))

    (func (export "start") (canon lift (core func $i "start")))
    (func (export "finish") async (result u32) (canon lift (core func $i "finish")))
  )

  (instance $a (instantiate $A))
  (instance $b (instantiate $B (with "a" (instance $a))))
  (export "start" (func $b "start"))
  (export "finish" (func $b "finish"))
)

(assert_return (invoke "start"))
(assert_return (invoke "finish") (u32.const 0))

which runs as:

$ wasmtime wast foo.wast
Error: failed to run script file 'foo.wast'

Caused by:
    0: failed directive on foo.wast:136
    1: result 0 didn't match
    2: expected                  0 / 0x0000000000000000
actual                    2 / 0x0000000000000002

The basic premise of this script is that component $B has an export "start" which starts a subtask in $A. This subtask returns a string and isn't immediately ready yet. Next a new task in $B, the "finish" function, flags $A as "please complete now" which then leads to the completion of $A's original task. This invokes realloc where the definition here stores context.get in a known location in memory. The $B "finish" subtask then completes with what was returned by realloc at that location in memory. Currently Wasmtime runs realloc in the context of the "finish" task in $B, but I'm not sure if this is correct or the desired outcome here.

To me the "start" subtask seems wrong since it's exited. The "finish" subtask also seems wrong because it hasn't even said anything about $A's subtask yet. Should this be something where a fresh, temporary, task is created when $A returns to $B? This fresh, temporary, task is basically scoped to realloc invocations and is disposed of quickly afterwards.

This is where I got lost enough myself in the "right" answer here, hence the issue here.