aboutsummaryrefslogtreecommitdiffstats
path: root/library/cpp/actors/util/unordered_cache.h
blob: 40794fc04bfb5c2a9af13d544c55b4ad9c24eb37 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
#pragma once

#include "defs.h"
#include "queue_chunk.h"

template <typename T, ui32 Size = 512, ui32 ConcurrencyFactor = 1, typename TChunk = TQueueChunk<T, Size>>
class TUnorderedCache : TNonCopyable {
    static_assert(std::is_integral<T>::value || std::is_pointer<T>::value, "expect std::is_integral<T>::value || std::is_pointer<T>::value");

public:
    static constexpr ui32 Concurrency = ConcurrencyFactor * 4;

private:
    struct TReadSlot {
        TChunk* volatile ReadFrom;
        volatile ui32 ReadPosition;
        char Padding[64 - sizeof(TChunk*) - sizeof(ui32)]; // 1 slot per cache line
    };

    struct TWriteSlot {
        TChunk* volatile WriteTo;
        volatile ui32 WritePosition;
        char Padding[64 - sizeof(TChunk*) - sizeof(ui32)]; // 1 slot per cache line
    };

    static_assert(sizeof(TReadSlot) == 64, "expect sizeof(TReadSlot) == 64");
    static_assert(sizeof(TWriteSlot) == 64, "expect sizeof(TWriteSlot) == 64");

private:
    TReadSlot ReadSlots[Concurrency];
    TWriteSlot WriteSlots[Concurrency];

    static_assert(sizeof(TChunk*) == sizeof(TAtomic), "expect sizeof(TChunk*) == sizeof(TAtomic)");

private:
    struct TLockedWriter {
        TWriteSlot* Slot;
        TChunk* WriteTo;

        TLockedWriter()
            : Slot(nullptr)
            , WriteTo(nullptr)
        { }

        TLockedWriter(TWriteSlot* slot, TChunk* writeTo)
            : Slot(slot)
            , WriteTo(writeTo)
        { }

        ~TLockedWriter() noexcept {
            Drop();
        }

        void Drop() {
            if (Slot) {
                AtomicStore(&Slot->WriteTo, WriteTo);
                Slot = nullptr;
            }
        }

        TLockedWriter(const TLockedWriter&) = delete;
        TLockedWriter& operator=(const TLockedWriter&) = delete;

        TLockedWriter(TLockedWriter&& rhs)
            : Slot(rhs.Slot)
            , WriteTo(rhs.WriteTo)
        {
            rhs.Slot = nullptr;
        }

        TLockedWriter& operator=(TLockedWriter&& rhs) {
            if (Y_LIKELY(this != &rhs)) {
                Drop();
                Slot = rhs.Slot;
                WriteTo = rhs.WriteTo;
                rhs.Slot = nullptr;
            }
            return *this;
        }
    };

private:
    TLockedWriter LockWriter(ui64 writerRotation) {
        ui32 cycle = 0;
        for (;;) {
            TWriteSlot* slot = &WriteSlots[writerRotation % Concurrency];
            if (AtomicLoad(&slot->WriteTo) != nullptr) {
                if (TChunk* writeTo = AtomicSwap(&slot->WriteTo, nullptr)) {
                    return TLockedWriter(slot, writeTo);
                }
            }
            ++writerRotation;

            // Do a spinlock pause after a full cycle
            if (++cycle == Concurrency) {
                SpinLockPause();
                cycle = 0;
            }
        }
    }

    void WriteOne(TLockedWriter& lock, T x) {
        Y_DEBUG_ABORT_UNLESS(x != 0);

        const ui32 pos = AtomicLoad(&lock.Slot->WritePosition);
        if (pos != TChunk::EntriesCount) {
            AtomicStore(&lock.Slot->WritePosition, pos + 1);
            AtomicStore(&lock.WriteTo->Entries[pos], x);
        } else {
            TChunk* next = new TChunk();
            AtomicStore(&next->Entries[0], x);
            AtomicStore(&lock.Slot->WritePosition, 1u);
            AtomicStore(&lock.WriteTo->Next, next);
            lock.WriteTo = next;
        }
    }

public:
    TUnorderedCache() {
        for (ui32 i = 0; i < Concurrency; ++i) {
            ReadSlots[i].ReadFrom = new TChunk();
            ReadSlots[i].ReadPosition = 0;

            WriteSlots[i].WriteTo = ReadSlots[i].ReadFrom;
            WriteSlots[i].WritePosition = 0;
        }
    }

    ~TUnorderedCache() {
        Y_ABORT_UNLESS(!Pop(0));

        for (ui64 i = 0; i < Concurrency; ++i) {
            if (ReadSlots[i].ReadFrom) {
                delete ReadSlots[i].ReadFrom;
                ReadSlots[i].ReadFrom = nullptr;
            }
            WriteSlots[i].WriteTo = nullptr;
        }
    }

    T Pop(ui64 readerRotation) noexcept {
        ui64 readerIndex = readerRotation;
        const ui64 endIndex = readerIndex + Concurrency;
        for (; readerIndex != endIndex; ++readerIndex) {
            TReadSlot* slot = &ReadSlots[readerIndex % Concurrency];
            if (AtomicLoad(&slot->ReadFrom) != nullptr) {
                if (TChunk* readFrom = AtomicSwap(&slot->ReadFrom, nullptr)) {
                    const ui32 pos = AtomicLoad(&slot->ReadPosition);
                    if (pos != TChunk::EntriesCount) {
                        if (T ret = AtomicLoad(&readFrom->Entries[pos])) {
                            AtomicStore(&slot->ReadPosition, pos + 1);
                            AtomicStore(&slot->ReadFrom, readFrom); // release lock with same chunk
                            return ret;                             // found, return
                        } else {
                            AtomicStore(&slot->ReadFrom, readFrom); // release lock with same chunk
                        }
                    } else if (TChunk* next = AtomicLoad(&readFrom->Next)) {
                        if (T ret = AtomicLoad(&next->Entries[0])) {
                            AtomicStore(&slot->ReadPosition, 1u);
                            AtomicStore(&slot->ReadFrom, next); // release lock with next chunk
                            delete readFrom;
                            return ret;
                        } else {
                            AtomicStore(&slot->ReadPosition, 0u);
                            AtomicStore(&slot->ReadFrom, next); // release lock with new chunk
                            delete readFrom;
                        }
                    } else {
                        // nothing in old chunk and no next chunk, just release lock with old chunk
                        AtomicStore(&slot->ReadFrom, readFrom);
                    }
                }
            }
        }

        return 0; // got nothing after full cycle, return
    }

    void Push(T x, ui64 writerRotation) {
        TLockedWriter lock = LockWriter(writerRotation);
        WriteOne(lock, x);
    }

    void PushBulk(T* x, ui32 xcount, ui64 writerRotation) {
        for (;;) {
            // Fill no more then one queue chunk per round
            const ui32 xround = Min(xcount, (ui32)TChunk::EntriesCount);

            {
                TLockedWriter lock = LockWriter(writerRotation++);
                for (T* end = x + xround; x != end; ++x)
                    WriteOne(lock, *x);
            }

            if (xcount <= TChunk::EntriesCount)
                break;

            xcount -= TChunk::EntriesCount;
        }
    }
};