src/butil/thread_key.h

// 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.

#ifndef  BUTIL_THREAD_KEY_H
#define  BUTIL_THREAD_KEY_H

#include <limits>
#include <pthread.h>
#include <stdlib.h>
#include <vector>
#include "butil/scoped_lock.h"
#include "butil/type_traits.h"

namespace butil {

typedef void (*DtorFunction)(void *);

class ThreadKey {
public:
    friend int thread_key_create(ThreadKey& thread_key, DtorFunction dtor);
    friend int thread_key_delete(ThreadKey& thread_key);
    friend int thread_setspecific(ThreadKey& thread_key, void* data);
    friend void* thread_getspecific(ThreadKey& thread_key);

    static constexpr size_t InvalidID = std::numeric_limits<size_t>::max();
    static constexpr size_t InitSeq = 0;

    constexpr ThreadKey() : _id(InvalidID), _seq(InitSeq) {}

    ~ThreadKey() {
        Reset();
    }

    ThreadKey(ThreadKey&& other) noexcept
        : _id(other._id)
        , _seq(other._seq) {
        other.Reset();
    }

    ThreadKey& operator=(ThreadKey&& other) noexcept;

    ThreadKey(const ThreadKey& other) = delete;
    ThreadKey& operator=(const ThreadKey& other) = delete;

    bool Valid() const;

    void Reset() {
        _id = InvalidID;
        _seq = InitSeq;
    }

private:
    size_t _id; // Key id.
    // Sequence number form g_thread_keys set in thread_key_create.
    size_t _seq;
};

struct ThreadKeyInfo {
    ThreadKeyInfo() : seq(0), dtor(NULL) {}

    size_t seq; // Already allocated?
    DtorFunction dtor; // Destruction routine.
};

struct ThreadKeyTLS {
    ThreadKeyTLS() : seq(0), data(NULL) {}

    // Sequence number form ThreadKey,
    // set in `thread_setspecific',
    // used to check if the key is valid in `thread_getspecific'.
    size_t seq;
    void* data; // User data.
};

// pthread_key_xxx implication without num limit of key.
int thread_key_create(ThreadKey& thread_key, DtorFunction dtor);
int thread_key_delete(ThreadKey& thread_key);
int thread_setspecific(ThreadKey& thread_key, void* data);
void* thread_getspecific(ThreadKey& thread_key);


template <typename T>
class ThreadLocal {
public:
    ThreadLocal() : ThreadLocal(false) {}

    explicit ThreadLocal(bool delete_on_thread_exit);

    ~ThreadLocal();

    // non-copyable
    ThreadLocal(const ThreadLocal&) = delete;
    ThreadLocal& operator=(const ThreadLocal&) = delete;

    T* get();

    T* operator->() { return get(); }

    T& operator*() { return *get(); }

    // Iterate through all thread local objects.
    // Callback, which must accept Args params and return void,
    // will be called under a thread lock.
    template <typename Callback>
    void for_each(Callback&& callback) {
        BAIDU_CASSERT(
            (is_result_void<Callback, T*>::value),
            "Callback must accept Args params and return void");
        BAIDU_SCOPED_LOCK(_mutex);
        for (auto ptr : ptrs) {
            callback(ptr);
        }
    }

    void reset(T* ptr);

    void reset() {
        reset(NULL);
    }

private:
    static void DefaultDtor(void* ptr) {
        if (ptr) {
            delete static_cast<T*>(ptr);
        }
    }

    ThreadKey _key;
    pthread_mutex_t _mutex;
    // All pointers of data allocated by the ThreadLocal.
    std::vector<T*> ptrs;
    // Delete data on thread exit or destructor of ThreadLocal.
    bool _delete_on_thread_exit;
};

template <typename T>
ThreadLocal<T>::ThreadLocal(bool delete_on_thread_exit)
        : _mutex(PTHREAD_MUTEX_INITIALIZER)
        , _delete_on_thread_exit(delete_on_thread_exit) {
    DtorFunction dtor = _delete_on_thread_exit ? DefaultDtor : NULL;
    thread_key_create(_key, dtor);
}


template <typename T>
ThreadLocal<T>::~ThreadLocal() {
    thread_key_delete(_key);
    if (!_delete_on_thread_exit) {
        BAIDU_SCOPED_LOCK(_mutex);
        for (auto ptr : ptrs) {
            DefaultDtor(ptr);
        }
    }
    pthread_mutex_destroy(&_mutex);
}

template <typename T>
T* ThreadLocal<T>::get() {
    T* ptr = static_cast<T*>(thread_getspecific(_key));
    if (!ptr) {
        ptr = new (std::nothrow) T;
        if (!ptr) {
            return NULL;
        }
        int rc = thread_setspecific(_key, ptr);
        if (rc != 0) {
            DefaultDtor(ptr);
            return NULL;
        }
        {
            BAIDU_SCOPED_LOCK(_mutex);
            ptrs.push_back(ptr);
        }
    }
    return ptr;
}

template <typename T>
void ThreadLocal<T>::reset(T* ptr) {
    T* old_ptr = get();
    if (ptr == old_ptr) {
        return;
    }
    if (thread_setspecific(_key, ptr) != 0) {
        return;
    }
    {
        BAIDU_SCOPED_LOCK(_mutex);
        if (ptr) {
            ptrs.push_back(ptr);
        }
        // Remove and delete old_ptr.
        if (old_ptr) {
            auto iter = std::remove(ptrs.begin(), ptrs.end(), old_ptr);
            if (iter != ptrs.end()) {
                ptrs.erase(iter, ptrs.end());
            }
            DefaultDtor(old_ptr);
        }
    }
}

}


#endif // BUTIL_THREAD_KEY_H