【備忘録】Lock-freeなスタックを作る際の落とし穴

今回はカスタムアロケータを作る際にハマったLock-freeスタックの落とし穴について簡単に書き残しておきます。
# ハマった実装
ここでは解説ポイントを絞るためにかなり単純化したアロケータを作成しています。

```c++
#include <atomic>
#include <thread>
#include <vector>
using namespace std;

// 型Tのオブジェクトを1個ずつ確保したり解放したりするだけのアロケータ
// 解放されたメモリはリンクリストに繋げられ次回以降の確保時に取り出される
// atomicを使ってロックフリーとして実装しているので複数スレッドからの同時利用が可能（落とし穴）
template<typename T>
class Allocator {
    union Link {
        T value;
        Link* next;
    };
    atomic<Link*> root{};
public:
    ~Allocator() {
        for (Link* l = root; l;) {
            auto next = l->next;
            delete l;
            l = next;
        }
    }

    // 後入れ先出しでフリーメモリをリスト管理
    T* alloc() {
        Link* l = root;
        do {
            if (l == nullptr) return &(new Link)->value;
        } while (!root.compare_exchange_weak(l, l->next));
        return &l->value;
    }

    void dealloc(void* p) {
        Link* l = root;
        auto dealloced = reinterpret_cast<Link*>(p);
        do {
            dealloced->next = l;
        } while (!root.compare_exchange_weak(l, dealloced));
    }
};

int main() {
    // スレッドたくさん作ってひたすらメモリの確保と解放を繰り返すだけ
    vector<thread> workers(thread::hardware_concurrency());
    Allocator<int> intPool;
    for (auto& w : workers) {
        w = thread([&] {
            for (int i = 0; i < 100000; i++)
                intPool.dealloc(intPool.alloc());
            });
    }
    for (auto& w : workers) {
        if (w.joinable())w.join();
    }
    return 0;// !この後ハング
}
```
こちらのアロケータですが実はそこそこ問題なく使えてしまいます。パっと見た感じでも問題なさそうに見えますが・・・
しかし、上記のように意図的に確保と解放を超高頻度で実行すると解法済みリストへの二重登録が発生してしまい、Allocatorのデストラクタで二重解放が起きてしまうのです。
# 問題の分析
何が起きたのか紐解いていくと・・・
管理されるメモリにA,B,C...と名前を付けるとして

```一部抜粋.cpp
    T* alloc() {
        Link* l = root;
        do {
            if (l == nullptr) return &(new Link)->value;
        } while (!root.compare_exchange_weak(l, l->next)); // ※
        return &l->value;
    }

    void dealloc(void* p) {
        ...
    }
```

| スレッド1 | スレッド2 | スレッド3 | リスト |
|:-:|:-:|:-:|:-:|:-:|
| alloc開始<br>※で関数の引数であるA, Bが読み込まれ<br>compare_exchange_weakが実行される直前 ||| A, B, C|
|| alloc開始<br>Aを取得 || B, C |
|||alloc開始<br>Bを取得| C |
||dealloc(A)||A, C|
|ここでcompare_exchange_weak(A, B)実行<br>Aを取得|||B, C|
|||dealloc(B)|B, B, C|
と、このように抜粋コードの※印の個所に問題がありました。
スレッドの切り替えはいかなるCPU命令間でも起きるので関数の引数をロードしてから実際に関数が呼ばれるまでの間でもスレッドが切り替わって状態が変わり得るということです。
そのせいでメモリBがリストに二重登録されてしまい不具合となったのです。

# 解決策
上記の例ではスレッド1によって保持されていたA->nextがCAS操作前に書き換えられていたにもかかわらずCASが成功してしまったために不整合が起きていたのでした。では、CASの比較を「**アドレス＋タグ**」で行うようにしてみてはどうでしょうか。
**タグ**はアロケートの度にインクリメントされる通し番号のようなものです。こうすれば同じアドレスであってもタグが異なる場合nextの値も書き変わっている可能性があることを検出できるので整合性を保つことができます。
ちなみにポインタサイズ以上の型でatomicがロックフリー動作するのか気になると思いますが、その場合は`atomic<T>::is_always_lock_free`や`atomic<T>::is_lock_free()`にてチェックしてください。環境によりますがポインタ2つ分まではロックフリーに行える環境が多いと思います。
以上を踏まえた最終実装がこちらになります。

:::note info
少なくともMSVCではこちらの解説は不完全でした。追記に情報あるので合わせてお読みください。
:::
```c++
template<typename T>
class Allocator {
    union Link {
        T value;
        Link* next;
    };
    struct TaggedPtr {
        Link* ptr;
        size_t tag;
        Link* operator->() const {
            return ptr;
        }
    };
    atomic<TaggedPtr> root{};
public:
    ~Allocator() {
        for (auto l = root.load().ptr; l;) {
            auto next = l->next;
            delete l;
            l = next;
        }
    }

    T* alloc() {
        TaggedPtr l = root;
        do {
            if (l.ptr == nullptr) return &(new Link)->value;
        } while (!root.compare_exchange_weak(l, { l->next, l.tag + 1u}));
        return &l->value;
    }

    void dealloc(void* p) {
        TaggedPtr l = root;
        auto dealloced = reinterpret_cast<Link*>(p);
        do {
            dealloced->next = l.ptr;
        } while (!root.compare_exchange_weak(l, { dealloced, l.tag }));
    }
};
```
アロケータが持つポインタのサイズがsize_t分増えただけで元の実装と大差ないですがこれで完璧にLock-freeなスタックによるアロケータが実装出来ました。
# 追記：atomic_ref
この記事を書いた後改めてMSVCで確認したら`atomic<TaggedPtr>`はlock freeになってませんでした。
（昔うまくいってた気がしてたんだが勘違いだったのか・・？）
とりあえずMSVCでc++20以降であれば[atomic_ref](https://cpprefjp.github.io/reference/atomic/atomic_ref.html)クラスが使えます。atomic_refだけは16byteの時でもlock freeで動作するみたいでした。（実行環境がサポートしていれば）
それを踏まえてatomic_ref版の実装を以下に示します。

```c++
template<typename T>
class Allocator {
    union Link {
        T value;
        Link* next;
    };
    struct alignas(16) TaggedPtr {
        Link* ptr{};
        size_t tag;
        Link* operator->() const {
            return ptr;
        }
    };
    TaggedPtr root;
public:
    ~Allocator() {
        for (auto l = root.ptr; l;) {
            auto next = l->next;
            delete l;
            l = next;
        }
    }

    T* alloc() {
        TaggedPtr l = atomic_ref{ root };
        do {
            if (l.ptr == nullptr) return &(new Link)->value;
        } while (!atomic_ref{ root }.compare_exchange_weak(l, { l->next, l.tag + 1u}));
        return &l->value;
    }

    void dealloc(void* p) {
        TaggedPtr l = atomic_ref{ root };
        auto dealloced = reinterpret_cast<Link*>(p);
        do {
            dealloced->next = l.ptr;
        } while (!atomic_ref{ root }.compare_exchange_weak(l, { dealloced, l.tag }));
    }
};
```
まずTaggedPtrですが、128bitCAS命令を使えるようにするためにalignas(16)しています。これをしておかないと実行時にハングします。
そしてrootですが、こちらはatomic_ref型ではなく生のTaggedPtrとして保持します。
atomic_refはコンストラクタでアトミック操作対象の参照を取りそのアドレスに対して処理を行うのですが、コンストラクタ自体は何もしていないのでアトミック操作したくなった時に適宜atomic_ref{ root }と囲えばよいのです。
また注意点として排他的にloadしたい場合も生のrootからコピーを取らずにいったんatomic_refで囲ってからload()しないとアトミック読み込みになりません。（読み取り中に部分的に別スレッドに書き換えられるかもしれない）

## is_lock_free(),is_always_lock_free
この確認をしているときにロックフリー確認用のAPIがどんな値を示すか確認したんですが、MSVCのis_always_lock_freeは型Tのサイズが1,2,4,8バイトでないとtrueになりませんでしたが、is_lock_free()はサポートチェック関数を呼び出して16バイトでもtrueを返しました。
このあたりの挙動にも要注意ですね。

# 感想
最初は全く問題点発見できなくて発狂寸前でした。見つけた時はもはやすっきりするというより「うわぁ・・・」って感じでしんどかったです。
どうやらこれABA問題と呼ばれる有名な現象のようです。Wikiにがっつりこの記事と同じこと書いてありました(笑)。それと今回は大きいサイズのアトミック操作を使ってロックフリーを実現しましたが、パフォーマンスという点ではスピンロックなどの方が上という話もあるみたいなのでパフォーマンスを重視したい場合はやはり実測で最適な手法を選ぶしかないですね。