TLS in `wasm32-wasip3` is broken in shared libraries

[alexcrichton]

In debugging with @dicej today we've realized that TLS in wasi-libc is broken on
the wasm32-wasip3 target with shared libraries. This is not surfaced in
testing at all today due to #813 where we don't test shared libraries at all.
The problems we've identified are fundamental enough such that it won't be
possible to fix them in just wasi-libc, but after discussion we think that we
can solve everything in the context of wit-component and wasm-tools -- no
LLVM changes needed.

The major problems we've identified today are:

• The codegen of LLVM/wasm-ld requires that each library gets a unique TLS
  base pointer. Currently in wasi-libc, however, the same TLS base pointer
  (context.get 1) is used for all libraries. This means that all libraries
  stomp over each other with their TLS blocks.

• Spawned threads have no means of initializing the TLS of other shared
  libraries. Right now they initialize the TLS of the shared object that libc
  itself is linked into, but this doesn't include the main program nor other
  shared libraries.

• Spawned threads do not allocate space for TLS in other shared libraries. The
  __builtin_wasm_tls_{size,align} functions which wasi-libc uses report the
  size/align for libc itself, but this doesn't include the alignment/size for
  any other library.

With that outlined, I'd next like to describe toolchain/convention/etc changes
that Joel and I are thinking of to fix this in wasi-libc. In designing this
solution we primarily considered the 2x2 matrix of {no coop threads, coop
threads} x {single-module, multi-module}. Here wasi-libc can be built either
with or without coop threads (the first axis), and when linking into a component
there is either one big statically linked module (single-module) or many dynamic
libraries linked together (multi-module). Each of these use cases has slightly
different constraints which helped inform the exact design we've landed on here.

One note is that below I'm going to refer to wit-component which is the
underlying implementation of creating a component in wasm-tools. This applies
to wasm-component-ld, the default linker to create a component coming out of
LLVM, as well as wasm-tools component link which is how dynamic libraries are
linked together.

One final note is that everything here only applies to the wasip3 target. Other
targets will continue to deal with TLS as they do today which works ok for all
of these use cases.

Changing __wasm_{get,set}_tls_base definition location

Right now the logic/knowledge of how the TLS base is set/stored is defined
exclusively within wasi-libc itself. The plan is to share this with
wit-component and change the responsibility of defining
__wasm_{get,set}_tls_base to wit-component itself. This means that
wasi-libc, and all other shared objects, will import
__wasm_{get,set}_tls_base regardless of whether coop threads are enabled or
not. This is enabled for coop threads already through #856 and the thinking is
that we'll adjust the non-coop-threads build to do the same thing.

The main rationale for this is that it's basically just going to be easier to
think about this at the wit-component layer than in wasi-libc. While it might
be possible to have various #defines to get things lined up just right it's
predicted to be conceptually easiest to say "this is a linker thing". How
exactly these are defined is below.

Changing the conceptual model of "tls base"

One of the main problems being solved here is to align the assumption in LLVM,
that every module gets a unique TLS base pointer, with the implementation in
wasi-libc. This is not trivially mapped to the component model's context.get 1
slot as originally anticipated because that only works for a single module. An
additional constraint that we have, without changing LLVM, is that every
module's main-thread-TLS is placed somewhere in linear memory when a module is
instantiated. These locations are disjoint for each library loaded, meaning it's
not possible to assume a contiguous set of TLS for all libraries.

Given that the TLS base definition is moving to wit-component, here's what
Joel and I were thinking:

• For the single-module use case this will largely remain the same as it is
  today. The context.get 1 slot of tasks, for example, is the base pointer of
  TLS for the single module and everything works as-is.

• For the multi-module use case, however, the context.get 1 slot will now
  instead be a pointer to an array of pointers in wasm linear memory. Each
  thread will get its own array of pointers, and this means that the TLS base
  for a particular module is effectively defined with another level of
  indirection here.

The exact definition and initialization of this array-of-pointers and how
exactly this gets hooked up for single-threaded use cases is described below.

Changes in wit-component

Ok, with that background out of the way, here's what wit-component will be
gaining:

• When linking it'll be detected whether thread.new-indirect is used anywhere.
  This is used to determine whether it's possible that a module could spawn a
  thread and thus might be using coop threads. This is a bit of a hack for now
  to avoid using context.get 1 when coop threads are disabled. The main
  rationale for this is that context.get 1 is considered invalid without the
  -W component-model-threading feature to Wasmtime, and the wasm32-wasip3
  target should work without flags in Wasmtime. Eventually this detection can go
  away once -W component-model-threading is on-by-default.

• For single-module use cases wit-component will synthesize functions to
  satisfy the storage of the TLS base pointer. Note that an array-of-pointers is
  not used here as it is not necessary as there's just a single module.

  • If the module is using coop threading (imports thread.new-indirect), then
    __wasm_{get,set}_tls_base are hooked directly up to context.{get,set} 1
    exactly as happens today (after Route __wasm_get_stack_pointer as an import #856).
  • If the module is not using coop threading then new functions will be
    synthesized that get/set a global. This will be included in one of
    wit-component's generated modules it already has, such as one of the shim
    modules. It's not expected that this'll be too hard to add, and this also
    exactly matches what wasi-libc does today for non-coop-threads use cases.

• For multi-module use cases the implementation is more involved as this is what
  we're actually trying to solve here. With multi-module use cases, however,
  wit-component can inject new regions of linear memory and generally allocate
  data in linear memory (everything is position independent). Here
  wit-component will synthesize a function-per-module for
  __wasm_{get,set}_tls_base imports because each module needs a unique base to
  operate on. The generated functions will look like:

  • If any module is using coop threading, then generated functions will use
    context.get 1 as a base pointer and then load/store at a static offset
    from that base pointer. It'll be assumed that wasi-libc itself is managing
    context.get 1 as part of the ABI, and wit-component is just synthesizing
    functions that load/store from this location. One small caveat here will be
    that __wasm_set_tls_base will conditionally test if context.get 1 is 0.
    If it's null then it'll store into the main_thread_tls_base array
    described below, otherwise it operates on the pointer in context.get 1.
    This caveat is needed to handle start functions which otherwise don't have
    context.get 1 configured and no special handling is needed for
    __wasm_get_tls_base because that's only ever called after
    __wasm_init_task is called.
  • If no module is using coop threading, then wit-component will allocate a
    static section of memory for this array-of-pointers. The synthesized
    functions will then load/store the addresses provided in memory at this
    static offset. While a fixed set of globals could be used similar to
    no-coop-threads above it's predicted to be easier to do this in-memory the
    same way as with-coop-threads above.

• Additionally for multi-module use cases wit-component will synthesize an
  in-memory data structure for wasi-libc to reflect on. The thinking is that
  this will look and behave similarly to how dlopen does today where
  wit-component creates a data structure that wasi-libc uses at runtime. In
  this case wit-component will create a data structure that looks like:

  struct {
      size_t num_libraries;
      struct {
          size_t __tls_size;
          size_t __tls_align;
          void (*__wasm_init_tls)(void*);
      } *library_info;
      void **main_thread_tls_base;
  } __wasm_program_tls_info;

  This will require changes in wasi-libc to export the __tls_{size,align}
  globals as well as the __wasm_init_tls function, and it's expected this'll
  be easy enough through the crt1-*.o objects. Using this information
  wit-component will be able to synthesize this in-memory data structure for
  wasi-libc to use, and this will primarily be used when spawning threads.

These changes alone aren't able to fully solve the problem, however, so this
will require more coordination from wasi-libc's side of things as well.

Changes in wasi-libc

In addition to the previous mention of removing __wasm_{get,set}_tls_base
definitions from wasi-libc, mostly when coop threads are disabled, there are
other various updates necessary to ensure that everything works. Specifically:

• The __wasm_init_task and __wasm_init_async_task hooks need to be updated
  in how they initialize TLS. TLS initialization will be deferred to a C
  function in wasi-libc instead of being pure inline assembly as-is today, and
  this function will have a weak import of the __wasm_program_tls_info symbol
  synthesized by wit-component. If this symbol is NULL then wasi-libc knows
  that __builtin_wasm_tls_{size,align} accurately reflect TLS information.
  Additionally this indicates whether context.set 1 is the raw TLS base or the
  indirect list-of-pointers for all modules.
  • If coop threads are disabled, nothing happens. In this situation the
    get/set TLS base functions statically bake in the address within
    wit-component, so no action is needed.
  • If coop threads are enabled then context.set 1 needs to be configured, and
    the exact value depends on __wasm_program_tls_info:
    • When __wasm_program_tls_info is NULL, then context.set 1 is called
      with __init_tls_base's value.
    • When __wasm_program_tls_info is not NULL, context.set 1 is set to
      __wasm_program_tls_info.main_thread_tls_base.
• When a thread is created wasi-libc will also use the presence or absence of
  __wasm_program_tls_info to determine how exactly to allocate TLS.
  • When __wasm_program_tls_info is NULL, then wasi-libc will allocate TLS
    as it does today using __builtin_wasm_tls_{size,align} and the local
    definition of __wasm_init_tls.
  • When __wasm_program_tls_info is not NULL, then wasi-libc will use the
    listed information to allocate an appropriately-sized/aligned region of
    memory to store all TLS information itself as well as the array-of-pointers.
    The array-of-pointers is initialized by libc itself and then each library's
    own __wasm_init_tls callback function is invoked to initialize its own
    block of TLS.

The majority of changes here are in pthread_create and thread initialization,
but there are also minor changes to __wasm_init_task to handle a few more
cases and management of context.set 1 as necessary.

How each case works

Originally I described a 2x2 matrix of {coop threads off,on} x {single,multi
module}. Here I want to outline how these cases all work with the changes
outlined above:

No coop threads & Single module

In this case __wasm_init_task doesn't do anything related to TLS. If TLS is
used the accessors will get synthesized by wit-component and will frob a
global in a generated function/module. Thread creation isn't supported here
with coop threads disabled, and __wasm_program_tls_info is not defined.

No coop threads & Multi module

In this case __wasm_init_task doesn't do anything related to TLS. While
__wasm_program_tls_info could be generated and used, there will be no usage of
this in wasi-libc so it won't end up getting generated in the end. In
wit-component though a separate get/set TLS base for each module will be
generated which will frob a location in memory injected by wit-component as
well. Thread creation isn't supported here with coop threads disabled.

Coop threads & Single module

In this case during module instantiation the main thread's initial TLS in linear
memory is initialized, the base pointer is put in __init_tls_base, and
__wasm_set_tls_base is called. Here __wasm_{get,set}_tls_base is directly
hooked up to context.{get,set} 1. The __wasm_init_task function will see
that __wasm_program_tls_info is NULL and context.set 1 will be invoked
with the value of __init_tls_base.

When a new thread is spawned the __wasm_program_tls_info symbol is NULL and
so thread creation will proceed as it does today. The program's TLS is allocated
and initialized in the child thread, finally passing the base pointer to
__wasm_set_tls_base.

Coop threads & Multi module

During component instantiation all modules are instantiated as well. Each module
initializes its main-thread TLS, sets its own personal __init_tls_base (which
is then never read again), and then calls __wasm_set_tls_base. This
function is synthesized by wit-component which initializes
main_thread_tls_base because context.get 1 is 0.

The __wasm_init_task function will see that __wasm_program_tls_info is not
NULL and will then call context.set 1 with the address of
main_thread_tls_base, which was populated during module instantiation.

The __wasm_get_tls_base function imported for each module is synthesized by
wit-component which loads from a static offset from context.get 1.

When a new thread is spawned the __wasm_program_tls_info symbol is present, so
TLS allocation will look quite different. This'll have to do all sorts of math
to align everything appropriately and such, but the end result is that there'll
be a size/align for the block of TLS for the whole program, and then this is
handled similarly to today. When the thread starts it'll call context.set 1
with the block of pointers allocated for the thread, and then TLS for each
library is initialized with the various function pointers in
__wasm_program_tls_info.

Future changes

One minor change possible in the future is to extend DynamicLinking.md in the
upstream tool-conventions repository to update the dylink.0 custom section
with information about a library's TLS size/align. Currently this needs to be
inferred from the __tls_{size,align} globals which isn't too bad but would be
nicer if we didn't have to do this. This in theory would be a pretty minor
extension and I suspect wouldn't be too controversial, and it'd help clean up a
bit.

Otherwise though this is all designed on the premise that changing LLVM's
codegen isn't desired. This might be able to all be simplified if that was
pulled into the picture, but I'm also not exactly sure what this would look
like. As a possible argument against changing LLVM it's kind of nice that
single-module use cases get to avoid the level of indirection that multi-module
use cases have, and this might necessitate not changing LLVM in the long
run. Ultimately though I mostly just wanted to point out that this was all
designed through the lens of "let's not change LLVM", but in the long-run
there's probably at least some minor tweak that could be made to simplify this
overall story. At this time, though, it's expected that this is only necessary
for cleanup and not for functionality gains.

---

cc @TartanLlama as you're likely to have thoughts on this too

[TartanLlama]

This approach makes sense to me, and it seems practical to try and find a solution that doesn't require modifying LLVM.

[alexcrichton]

First wrinkle: this scheme means that cross-library references to thread-locals are still broken. For example right now there's extern _Thread_local int errno; which defines how all libraries access errno. The scheme defined in this issue means that cross-library references to errno are broken because it's always relative to the library-accessing-errno's tls base, which is incorrect for accessing wasi-libc's tls base.

This can be mitigated by switching to __errno_location() for errno accesses, but it'd still be a toolchain bug for any other externs. This, however, I think may only be fixable with LLVM changes.

[alexcrichton]

Ok I've done an initial attempt at implementing all of this. Locally all the wasi-libc tests pass, but CI doesn't show that because of the split between repos. The changes are:

• Implement new handling of TLS base/symbols bytecodealliance/wasm-tools#2590
• Support shared-library TLS #858

The wrinkle I described above mostly ended up being the only one. I think Windows has the same limitation for cross-library TLS so it may not be all that bad in the long run.

The only other wrinkle I encountered was getting __tls_{size,align} exported with __wasm_init_tls as well. Right now that's custom linker flags passed in wasi-libc's build system. I think we're going to want to update Rust/Clang compiler drivers to pass those arguments by default. Exporting them shouldn't be a problem, and it's easy to be wasip3-specific. Other possible solutions all seemed worse, so I think that's the best for now.

@dicej I know you mentioned you were going to take a stab at this next week so sorry about that :(. I may have been a bit too eager to prove this all out. No rush reviewing/landing though, I'm fine to leave this open for a bit.