Thread Local Storage Matrix

Thread Local Storage Matrix

Ordinum's implementation of Thread Local Storage and how the storage engine interacts with it

This journal walks through the design challenges and decisions faced with when using thread local storage for the Ordinum storage engine.

At a high level, thread local storage is quite simply what it's title suggests. A per thread process with storage local to the thread that can be called into and accessed for that thread only, independent of other thread processes. In rust this comes in the form of the thread_local!() macro

macro_rules! thread_local {
    () => { ... };
    ($($tt:tt)+) => { ... };
}

It implements std::thread::LocalKey which is a key into the underlying storage for tls on the target platform. For example, on Linux that maybe the TLS support of the ELF ABI (..).

LocalKey uses the fastest implementation available on the target platform and is constructed with the thread_local! macros as described above.

Some interesting points when it comes to TLS and LocalKey:

• Initialisation is done dynamically/lazily on the first call to a setter X.with(..)
• Although TLS is a single thread primitive, it is possible for thread local state to be shared with other threads so the implementation detail specifies that only &T references may be obtained. It is therefore necessary to encapsulate tls fields with interior mutability primitives such as Cell<>, UnsafeCell<>, RefCell<> etc. if mutability is required.
• Destructors are 'best effort' and platform specific. A number of caveats are known for where destructors are not run (..)

Here is an example of how we would initialise and interact with thread local storage in rust:

use std::cell::{Cell, RefCell};

thread_local! {
    pub static FOO: Cell<u32> = const { Cell::new(1) };

    static BAR: RefCell<Vec<f32>> = RefCell::new(vec![1.0, 2.0]);
}

assert_eq!(FOO.get(), 1);
BAR.with_borrow(|v| assert_eq!(v[1], 2.0));

The book Rust Atomics and Locks by Mara Bos has an excellent introduction into threads and concurrency for low level systems. Highly worth a read.

Ordinum makes use of thread local storage quite heavily, not just for capturing metrics but as subsystems for optimisations and efficiencies.

Problem Statement

Ordinum has a number of sub-systems which require the use of thread local storage to speed up processes and behaviour, and to reduce strain on global structures. It also needs thread local storage to capture per process metrics and local state which extend for the lifetime of the thread. Both of these are orthogonal to each other. The former must ineract with the program with state having different lifetimes and accessors based on the program logic, for example scoped to per database instances. The latter, stretches for the length of the thread process and is purely local to that thread.

For the storage state which must interact with the program, the problem becomes, how do we effectively separate state from different instances of the program and protect cross thread interaction.

Those are the two axis which we are to focus on for this.

If we cast our mind back up to the TLS/LocalKey invariants, it is mentioned that we only are given &T references back from TLS and that we must carefully address the fact that other threads can (and in our case, will) touch thread local storage and more importantly may mutate thread entry state in certain cases.

Ordinum will have subsystems of varying complexity and needs which will need to utilise tls and this problem space is what we'll address further in the journal.

Naive Implementation

Where possible, it is recommended to start with the naive implementation. Although my perfectionism often overules this and forces me into an optimisation blender. For Ordinum's thread local storage, we truly started with the basic implementation and discovered along the way what needed to be changed based on the problem evolving as we introduced more complex subsystems and invariants.

We will talk to 3 subsystems (non-exhaustive), each with their own needs and each covering the different problems as described.

1. PerfContext

- Local thread metrics

1. BatchCache

- A cache local to the thread which stores an array of batches for grouped writes

1. SuperVersion Cache

- A cached pointer to the superversion subsystem for snapshot reads

To begin, it is useful to briefly outline why we decide to use thread local storage, the decision as to why we would want to use it and how Ordinum can makes use of it to create efficiencies.

Why We use Thread Local Storage

We can see how thread local storage might be useful by starting with the example of the write batch. For writers writing to the database engine, at it's simplist, we can imagine a single thread/process which takes bytes and calls into the storage engine to write them.

The storage engines job is to carry out the write. Our job is to make sure the storage engine does this efficiently and the writer is not stuck or waiting a long time during this process.

The storage engine will seek to batch the writes, again thinking simply, this will take the form of allocating a contigous chunk of memory like a Vec<U8> to house the batched writes and run it through it's pipeline process

Now, clearly looking at this you may think, 'Well, why not create a pool of batches?' and you'd be right! So the next optimisation would be to create a global pool scoped to the database instance. Of course, there'd be many different ways you could go about this, ways which are outside the scope of this journal.

We create a global pool whereby we lazily allocate batches and return them to the pool on drop.

This works fine, and for most workloads this is an already optimal solution. But, thinking about thread local storage, can we do better?

Yes! We can and we shall.

If we imagine a very basic batch pool looking somehting like:

struct BatchPool {
  batches: Mutex<Vec<Batch>>,
}

We have a shared structure which many processes will be hitting. We have a Mutex<_> which all those processes will be trying to acquire and will be waiting for if they do not have the lock.

The Write Path is most definately a hot path and so we want to reduce contention as much as possible, where possible.

The solution is to lean on thread local storage to allow threads to avoid hitting global structures on the hot path and thereby avoid contention. We implement a caching structure in thread local storage where allocated batches are cached and can be reused by threads on writes.

The process flow can be reduced to:

1. Thread checks it's own cache
2. If empty, check the global pool
3. If empty, allocate new batch

And on a thread finishing with a batch, the reverse of this is:

1. Try to return to thread local cache
2. If full, try to return to global
3. If full, destroy the batch

Building Thread Local Storage for Ordinum

The first introduction I had to thread-local storage (TLS) came while tackling what is probably every systems programmer's worst nightmare: memory reclamation.

I was working on a problem that required the use of Epoch based reclamation While that's outside the scope of this journal, it was certainly a rude awakening to the realities of concurrent state, synchronisation, and designing data structures that can safely outlive the threads accessing them.

For Ordinum, the first iteration of TLS looked almost exactly like the introductory example above. We had a simple registry module containing a thread context that stored any per-thread objects we wanted to cache or reuse.

This implementation served its purpose well. It allowed the design to evolve naturally while answering questions such as: Why are we using TLS here? When does it make sense to cache data on a per-thread basis? When is TLS the wrong solution?

Before long, however, a more fundamental architectural question emerged:

What happens if an application opens multiple database instances?

This was the first major design decision surrounding the TLS implementation. Should Ordinum even support multiple database instances?

There are perfectly reasonable arguments for supporting only a single database per process, and many applications will never need anything more. However, allowing multiple independent database instances provides an important level of operational isolation. Each database owns its own storage engine, write pipeline, WAL, compaction scheduler, caches, and configuration, allowing different workloads to coexist without interfering with one another.

This naturally raises another question: couldn't these simply be column families instead?

The answer is that column families solve a different problem.

Database Instances vs Column Families

A column family is a logical namespace within a database. It has its own memtable and LSM tree, but shares the database's infrastructure.

A database instance is a completely independent storage engine with its own WAL, write pipeline, background workers, caches, and metadata. (Although the benefit of database instances on the same machine is that they can still share and utilise program wide infrastructure such as thread pooling).

We would ideally use a comlumn familty when the datasets belong together and should share infrastructure.

• Users
• Orders
• Products

And use a database when the datasets are operationally independent.

For example, if the cache database fills the disk or experiences heavy write stalls, the user database can continue operating unaffected.

But again, this is purely a matter of opinion, and only you know the needs of your workload, Ordinum just supplies the concise tools and engine to service those.

I digress. Ordinum ultimately chose to support mutliple database instances, and this brings us to the first iteration of that.

Thread Local Matrix

thread_local! {
    static THREAD_CTX: UnsafeCell<ThreadContext> =
        UnsafeCell::new(ThreadContext::new());
}

/// Main thread local storage structure
struct ThreadContext {
  //
  instances: Mutex<HashMap<usize, DBInstanceContext>>,
  //
  // Other fields not relevant ...
}

/// Structure for each database instance stored within thread local storage
struct DBInstanceContext {
  //
  batch_cache: Vec<Batch>,
  //
  // Other tls sub sytems ...
}

Our first approach was to simply define the database instance as a seperate structure which we would store in a HashMap inside the ThreadContext this gives the benefit of being very simple and very clear in it's intention.

In the code example we have the DBInstanceContext struct which houses the batch_cache. If we were to put another sub-system in there we can see how we might be encapsulating state within the context of the db instance sort of like a registry.

/// Structure for each database instance stored within thread local storage
struct DBInstanceContext {
  //
  batch_cache: Vec<Batch>,
  //
  superversion_cache: SVCache,
  perf_context: PerfContext,
  //...
}

For this approach we needed an ID to be able to access each database instance and retrieve it's context of sub-systems. This is simple enough, on each DB::Open() we increment a global const Atomic number DB_ID.fetch_add(1, Ordering::Release). And on each thread we lazily add to the HashMap on first access to TLS.

The key mental model to hold is that we (at this point) are storing the DBInstanceContext object as a whole, including all sub-systems within. So when we access thread local data, we must go through the context. This is ok for simple cases such as referencing/reading the data inside but for operations on sub-systems that might require different lifetimes or mutability contracts, we run into problems.

The "aha!" (or perhaps more accurately, the "oh no... time to refactor") moment came while thinking about how the objects inside DBInstanceContext would actually be destroyed.

At some point the database instance is closed, meaning every thread-local subsystem stored within the context must eventually be dropped. The question quickly became: when is it actually safe to do that?

Initially it seemed reasonable that the DBInstanceContext itself should own this responsibility. However, the more I thought about it, the less practical that became. Different subsystems have completely different lifetime requirements. A thread-local batch cache, for example, has a very different notion of "safe to destroy" than a SuperVersion cache holding protected pointers.

To make destruction safe from the context itself would require introducing in-flight operation counters, additional reference counting, or other coordination mechanisms so that the context could somehow know when every subsystem had reached a quiescent state. Even then, the context would be making the flawed assumption that every subsystem follows the same shutdown semantics.

That was the realisation: DBInstanceContext was owning far too much. Rather than treating every subsystem as though it shared a common lifetime model, each subsystem should own its own lifecycle and define for itself what "safe to destroy" actually means.

Luckily, this problem had already been encountered and implemented by RocksDB who design the thread local storage structure as a matrix. (thread_local.cc)

This is the structure that is declared as "thread_local" storage.
The vector keep list of atomic pointer for all instances for "current"
thread. The vector is indexed by an Id that is unique in process and
associated with one ThreadLocalPtr instance. The Id is assigned by a
global StaticMeta singleton. So if we instantiated 3 ThreadLocalPtr
instances, each thread will have a ThreadData with a vector of size 3:
     ---------------------------------------------------
     |          | instance 1 | instance 2 | instance 3 |
     ---------------------------------------------------
     | thread 1 |    void*   |    void*   |    void*   | <- ThreadData
     ---------------------------------------------------
     | thread 2 |    void*   |    void*   |    void*   | <- ThreadData
     ---------------------------------------------------
     | thread 3 |    void*   |    void*   |    void*   | <- ThreadData
     ---------------------------------------------------

Similar to the diagram at the beginning of this journal, each thread is represented by a row in the matrix. The important difference is that the meaning of each column has changed. Previously, each column represented a database instance, with a DBInstanceContext acting as a container for every thread-local subsystem associated with that database.

Instead, each column now represents a single thread-local subsystem instance. Rather than assigning a DB_ID, each subsystem registers itself and is assigned a unique TLS_ID. This TLS_ID is generated by a global allocator and used as the index into each thread's entries vector. On first access, the vector is lazily resized and a pointer to that subsystem's thread-local state is stored directly in the corresponding slot.

The immediate benefit is that we no longer have to access a subsystem by first retrieving a DBInstanceContext. Each subsystem can be accessed directly through its own TLS_ID, allowing it to own its own initialization, lifetime, and destruction semantics independently of every other subsystem.

This may initially seem like a subtle distinction, but it becomes much more powerful as additional thread-local subsystems are introduced. Consider the BatchCache. Since each database owns a single batch cache, it is natural for each database instance to contribute one column to the matrix.

The advantage becomes much clearer with the SuperVersionCache. Unlike the batch cache, SuperVersions exist on a per-column-family basis. A single database may contain many column families, each requiring its own cached SuperVersion. Under the original DBInstanceContext design, these independent caches would all have been hidden behind a single context object. With the new model, each cache simply registers its own TLS_ID, resulting in one column per cached SuperVersion instance. The thread-local matrix does not need to understand what the subsystem represents it simply provides fast, direct access to thread-local state.

                         Database 1
                  ┌──────────────────────┐
                  │ BatchCache           │
                  │ SVCache(CF0)         │
                  │ SVCache(CF1)         │
                  │ SVCache(CF2)         │
                  └─────────┬────────────┘
                            │
                            ▼

        ------------------------------------------------------------------------------------------------------------
        |          | BatchCache(DB1) | SVCache(CF0) | SVCache(CF1) | SVCache(CF2) | BatchCache(DB2) | ... |
        ------------------------------------------------------------------------------------------------------------
        | thread 1 |      void*      |    void*     |    void*     |    void*     |      void*      |     |
        ------------------------------------------------------------------------------------------------------------
        | thread 2 |      void*      |    void*     |    void*     |    void*     |      void*      |     |
        ------------------------------------------------------------------------------------------------------------
        | thread 3 |      void*      |    void*     |    void*     |    void*     |      void*      |     |
        ------------------------------------------------------------------------------------------------------------

The next problem is 'how do we implement this?'

Implementation

We need to think about the two axes of the thread-local matrix: columns and rows. Each represents a different access path and therefore serves a different purpose.

As we've discussed already, a column represents a distinct subsystem (or vertical), independent of its neighbouring columns. A row represents a thread and contains that thread's data for each registered subsystem.

An easy way to visualise this is as a spreadsheet. To access a particular cell, we must first identify the column (the subsystem), then traverse to the required row (the thread). The intersection of the two gives us the object we wish to access.

This naturally gives us two traversal directions:

• Across rows – iterate over every thread for a particular subsystem.
• Across columns – iterate over every subsystem belonging to a particular thread.

Before diving into the implementation, it's worth understanding why these traversal paths are necessary.

We'll use two examples:

1. Database shutdown
2. Thread exit

Database Shutdown:

In our earlier BatchPool example, we described how each pool is scoped to a database instance. Multiple threads may be accessing that pool simultaneously, meaning each thread's row contains a cell for that pool's thread-local cache.

When the database shuts down, the pool itself is destroyed. Consequently, the entire column representing that subsystem must be removed to prevent threads from accessing caches belonging to a database that no longer exists.

To achieve this, the shutdown thread traverses every row and invokes the registered destructor for that column. Only the cells belonging to the shutting-down database are removed; all other columns and thread-local subsystems remain untouched.

Thread Exit:

Thread exit is where we can expect to experience the highest churn. Many threads may come and go, be spun up and destroyed, or equally be long running.

When a thread exits, we instead traverse across that thread's row. For every populated cell, the registered destructor is invoked before the entry is cleared. Once complete, the entire row can be safely reclaimed.

This operation is performed while holding the global thread registry mutex. The mutex serialises thread registration and removal, ensuring that database shutdown and thread exit cannot race while traversing the registry.

To start with, we can begin looking at what data threads will store and how it is structured. We begin with two high level objects:

1. ThreadMetaGlobal
2. ThreaData

..

Each row is owned by a particular thread but is discoverable by other threads through the global registry. Accessing another thread's row is an uncommon operation used for coordination tasks (e.g. cache invalidation or reclamation) and therefore requires carefully designed synchronization and ownership invariants.

To do this we use a Doubly Linked List where we start at a sentinel node and traverse registered threads