Thursday, May 20, 2010

Tidying up TLS

So, I wanted a nice place to tear down TLS data, and found (yet another) piece of Win32 that isn’t very nicely designed.

Whilst the (newer) fiber-local storage allows you to specify a clean-up callback, the same isn’t true of TLS. So if you want to write a library that uses TLS behind the scenes you have an annoying problem. Either leak the TLS data whenever the thread exits, or require the library user to explicitly call a clean-up function.

One solution is to package your code as a DLL. DLLs get notified when each thread exits, providing a neat place to tear down the TLS data. But sometimes, I want a static .lib instead of a DLL. For installer-less programs, I typically want everything statically compiled into a single .exe.

There are also situations where I’m writing a program (not a library) but don’t have any particular control over threads; for example, when using thread pooling. Yes, my work items could each tear down the TLS data, but that’s not necessarily what I want. Thread-specific caches (for memory allocators, say) shouldn’t be torn down at the end of each work item: they should be allowed to exist right up until the thread pool ends the thread.

Neither of these is particularly appealing. However, it turns out that the PE file format has a solution. PEs contain a .tls section. This section contains a set of pointers to callback functions. The callback functions are essentially equivalent to the DllMain callbacks.

So, to get automatic cleaning of TLS, we just need to put a suitable callback in the .tls section of the executable.

The difficult bit is to get the linker to put our function pointers in to the .tls section. The VC++ C Runtime actually knows about the .tls section (even though it doesn’t generally use it), and so we can get our callbacks registered by putting function pointers in any section .CRT$XLx (where x is A-Z).

The code to do this would look something like:

void NTAPI on_tls_callback(void* dll, DWORD reason, void* reserved)
{
    UNREFERENCED_PARAMETER(dll);
    UNREFERENCED_PARAMETER(reserved);
    switch(reason)
    {
    case DLL_PROCESS_ATTACH:
        printf("process_attachn");
        break;
    case DLL_PROCESS_DETACH:
        printf("process_detachn");
        break;
    case DLL_THREAD_ATTACH:
        printf("thread_attachn");
        break;
    case DLL_THREAD_DETACH:
        printf("thread_detachn");
        break;
    }
}

#ifdef _M_IX86
#pragma comment (linker, "/INCLUDE:__tls_used")
#pragma comment (linker, "/INCLUDE:__xl_b")
#else
#pragma comment (linker, "/INCLUDE:_tls_used")
#pragma comment (linker, "/INCLUDE:_xl_b")
#endif

#ifdef _M_X64
#pragma const_seg(".CRT$XLB")
EXTERN_C const
#else
#pragma data_seg(".CRT$XLB")
EXTERN_C
#endif

PIMAGE_TLS_CALLBACK _xl_b = on_tls_callback;

#ifdef _M_X64
#pragma const_seg()
#else
#pragma data_seg()
#endif