compiler: perform BuiltinFunctionCheckExitCode natively

[mathetake]

wazeroir's BuiltinFunctionCheckExitCode operation has been introduced in the context cancelation integration in #1108. Basically, this operation is supposed to perform the check on ModuleInstance.Closed value, and exit the execution if the value is not zero.

Currently, the compiler's implementation is pretty naive like simply exiting the native world and call ModuleInstance.FailIfClosed.

wazero/internal/engine/compiler/impl_amd64.go

Lines 271 to 272 in 86444c6

	func (c *amd64Compiler) compileBuiltinFunctionCheckExitCode() error {
	if err := c.compileCallBuiltinFunction(builtinFunctionIndexCheckExitCode); err != nil {

wazero/internal/engine/compiler/engine.go

Lines 1012 to 1016 in 86444c6

	case builtinFunctionIndexCheckExitCode:
	// Note: this operation must be done in Go, not native code. The reason is that
	// native code cannot be preempted and that means it can block forever if there are not
	// enough OS threads (which we don't have control over).
	if err := m.FailIfClosed(); err != nil {

but this comes with the cost - BuiltinFunctionCheckExitCode is inserted at all of the loop instructions and each time we call a Go function. FailIfClosed function checks the exit code atomically, but it is not necessary as in anyway the implementation is not synchronous (e.g. we don't check the code until the Wasm VM reaches loop instruction even when the context reaches timeout 10 years ago in extreme case).

Given that assumption we do not need to check it synchronously, we could optimize the translation of BuiltinFunctionCheckExitCode by simply checking the exit code natively.

The proposed change should come with

• Native check of ExitCode in compiler.compileBuiltinFunctionCheckExitCode function for both arm64 and amd64
  • native code has access to it via callEngine.moduleInstance
• The benchmark on the cost of context cancelation integration

[mathetake]

@evacchi could you take on after #1406

[evacchi]

Nice!👍

[evacchi]

If I understand correctly this boils down to c.compileExitFromNativeCode(exitCode) when exitCode is nonzero; where exitCode is extracted from (*moduleInstance).Closed (closed >> 32); Closed is at the address callEngineModuleContextModuleInstanceOffset because it is at offset 0 of the struct

EDIT: well not exactly that because c.compileExitFromNativeCode(status) wants a value, while we are going to have a register instead

[mathetake]

no, you misunderstood compileExitFromNativeCode(exitCode). That is compile time function and it accepts nativeCallStatusCode which is different from "exitCode" stored at ModuleInstance.Closed. See the existing code base on how it works. You don't need to "pass" the value on ModuleInstance.Closed, just let it exit the execution regardless of the value.

[mathetake]

ps the new benchmark's case should look like for i := 0; i < 10000; i++ {} in wat where basically execute the many loop in-Wasm so that we can see the cost of each loop header without being offset by the other noises

[mathetake]

Ok, I just realized that this is in theory impossible - if we do this check completely in "native world", then Goroutine switch doesn't/may not happen and therefore the ExitCode might not get updated forever (See compiler/RATIONALE.md).

That's why @evacchi is experiencing the forever-hung whatever change you make when t.Parallel was enabled in the tests in #1413.

With that said, I am closing this as not planned. Thanks @evacchi for your investigation and try!