/src/brpc/src/bthread/key.cpp

Source
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

// bthread - An M:N threading library to make applications more concurrent.

// Date: Sun Aug  3 12:46:15 CST 2014

#include <pthread.h>
#include <gflags/gflags.h>

#include "bthread/bthread.h"     // bthread_create_span_fn and related types
#include "bthread/errno.h"       // EAGAIN
#include "bthread/task_group.h"  // TaskGroup
#include "butil/atomicops.h"
#include "butil/macros.h"
#include "butil/thread_key.h"
#include "butil/thread_local.h"
#include "bvar/passive_status.h"

// Implement bthread_key_t related functions

namespace bthread {

DEFINE_uint32(key_table_list_size, 4000,
              "The maximum length of the KeyTableList. Once this value is "
              "exceeded, a portion of the KeyTables will be moved to the "
              "global free_keytables list.");

DEFINE_uint32(borrow_from_globle_size, 200,
              "The maximum number of KeyTables retrieved in a single operation "
              "from the global free_keytables when no KeyTable exists in the "
              "current thread's keytable_list.");

EXTERN_BAIDU_VOLATILE_THREAD_LOCAL(TaskGroup*, tls_task_group);

class KeyTable;

// defined in task_group.cpp
extern __thread LocalStorage tls_bls;
static __thread bool tls_ever_created_keytable = false;

// We keep thread specific data in a two-level array. The top-level array
// contains at most KEY_1STLEVEL_SIZE pointers to dynamically allocated
// arrays of at most KEY_2NDLEVEL_SIZE data pointers. Many applications
// may just occupy one or two second level array, thus this machanism keeps
// memory footprint smaller and we can change KEY_1STLEVEL_SIZE to a
// bigger number more freely. The tradeoff is an additional memory indirection:
// negligible at most time.
static const uint32_t KEY_2NDLEVEL_SIZE = 32;

// Notice that we're trying to make the memory of second level and first
// level both 256 bytes to make memory allocator happier.
static const uint32_t KEY_1STLEVEL_SIZE = 31;

// Max tls in one thread, currently the value is 992 which should be enough
// for most projects throughout baidu.
static const uint32_t KEYS_MAX = KEY_2NDLEVEL_SIZE * KEY_1STLEVEL_SIZE;

// destructors/version of TLS.
struct KeyInfo {
    uint32_t version;
    void (*dtor)(void*, const void*);
    const void* dtor_args;
};
static KeyInfo s_key_info[KEYS_MAX] = {};

// For allocating keys.
static pthread_mutex_t s_key_mutex = PTHREAD_MUTEX_INITIALIZER;
static size_t nfreekey = 0;
static size_t nkey = 0;
static uint32_t s_free_keys[KEYS_MAX];

// Stats.
static butil::static_atomic<size_t> nkeytable = BUTIL_STATIC_ATOMIC_INIT(0);
static butil::static_atomic<size_t> nsubkeytable = BUTIL_STATIC_ATOMIC_INIT(0);

// The second-level array.
// Align with cacheline to avoid false sharing.
class BAIDU_CACHELINE_ALIGNMENT SubKeyTable {
public:
    SubKeyTable() {
        memset(_data, 0, sizeof(_data));
        nsubkeytable.fetch_add(1, butil::memory_order_relaxed);
    }

    // NOTE: Call clear first.
    ~SubKeyTable() {
        nsubkeytable.fetch_sub(1, butil::memory_order_relaxed);
    }

    void clear(uint32_t offset) {
        for (uint32_t i = 0; i < KEY_2NDLEVEL_SIZE; ++i) {
            void* p = _data[i].ptr;
            if (p) {
                // Set the position to NULL before calling dtor which may set
                // the position again.
                _data[i].ptr = NULL;

                KeyInfo info = bthread::s_key_info[offset + i];
                if (info.dtor && _data[i].version == info.version) {
                    info.dtor(p, info.dtor_args);
                }
            }
        }
    }

    bool cleared() const {
        // We need to iterate again to check if every slot is empty. An
        // alternative is remember if set_data() was called during clear.
        for (uint32_t i = 0; i < KEY_2NDLEVEL_SIZE; ++i) {
            if (_data[i].ptr) {
                return false;
            }
        }
        return true;
    }

    inline void* get_data(uint32_t index, uint32_t version) const {
        if (_data[index].version == version) {
            return _data[index].ptr;
        }
        return NULL;
    }
    inline void set_data(uint32_t index, uint32_t version, void* data) {
        _data[index].version = version;
        _data[index].ptr = data;
    }

private:
    struct Data {
        uint32_t version;
        void* ptr;
    };
    Data _data[KEY_2NDLEVEL_SIZE];
};

// The first-level array.
// Align with cacheline to avoid false sharing.
class BAIDU_CACHELINE_ALIGNMENT KeyTable {
public:
    KeyTable() : next(NULL) {
        memset(_subs, 0, sizeof(_subs));
        nkeytable.fetch_add(1, butil::memory_order_relaxed);
    }

    ~KeyTable() {
        nkeytable.fetch_sub(1, butil::memory_order_relaxed);
        for (int ntry = 0; ntry < PTHREAD_DESTRUCTOR_ITERATIONS; ++ntry) {
            for (uint32_t i = 0; i < KEY_1STLEVEL_SIZE; ++i) {
                if (_subs[i]) {
                    _subs[i]->clear(i * KEY_2NDLEVEL_SIZE);
                }
            }
            bool all_cleared = true;
            for (uint32_t i = 0; i < KEY_1STLEVEL_SIZE; ++i) {
                if (_subs[i] != NULL && !_subs[i]->cleared()) {
                    all_cleared = false;
                    break;
                }
            }
            if (all_cleared) {
                for (uint32_t i = 0; i < KEY_1STLEVEL_SIZE; ++i) {
                    delete _subs[i];
                }
                return;
            }
        }
        LOG(ERROR) << "Fail to destroy all objects in KeyTable[" << this << ']';
    }

    inline void* get_data(bthread_key_t key) const {
        const uint32_t subidx = key.index / KEY_2NDLEVEL_SIZE;
        if (subidx < KEY_1STLEVEL_SIZE) {
            const SubKeyTable* sub_kt = _subs[subidx];
            if (sub_kt) {
                return sub_kt->get_data(
                    key.index - subidx * KEY_2NDLEVEL_SIZE, key.version);
            }
        }
        return NULL;
    }

    inline int set_data(bthread_key_t key, void* data) {
        const uint32_t subidx = key.index / KEY_2NDLEVEL_SIZE;
        if (subidx < KEY_1STLEVEL_SIZE &&
            key.version == s_key_info[key.index].version) {
            SubKeyTable* sub_kt = _subs[subidx];
            if (sub_kt == NULL) {
                sub_kt = new (std::nothrow) SubKeyTable;
                if (NULL == sub_kt) {
                    return ENOMEM;
                }
                _subs[subidx] = sub_kt;
            }
            sub_kt->set_data(key.index - subidx * KEY_2NDLEVEL_SIZE,
                             key.version, data);
            return 0;
        }
        CHECK(false) << "bthread_setspecific is called on invalid " << key;
        return EINVAL;
    }

public:
    KeyTable* next;
private:
    SubKeyTable* _subs[KEY_1STLEVEL_SIZE];
};

class BAIDU_CACHELINE_ALIGNMENT KeyTableList {
public:
    KeyTableList() :
        _head(NULL), _tail(NULL), _length(0) {}

    ~KeyTableList() {
        TaskGroup* g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
        KeyTable* old_kt = tls_bls.keytable;
        KeyTable* keytable = _head;
        while (keytable) {
            KeyTable* kt = keytable;
            keytable = kt->next;
            tls_bls.keytable = kt;
            if (g) {
                g->current_task()->local_storage.keytable = kt;
            }
            delete kt;
            if (old_kt == kt) {
                old_kt = NULL;
            }
            g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
        }
        tls_bls.keytable = old_kt;
        if (g) {
            g->current_task()->local_storage.keytable = old_kt;
        }
    }

    void append(KeyTable* keytable) {
        if (keytable == NULL) {
            return;
        }
        if (_head == NULL) {
            _head = _tail = keytable;
        } else {
            _tail->next = keytable;
            _tail = keytable;
        }
        keytable->next = NULL;
        _length++;
    }

    KeyTable* remove_front() {
        if (_head == NULL) {
            return NULL;
        }
        KeyTable* temp = _head;
        _head = _head->next;
        _length--;
        if (_head == NULL) {
            _tail = NULL;
        }
        return temp;
    }

    int move_first_n_to_target(KeyTable** target, uint32_t size) {
        if (size > _length || _head == NULL) {
            return 0;
        }

        KeyTable* current = _head;
        KeyTable* prev = NULL;
        uint32_t count = 0;
        while (current != NULL && count < size) {
            prev = current;
            current = current->next;
            count++;
        }
        if (prev != NULL) {
            if (*target == NULL) {
                *target = _head;
                prev->next = NULL;
            } else {
                prev->next = *target;
                *target = _head;
            }
            _head = current;
            _length -= count;
            if (_head == NULL) {
                _tail = NULL;
            }
        }
        return count;
    }

    inline uint32_t get_length() const {
        return _length;
    }

    // Only for test
    inline bool check_length() {
        KeyTable* current = _head;
        uint32_t count = 0;
        while (current != NULL) {
            current = current->next;
            count++;
        }
        return count == _length;
    }

private:
    KeyTable* _head;
    KeyTable* _tail;
    uint32_t _length;
};

KeyTable* borrow_keytable(bthread_keytable_pool_t* pool) {
    if (pool != NULL && (pool->list || pool->free_keytables)) {
        KeyTable* p;
        pthread_rwlock_rdlock(&pool->rwlock);
        auto list = (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
        if (list) {
            p = list->get()->remove_front();
            if (p) {
                pthread_rwlock_unlock(&pool->rwlock);
                return p;
            }
        }
        pthread_rwlock_unlock(&pool->rwlock);
        if (pool->free_keytables) {
            pthread_rwlock_wrlock(&pool->rwlock);
            p = (KeyTable*)pool->free_keytables;
            if (list) {
                for (uint32_t i = 0; i < FLAGS_borrow_from_globle_size; ++i) {
                    if (p) {
                        pool->free_keytables = p->next;
                        list->get()->append(p);
                        p = (KeyTable*)pool->free_keytables;
                        --pool->size;
                    } else {
                        break;
                    }
                }
                KeyTable* result = list->get()->remove_front();
                pthread_rwlock_unlock(&pool->rwlock);
                return result;
            } else {
                if (p) {
                    pool->free_keytables = p->next;
                    pthread_rwlock_unlock(&pool->rwlock);
                    return p;
                }
            }
            pthread_rwlock_unlock(&pool->rwlock);
        }
    }
    return NULL;
}

// Referenced in task_group.cpp, must be extern.
// Caller of this function must hold the KeyTable
void return_keytable(bthread_keytable_pool_t* pool, KeyTable* kt) {
    if (NULL == kt) {
        return;
    }
    if (pool == NULL) {
        delete kt;
        return;
    }
    pthread_rwlock_rdlock(&pool->rwlock);
    if (pool->destroyed) {
        pthread_rwlock_unlock(&pool->rwlock);
        delete kt;
        return;
    }
    auto list = (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
    list->get()->append(kt);
    if (list->get()->get_length() > FLAGS_key_table_list_size) {
        pthread_rwlock_unlock(&pool->rwlock);
        pthread_rwlock_wrlock(&pool->rwlock);
        if (!pool->destroyed) {
            int out = list->get()->move_first_n_to_target(
                (KeyTable**)(&pool->free_keytables),
                FLAGS_key_table_list_size / 2);
            pool->size += out;
        }
    }
    pthread_rwlock_unlock(&pool->rwlock);
}

static void cleanup_pthread(void* arg) {
    KeyTable* kt = static_cast<KeyTable*>(arg);
    if (kt) {
        delete kt;
        // After deletion: tls may be set during deletion.
        tls_bls.keytable = NULL;
    }
}

static void arg_as_dtor(void* data, const void* arg) {
    typedef void (*KeyDtor)(void*);
    return ((KeyDtor)arg)(data);
}

static int get_key_count(void*) {
    BAIDU_SCOPED_LOCK(bthread::s_key_mutex);
    return (int)nkey - (int)nfreekey;
}
static size_t get_keytable_count(void*) {
    return nkeytable.load(butil::memory_order_relaxed);
}
static size_t get_keytable_memory(void*) {
    const size_t n = nkeytable.load(butil::memory_order_relaxed);
    const size_t nsub = nsubkeytable.load(butil::memory_order_relaxed);
    return n * sizeof(KeyTable) + nsub * sizeof(SubKeyTable);
}

static bvar::PassiveStatus<int> s_bthread_key_count(
    "bthread_key_count", get_key_count, NULL);
static bvar::PassiveStatus<size_t> s_bthread_keytable_count(
    "bthread_keytable_count", get_keytable_count, NULL);
static bvar::PassiveStatus<size_t> s_bthread_keytable_memory(
    "bthread_keytable_memory", get_keytable_memory, NULL);

}  // namespace bthread

extern "C" {

int bthread_keytable_pool_init(bthread_keytable_pool_t* pool) {
    if (pool == NULL) {
        LOG(ERROR) << "Param[pool] is NULL";
        return EINVAL;
    }
    pthread_rwlock_init(&pool->rwlock, NULL);
    pool->list = new butil::ThreadLocal<bthread::KeyTableList>();
    pool->free_keytables = NULL;
    pool->size = 0;
    pool->destroyed = 0;
    return 0;
}

int bthread_keytable_pool_destroy(bthread_keytable_pool_t* pool) {
    if (pool == NULL) {
        LOG(ERROR) << "Param[pool] is NULL";
        return EINVAL;
    }
    bthread::KeyTable* saved_free_keytables = NULL;
    pthread_rwlock_wrlock(&pool->rwlock);
    pool->destroyed = 1;
    pool->size = 0;
    delete (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
    saved_free_keytables = (bthread::KeyTable*)pool->free_keytables;
    pool->list = NULL;
    pool->free_keytables = NULL;
    pthread_rwlock_unlock(&pool->rwlock);

    // Cheat get/setspecific and destroy the keytables.
    bthread::TaskGroup* g =
        bthread::BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
    bthread::KeyTable* old_kt = bthread::tls_bls.keytable;
    while (saved_free_keytables) {
        bthread::KeyTable* kt = saved_free_keytables;
        saved_free_keytables = kt->next;
        bthread::tls_bls.keytable = kt;
        if (g) {
            g->current_task()->local_storage.keytable = kt;
        }
        delete kt;
        g = bthread::BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
    }
    bthread::tls_bls.keytable = old_kt;
    if (g) {
        g->current_task()->local_storage.keytable = old_kt;
    }
    // TODO: return_keytable may race with this function, we don't destroy
    // the mutex right now.
    // pthread_mutex_destroy(&pool->mutex);
    return 0;
}

int bthread_keytable_pool_getstat(bthread_keytable_pool_t* pool,
                                  bthread_keytable_pool_stat_t* stat) {
    if (pool == NULL || stat == NULL) {
        LOG(ERROR) << "Param[pool] or Param[stat] is NULL";
        return EINVAL;
    }
    pthread_rwlock_wrlock(&pool->rwlock);
    stat->nfree = pool->size;
    pthread_rwlock_unlock(&pool->rwlock);
    return 0;
}

int get_thread_local_keytable_list_length(bthread_keytable_pool_t* pool) {
    if (pool == NULL) {
        LOG(ERROR) << "Param[pool] is NULL";
        return EINVAL;
    }
    int length = 0;
    pthread_rwlock_rdlock(&pool->rwlock);
    if (pool->destroyed) {
        pthread_rwlock_unlock(&pool->rwlock);
        return length;
    }
    auto list = (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
    if (list) {
        length = (int)(list->get()->get_length());
        if (!list->get()->check_length()) {
            LOG(ERROR) << "Length is not equal";
        }
    }
    pthread_rwlock_unlock(&pool->rwlock);
    return length;
}

// TODO: this is not strict `reserve' because we only check #free.
// Currently there's no way to track KeyTables that may be returned
// to the pool in future.
void bthread_keytable_pool_reserve(bthread_keytable_pool_t* pool,
                                   size_t nfree,
                                   bthread_key_t key,
                                   void* ctor(const void*),
                                   const void* ctor_args) {
    if (pool == NULL) {
        LOG(ERROR) << "Param[pool] is NULL";
        return;
    }
    bthread_keytable_pool_stat_t stat;
    if (bthread_keytable_pool_getstat(pool, &stat) != 0) {
        LOG(ERROR) << "Fail to getstat of pool=" << pool;
        return;
    }
    for (size_t i = stat.nfree; i < nfree; ++i) {
        bthread::KeyTable* kt = new (std::nothrow) bthread::KeyTable;
        if (kt == NULL) {
            break;
        }
        void* data = ctor(ctor_args);
        if (data) {
            kt->set_data(key, data);
        }  // else append kt w/o data.

        pthread_rwlock_wrlock(&pool->rwlock);
        if (pool->destroyed) {
            pthread_rwlock_unlock(&pool->rwlock);
            delete kt;
            break;
        }
        kt->next = (bthread::KeyTable*)pool->free_keytables;
        pool->free_keytables = kt;
        ++pool->size;
        pthread_rwlock_unlock(&pool->rwlock);
        if (data == NULL) {
            break;
        }
    }
}

int bthread_key_create2(bthread_key_t* key,
                        void (*dtor)(void*, const void*),
                        const void* dtor_args) {
    uint32_t index = 0;
    {
        BAIDU_SCOPED_LOCK(bthread::s_key_mutex);
        if (bthread::nfreekey > 0) {
            index = bthread::s_free_keys[--bthread::nfreekey];
        } else if (bthread::nkey < bthread::KEYS_MAX) {
            index = bthread::nkey++;
        } else {
            return EAGAIN;  // what pthread_key_create returns in this case.
        }
    }
    bthread::s_key_info[index].dtor = dtor;
    bthread::s_key_info[index].dtor_args = dtor_args;
    key->index = index;
    key->version = bthread::s_key_info[index].version;
    if (key->version == 0) {
        ++bthread::s_key_info[index].version;
        ++key->version;
    }
    return 0;
}

int bthread_key_create(bthread_key_t* key, void (*dtor)(void*)) {
    if (dtor == NULL) {
        return bthread_key_create2(key, NULL, NULL);
    } else {
        return bthread_key_create2(key, bthread::arg_as_dtor, (const void*)dtor);
    }
}

int bthread_key_delete(bthread_key_t key) {
    if (key.index < bthread::KEYS_MAX &&
        key.version == bthread::s_key_info[key.index].version) {
        BAIDU_SCOPED_LOCK(bthread::s_key_mutex);
        if (key.version == bthread::s_key_info[key.index].version) {
            if (++bthread::s_key_info[key.index].version == 0) {
                ++bthread::s_key_info[key.index].version;
            }
            bthread::s_key_info[key.index].dtor = NULL;
            bthread::s_key_info[key.index].dtor_args = NULL;
            bthread::s_free_keys[bthread::nfreekey++] = key.index;
            return 0;
        }
    }
    CHECK(false) << "bthread_key_delete is called on invalid " << key;
    return EINVAL;
}

// NOTE: Can't borrow_keytable in bthread_setspecific, otherwise following
// memory leak may occur:
//  -> bthread_getspecific fails to borrow_keytable and returns NULL.
//  -> bthread_setspecific succeeds to borrow_keytable and overwrites old data
//     at the position with newly created data, the old data is leaked.
int bthread_setspecific(bthread_key_t key, void* data) {
    bthread::KeyTable* kt = bthread::tls_bls.keytable;
    if (NULL == kt) {
        kt = new (std::nothrow) bthread::KeyTable;
        if (NULL == kt) {
            return ENOMEM;
        }
        bthread::tls_bls.keytable = kt;
        bthread::TaskGroup* const g = bthread::BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
        if (g) {
            g->current_task()->local_storage.keytable = kt;
        } else {
            // Only cleanup keytable created by pthread.
            // keytable created by bthread will be deleted
            // in `return_keytable' or `bthread_keytable_pool_destroy'.
            if (!bthread::tls_ever_created_keytable) {
                bthread::tls_ever_created_keytable = true;
                CHECK_EQ(0, butil::thread_atexit(bthread::cleanup_pthread, kt));
            }
        }
    }
    return kt->set_data(key, data);
}

void* bthread_getspecific(bthread_key_t key) {
    bthread::KeyTable* kt = bthread::tls_bls.keytable;
    if (kt) {
        return kt->get_data(key);
    }
    bthread::TaskGroup* const g = bthread::BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
    if (g) {
        bthread::TaskMeta* const task = g->current_task();
        kt = bthread::borrow_keytable(task->attr.keytable_pool);
        if (kt) {
            g->current_task()->local_storage.keytable = kt;
            bthread::tls_bls.keytable = kt;
            return kt->get_data(key);
        }
    }
    return NULL;
}

void bthread_assign_data(void* data) {
    bthread::tls_bls.assigned_data = data;
}

void* bthread_get_assigned_data() {
    return bthread::tls_bls.assigned_data;
}

}  // extern "C"

Coverage Report

Created: 2026-07-03 06:49

Line	Count	Source
1		// Licensed to the Apache Software Foundation (ASF) under one
2		// or more contributor license agreements. See the NOTICE file
3		// distributed with this work for additional information
4		// regarding copyright ownership. The ASF licenses this file
5		// to you under the Apache License, Version 2.0 (the
6		// "License"); you may not use this file except in compliance
7		// with the License. You may obtain a copy of the License at
8		//
9		// http://www.apache.org/licenses/LICENSE-2.0
10		//
11		// Unless required by applicable law or agreed to in writing,
12		// software distributed under the License is distributed on an
13		// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14		// KIND, either express or implied. See the License for the
15		// specific language governing permissions and limitations
16		// under the License.
17
18		// bthread - An M:N threading library to make applications more concurrent.
19
20		// Date: Sun Aug 3 12:46:15 CST 2014
21
22		#include <pthread.h>
23		#include <gflags/gflags.h>
24
25		#include "bthread/bthread.h" // bthread_create_span_fn and related types
26		#include "bthread/errno.h" // EAGAIN
27		#include "bthread/task_group.h" // TaskGroup
28		#include "butil/atomicops.h"
29		#include "butil/macros.h"
30		#include "butil/thread_key.h"
31		#include "butil/thread_local.h"
32		#include "bvar/passive_status.h"
33
34		// Implement bthread_key_t related functions
35
36		namespace bthread {
37
38		DEFINE_uint32(key_table_list_size, 4000,
39		"The maximum length of the KeyTableList. Once this value is "
40		"exceeded, a portion of the KeyTables will be moved to the "
41		"global free_keytables list.");
42
43		DEFINE_uint32(borrow_from_globle_size, 200,
44		"The maximum number of KeyTables retrieved in a single operation "
45		"from the global free_keytables when no KeyTable exists in the "
46		"current thread's keytable_list.");
47
48		EXTERN_BAIDU_VOLATILE_THREAD_LOCAL(TaskGroup*, tls_task_group);
49
50		class KeyTable;
51
52		// defined in task_group.cpp
53		extern __thread LocalStorage tls_bls;
54		static __thread bool tls_ever_created_keytable = false;
55
56		// We keep thread specific data in a two-level array. The top-level array
57		// contains at most KEY_1STLEVEL_SIZE pointers to dynamically allocated
58		// arrays of at most KEY_2NDLEVEL_SIZE data pointers. Many applications
59		// may just occupy one or two second level array, thus this machanism keeps
60		// memory footprint smaller and we can change KEY_1STLEVEL_SIZE to a
61		// bigger number more freely. The tradeoff is an additional memory indirection:
62		// negligible at most time.
63		static const uint32_t KEY_2NDLEVEL_SIZE = 32;
64
65		// Notice that we're trying to make the memory of second level and first
66		// level both 256 bytes to make memory allocator happier.
67		static const uint32_t KEY_1STLEVEL_SIZE = 31;
68
69		// Max tls in one thread, currently the value is 992 which should be enough
70		// for most projects throughout baidu.
71		static const uint32_t KEYS_MAX = KEY_2NDLEVEL_SIZE * KEY_1STLEVEL_SIZE;
72
73		// destructors/version of TLS.
74		struct KeyInfo {
75		uint32_t version;
76		void (dtor)(void, const void*);
77		const void* dtor_args;
78		};
79		static KeyInfo s_key_info[KEYS_MAX] = {};
80
81		// For allocating keys.
82		static pthread_mutex_t s_key_mutex = PTHREAD_MUTEX_INITIALIZER;
83		static size_t nfreekey = 0;
84		static size_t nkey = 0;
85		static uint32_t s_free_keys[KEYS_MAX];
86
87		// Stats.
88		static butil::static_atomic<size_t> nkeytable = BUTIL_STATIC_ATOMIC_INIT(0);
89		static butil::static_atomic<size_t> nsubkeytable = BUTIL_STATIC_ATOMIC_INIT(0);
90
91		// The second-level array.
92		// Align with cacheline to avoid false sharing.
93		class BAIDU_CACHELINE_ALIGNMENT SubKeyTable {
94		public:
95	1	SubKeyTable() {
96	1	memset(_data, 0, sizeof(_data));
97	1	nsubkeytable.fetch_add(1, butil::memory_order_relaxed);
98	1	}
99
100		// NOTE: Call clear first.
101	0	~SubKeyTable() {
102	0	nsubkeytable.fetch_sub(1, butil::memory_order_relaxed);
103	0	}
104
105	0	void clear(uint32_t offset) {
106	0	for (uint32_t i = 0; i < KEY_2NDLEVEL_SIZE; ++i) {
107	0	void* p = _data[i].ptr;
108	0	if (p) {
109		// Set the position to NULL before calling dtor which may set
110		// the position again.
111	0	_data[i].ptr = NULL;
112
113	0	KeyInfo info = bthread::s_key_info[offset + i];
114	0	if (info.dtor && _data[i].version == info.version) {
115	0	info.dtor(p, info.dtor_args);
116	0	}
117	0	}
118	0	}
119	0	}
120
121	0	bool cleared() const {
122		// We need to iterate again to check if every slot is empty. An
123		// alternative is remember if set_data() was called during clear.
124	0	for (uint32_t i = 0; i < KEY_2NDLEVEL_SIZE; ++i) {
125	0	if (_data[i].ptr) {
126	0	return false;
127	0	}
128	0	}
129	0	return true;
130	0	}
131
132	1	inline void* get_data(uint32_t index, uint32_t version) const {
133	1	if (_data[index].version == version) {
134	1	return _data[index].ptr;
135	1	}
136	0	return NULL;
137	1	}
138	1	inline void set_data(uint32_t index, uint32_t version, void* data) {
139	1	_data[index].version = version;
140	1	_data[index].ptr = data;
141	1	}
142
143		private:
144		struct Data {
145		uint32_t version;
146		void* ptr;
147		};
148		Data _data[KEY_2NDLEVEL_SIZE];
149		};
150
151		// The first-level array.
152		// Align with cacheline to avoid false sharing.
153		class BAIDU_CACHELINE_ALIGNMENT KeyTable {
154		public:
155	1	KeyTable() : next(NULL) {
156	1	memset(_subs, 0, sizeof(_subs));
157	1	nkeytable.fetch_add(1, butil::memory_order_relaxed);
158	1	}
159
160	0	~KeyTable() {
161	0	nkeytable.fetch_sub(1, butil::memory_order_relaxed);
162	0	for (int ntry = 0; ntry < PTHREAD_DESTRUCTOR_ITERATIONS; ++ntry) {
163	0	for (uint32_t i = 0; i < KEY_1STLEVEL_SIZE; ++i) {
164	0	if (_subs[i]) {
165	0	_subs[i]->clear(i * KEY_2NDLEVEL_SIZE);
166	0	}
167	0	}
168	0	bool all_cleared = true;
169	0	for (uint32_t i = 0; i < KEY_1STLEVEL_SIZE; ++i) {
170	0	if (_subs[i] != NULL && !_subs[i]->cleared()) {
171	0	all_cleared = false;
172	0	break;
173	0	}
174	0	}
175	0	if (all_cleared) {
176	0	for (uint32_t i = 0; i < KEY_1STLEVEL_SIZE; ++i) {
177	0	delete _subs[i];
178	0	}
179	0	return;
180	0	}
181	0	}
182	0	LOG(ERROR) << "Fail to destroy all objects in KeyTable[" << this << ']';
183	0	}
184
185	1	inline void* get_data(bthread_key_t key) const {
186	1	const uint32_t subidx = key.index / KEY_2NDLEVEL_SIZE;
187	1	if (subidx < KEY_1STLEVEL_SIZE) {
188	1	const SubKeyTable* sub_kt = _subs[subidx];
189	1	if (sub_kt) {
190	1	return sub_kt->get_data(
191	1	key.index - subidx * KEY_2NDLEVEL_SIZE, key.version);
192	1	}
193	1	}
194	0	return NULL;
195	1	}
196
197	1	inline int set_data(bthread_key_t key, void* data) {
198	1	const uint32_t subidx = key.index / KEY_2NDLEVEL_SIZE;
199	1	if (subidx < KEY_1STLEVEL_SIZE &&
200	1	key.version == s_key_info[key.index].version) {
201	1	SubKeyTable* sub_kt = _subs[subidx];
202	1	if (sub_kt == NULL) {
203	1	sub_kt = new (std::nothrow) SubKeyTable;
204	1	if (NULL == sub_kt) {
205	0	return ENOMEM;
206	0	}
207	1	_subs[subidx] = sub_kt;
208	1	}
209	1	sub_kt->set_data(key.index - subidx * KEY_2NDLEVEL_SIZE,
210	1	key.version, data);
211	1	return 0;
212	1	}
213	0	CHECK(false) << "bthread_setspecific is called on invalid " << key;
214	0	return EINVAL;
215	1	}
216
217		public:
218		KeyTable* next;
219		private:
220		SubKeyTable* _subs[KEY_1STLEVEL_SIZE];
221		};
222
223		class BAIDU_CACHELINE_ALIGNMENT KeyTableList {
224		public:
225		KeyTableList() :
226	0	_head(NULL), _tail(NULL), _length(0) {}
227
228	0	~KeyTableList() {
229	0	TaskGroup* g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
230	0	KeyTable* old_kt = tls_bls.keytable;
231	0	KeyTable* keytable = _head;
232	0	while (keytable) {
233	0	KeyTable* kt = keytable;
234	0	keytable = kt->next;
235	0	tls_bls.keytable = kt;
236	0	if (g) {
237	0	g->current_task()->local_storage.keytable = kt;
238	0	}
239	0	delete kt;
240	0	if (old_kt == kt) {
241	0	old_kt = NULL;
242	0	}
243	0	g = BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
244	0	}
245	0	tls_bls.keytable = old_kt;
246	0	if (g) {
247	0	g->current_task()->local_storage.keytable = old_kt;
248	0	}
249	0	}
250
251	0	void append(KeyTable* keytable) {
252	0	if (keytable == NULL) {
253	0	return;
254	0	}
255	0	if (_head == NULL) {
256	0	_head = _tail = keytable;
257	0	} else {
258	0	_tail->next = keytable;
259	0	_tail = keytable;
260	0	}
261	0	keytable->next = NULL;
262	0	_length++;
263	0	}
264
265	0	KeyTable* remove_front() {
266	0	if (_head == NULL) {
267	0	return NULL;
268	0	}
269	0	KeyTable* temp = _head;
270	0	_head = _head->next;
271	0	_length--;
272	0	if (_head == NULL) {
273	0	_tail = NULL;
274	0	}
275	0	return temp;
276	0	}
277
278	0	int move_first_n_to_target(KeyTable** target, uint32_t size) {
279	0	if (size > _length \|\| _head == NULL) {
280	0	return 0;
281	0	}
282
283	0	KeyTable* current = _head;
284	0	KeyTable* prev = NULL;
285	0	uint32_t count = 0;
286	0	while (current != NULL && count < size) {
287	0	prev = current;
288	0	current = current->next;
289	0	count++;
290	0	}
291	0	if (prev != NULL) {
292	0	if (*target == NULL) {
293	0	*target = _head;
294	0	prev->next = NULL;
295	0	} else {
296	0	prev->next = *target;
297	0	*target = _head;
298	0	}
299	0	_head = current;
300	0	_length -= count;
301	0	if (_head == NULL) {
302	0	_tail = NULL;
303	0	}
304	0	}
305	0	return count;
306	0	}
307
308	0	inline uint32_t get_length() const {
309	0	return _length;
310	0	}
311
312		// Only for test
313	0	inline bool check_length() {
314	0	KeyTable* current = _head;
315	0	uint32_t count = 0;
316	0	while (current != NULL) {
317	0	current = current->next;
318	0	count++;
319	0	}
320	0	return count == _length;
321	0	}
322
323		private:
324		KeyTable* _head;
325		KeyTable* _tail;
326		uint32_t _length;
327		};
328
329	0	KeyTable* borrow_keytable(bthread_keytable_pool_t* pool) {
330	0	if (pool != NULL && (pool->list \|\| pool->free_keytables)) {
331	0	KeyTable* p;
332	0	pthread_rwlock_rdlock(&pool->rwlock);
333	0	auto list = (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
334	0	if (list) {
335	0	p = list->get()->remove_front();
336	0	if (p) {
337	0	pthread_rwlock_unlock(&pool->rwlock);
338	0	return p;
339	0	}
340	0	}
341	0	pthread_rwlock_unlock(&pool->rwlock);
342	0	if (pool->free_keytables) {
343	0	pthread_rwlock_wrlock(&pool->rwlock);
344	0	p = (KeyTable*)pool->free_keytables;
345	0	if (list) {
346	0	for (uint32_t i = 0; i < FLAGS_borrow_from_globle_size; ++i) {
347	0	if (p) {
348	0	pool->free_keytables = p->next;
349	0	list->get()->append(p);
350	0	p = (KeyTable*)pool->free_keytables;
351	0	--pool->size;
352	0	} else {
353	0	break;
354	0	}
355	0	}
356	0	KeyTable* result = list->get()->remove_front();
357	0	pthread_rwlock_unlock(&pool->rwlock);
358	0	return result;
359	0	} else {
360	0	if (p) {
361	0	pool->free_keytables = p->next;
362	0	pthread_rwlock_unlock(&pool->rwlock);
363	0	return p;
364	0	}
365	0	}
366	0	pthread_rwlock_unlock(&pool->rwlock);
367	0	}
368	0	}
369	0	return NULL;
370	0	}
371
372		// Referenced in task_group.cpp, must be extern.
373		// Caller of this function must hold the KeyTable
374	0	void return_keytable(bthread_keytable_pool_t* pool, KeyTable* kt) {
375	0	if (NULL == kt) {
376	0	return;
377	0	}
378	0	if (pool == NULL) {
379	0	delete kt;
380	0	return;
381	0	}
382	0	pthread_rwlock_rdlock(&pool->rwlock);
383	0	if (pool->destroyed) {
384	0	pthread_rwlock_unlock(&pool->rwlock);
385	0	delete kt;
386	0	return;
387	0	}
388	0	auto list = (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
389	0	list->get()->append(kt);
390	0	if (list->get()->get_length() > FLAGS_key_table_list_size) {
391	0	pthread_rwlock_unlock(&pool->rwlock);
392	0	pthread_rwlock_wrlock(&pool->rwlock);
393	0	if (!pool->destroyed) {
394	0	int out = list->get()->move_first_n_to_target(
395	0	(KeyTable**)(&pool->free_keytables),
396	0	FLAGS_key_table_list_size / 2);
397	0	pool->size += out;
398	0	}
399	0	}
400	0	pthread_rwlock_unlock(&pool->rwlock);
401	0	}
402
403	0	static void cleanup_pthread(void* arg) {
404	0	KeyTable* kt = static_cast<KeyTable*>(arg);
405	0	if (kt) {
406	0	delete kt;
407		// After deletion: tls may be set during deletion.
408	0	tls_bls.keytable = NULL;
409	0	}
410	0	}
411
412	0	static void arg_as_dtor(void* data, const void* arg) {
413	0	typedef void (KeyDtor)(void);
414	0	return ((KeyDtor)arg)(data);
415	0	}
416
417	6	static int get_key_count(void*) {
418	6	BAIDU_SCOPED_LOCK(bthread::s_key_mutex);
419	6	return (int)nkey - (int)nfreekey;
420	6	}
421	6	static size_t get_keytable_count(void*) {
422	6	return nkeytable.load(butil::memory_order_relaxed);
423	6	}
424	6	static size_t get_keytable_memory(void*) {
425	6	const size_t n = nkeytable.load(butil::memory_order_relaxed);
426	6	const size_t nsub = nsubkeytable.load(butil::memory_order_relaxed);
427	6	return n * sizeof(KeyTable) + nsub * sizeof(SubKeyTable);
428	6	}
429
430		static bvar::PassiveStatus<int> s_bthread_key_count(
431		"bthread_key_count", get_key_count, NULL);
432		static bvar::PassiveStatus<size_t> s_bthread_keytable_count(
433		"bthread_keytable_count", get_keytable_count, NULL);
434		static bvar::PassiveStatus<size_t> s_bthread_keytable_memory(
435		"bthread_keytable_memory", get_keytable_memory, NULL);
436
437		} // namespace bthread
438
439		extern "C" {
440
441	0	int bthread_keytable_pool_init(bthread_keytable_pool_t* pool) {
442	0	if (pool == NULL) {
443	0	LOG(ERROR) << "Param[pool] is NULL";
444	0	return EINVAL;
445	0	}
446	0	pthread_rwlock_init(&pool->rwlock, NULL);
447	0	pool->list = new butil::ThreadLocal<bthread::KeyTableList>();
448	0	pool->free_keytables = NULL;
449	0	pool->size = 0;
450	0	pool->destroyed = 0;
451	0	return 0;
452	0	}
453
454	0	int bthread_keytable_pool_destroy(bthread_keytable_pool_t* pool) {
455	0	if (pool == NULL) {
456	0	LOG(ERROR) << "Param[pool] is NULL";
457	0	return EINVAL;
458	0	}
459	0	bthread::KeyTable* saved_free_keytables = NULL;
460	0	pthread_rwlock_wrlock(&pool->rwlock);
461	0	pool->destroyed = 1;
462	0	pool->size = 0;
463	0	delete (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
464	0	saved_free_keytables = (bthread::KeyTable*)pool->free_keytables;
465	0	pool->list = NULL;
466	0	pool->free_keytables = NULL;
467	0	pthread_rwlock_unlock(&pool->rwlock);
468
469		// Cheat get/setspecific and destroy the keytables.
470	0	bthread::TaskGroup* g =
471	0	bthread::BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
472	0	bthread::KeyTable* old_kt = bthread::tls_bls.keytable;
473	0	while (saved_free_keytables) {
474	0	bthread::KeyTable* kt = saved_free_keytables;
475	0	saved_free_keytables = kt->next;
476	0	bthread::tls_bls.keytable = kt;
477	0	if (g) {
478	0	g->current_task()->local_storage.keytable = kt;
479	0	}
480	0	delete kt;
481	0	g = bthread::BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
482	0	}
483	0	bthread::tls_bls.keytable = old_kt;
484	0	if (g) {
485	0	g->current_task()->local_storage.keytable = old_kt;
486	0	}
487		// TODO: return_keytable may race with this function, we don't destroy
488		// the mutex right now.
489		// pthread_mutex_destroy(&pool->mutex);
490	0	return 0;
491	0	}
492
493		int bthread_keytable_pool_getstat(bthread_keytable_pool_t* pool,
494	0	bthread_keytable_pool_stat_t* stat) {
495	0	if (pool == NULL \|\| stat == NULL) {
496	0	LOG(ERROR) << "Param[pool] or Param[stat] is NULL";
497	0	return EINVAL;
498	0	}
499	0	pthread_rwlock_wrlock(&pool->rwlock);
500	0	stat->nfree = pool->size;
501	0	pthread_rwlock_unlock(&pool->rwlock);
502	0	return 0;
503	0	}
504
505	0	int get_thread_local_keytable_list_length(bthread_keytable_pool_t* pool) {
506	0	if (pool == NULL) {
507	0	LOG(ERROR) << "Param[pool] is NULL";
508	0	return EINVAL;
509	0	}
510	0	int length = 0;
511	0	pthread_rwlock_rdlock(&pool->rwlock);
512	0	if (pool->destroyed) {
513	0	pthread_rwlock_unlock(&pool->rwlock);
514	0	return length;
515	0	}
516	0	auto list = (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
517	0	if (list) {
518	0	length = (int)(list->get()->get_length());
519	0	if (!list->get()->check_length()) {
520	0	LOG(ERROR) << "Length is not equal";
521	0	}
522	0	}
523	0	pthread_rwlock_unlock(&pool->rwlock);
524	0	return length;
525	0	}
526
527		// TODO: this is not strict `reserve' because we only check #free.
528		// Currently there's no way to track KeyTables that may be returned
529		// to the pool in future.
530		void bthread_keytable_pool_reserve(bthread_keytable_pool_t* pool,
531		size_t nfree,
532		bthread_key_t key,
533		void* ctor(const void*),
534	0	const void* ctor_args) {
535	0	if (pool == NULL) {
536	0	LOG(ERROR) << "Param[pool] is NULL";
537	0	return;
538	0	}
539	0	bthread_keytable_pool_stat_t stat;
540	0	if (bthread_keytable_pool_getstat(pool, &stat) != 0) {
541	0	LOG(ERROR) << "Fail to getstat of pool=" << pool;
542	0	return;
543	0	}
544	0	for (size_t i = stat.nfree; i < nfree; ++i) {
545	0	bthread::KeyTable* kt = new (std::nothrow) bthread::KeyTable;
546	0	if (kt == NULL) {
547	0	break;
548	0	}
549	0	void* data = ctor(ctor_args);
550	0	if (data) {
551	0	kt->set_data(key, data);
552	0	} // else append kt w/o data.
553
554	0	pthread_rwlock_wrlock(&pool->rwlock);
555	0	if (pool->destroyed) {
556	0	pthread_rwlock_unlock(&pool->rwlock);
557	0	delete kt;
558	0	break;
559	0	}
560	0	kt->next = (bthread::KeyTable*)pool->free_keytables;
561	0	pool->free_keytables = kt;
562	0	++pool->size;
563	0	pthread_rwlock_unlock(&pool->rwlock);
564	0	if (data == NULL) {
565	0	break;
566	0	}
567	0	}
568	0	}
569
570		int bthread_key_create2(bthread_key_t* key,
571		void (dtor)(void, const void*),
572	1	const void* dtor_args) {
573	1	uint32_t index = 0;
574	1	{
575	1	BAIDU_SCOPED_LOCK(bthread::s_key_mutex);
576	1	if (bthread::nfreekey > 0) {
577	0	index = bthread::s_free_keys[--bthread::nfreekey];
578	1	} else if (bthread::nkey < bthread::KEYS_MAX) {
579	1	index = bthread::nkey++;
580	1	} else {
581	0	return EAGAIN; // what pthread_key_create returns in this case.
582	0	}
583	1	}
584	1	bthread::s_key_info[index].dtor = dtor;
585	1	bthread::s_key_info[index].dtor_args = dtor_args;
586	1	key->index = index;
587	1	key->version = bthread::s_key_info[index].version;
588	1	if (key->version == 0) {
589	1	++bthread::s_key_info[index].version;
590	1	++key->version;
591	1	}
592	1	return 0;
593	1	}
594
595	1	int bthread_key_create(bthread_key_t* key, void (dtor)(void)) {
596	1	if (dtor == NULL) {
597	0	return bthread_key_create2(key, NULL, NULL);
598	1	} else {
599	1	return bthread_key_create2(key, bthread::arg_as_dtor, (const void*)dtor);
600	1	}
601	1	}
602
603	0	int bthread_key_delete(bthread_key_t key) {
604	0	if (key.index < bthread::KEYS_MAX &&
605	0	key.version == bthread::s_key_info[key.index].version) {
606	0	BAIDU_SCOPED_LOCK(bthread::s_key_mutex);
607	0	if (key.version == bthread::s_key_info[key.index].version) {
608	0	if (++bthread::s_key_info[key.index].version == 0) {
609	0	++bthread::s_key_info[key.index].version;
610	0	}
611	0	bthread::s_key_info[key.index].dtor = NULL;
612	0	bthread::s_key_info[key.index].dtor_args = NULL;
613	0	bthread::s_free_keys[bthread::nfreekey++] = key.index;
614	0	return 0;
615	0	}
616	0	}
617	0	CHECK(false) << "bthread_key_delete is called on invalid " << key;
618	0	return EINVAL;
619	0	}
620
621		// NOTE: Can't borrow_keytable in bthread_setspecific, otherwise following
622		// memory leak may occur:
623		// -> bthread_getspecific fails to borrow_keytable and returns NULL.
624		// -> bthread_setspecific succeeds to borrow_keytable and overwrites old data
625		// at the position with newly created data, the old data is leaked.
626	1	int bthread_setspecific(bthread_key_t key, void* data) {
627	1	bthread::KeyTable* kt = bthread::tls_bls.keytable;
628	1	if (NULL == kt) {
629	1	kt = new (std::nothrow) bthread::KeyTable;
630	1	if (NULL == kt) {
631	0	return ENOMEM;
632	0	}
633	1	bthread::tls_bls.keytable = kt;
634	1	bthread::TaskGroup* const g = bthread::BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
635	1	if (g) {
636	0	g->current_task()->local_storage.keytable = kt;
637	1	} else {
638		// Only cleanup keytable created by pthread.
639		// keytable created by bthread will be deleted
640		// in `return_keytable' or `bthread_keytable_pool_destroy'.
641	1	if (!bthread::tls_ever_created_keytable) {
642	1	bthread::tls_ever_created_keytable = true;
643	1	CHECK_EQ(0, butil::thread_atexit(bthread::cleanup_pthread, kt));
644	1	}
645	1	}
646	1	}
647	1	return kt->set_data(key, data);
648	1	}
649
650	2	void* bthread_getspecific(bthread_key_t key) {
651	2	bthread::KeyTable* kt = bthread::tls_bls.keytable;
652	2	if (kt) {
653	1	return kt->get_data(key);
654	1	}
655	1	bthread::TaskGroup* const g = bthread::BAIDU_GET_VOLATILE_THREAD_LOCAL(tls_task_group);
656	1	if (g) {
657	0	bthread::TaskMeta* const task = g->current_task();
658	0	kt = bthread::borrow_keytable(task->attr.keytable_pool);
659	0	if (kt) {
660	0	g->current_task()->local_storage.keytable = kt;
661	0	bthread::tls_bls.keytable = kt;
662	0	return kt->get_data(key);
663	0	}
664	0	}
665	1	return NULL;
666	1	}
667
668	0	void bthread_assign_data(void* data) {
669	0	bthread::tls_bls.assigned_data = data;
670	0	}
671
672	0	void* bthread_get_assigned_data() {
673	0	return bthread::tls_bls.assigned_data;
674	0	}
675
676		} // extern "C"