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
|
#include "ThreadLocalPtr.h"
#include <algorithm>
namespace caffe2 {
// meyer's singleton
AllThreadLocalHelperVector* getAllThreadLocalHelperVector() {
// leak the pointer to avoid dealing with destruction order issues
static auto* instance = new AllThreadLocalHelperVector();
return instance;
}
ThreadLocalHelper* getThreadLocalHelper() {
static thread_local ThreadLocalHelper instance;
return &instance;
}
// AllThreadLocalHelperVector
void AllThreadLocalHelperVector::push_back(ThreadLocalHelper* helper) {
std::lock_guard<std::mutex> lg(mutex_);
vector_.push_back(helper);
}
void AllThreadLocalHelperVector::erase(ThreadLocalHelper* helper) {
std::lock_guard<std::mutex> lg(mutex_);
vector_.erase(
std::remove(vector_.begin(), vector_.end(), helper), vector_.end());
}
void AllThreadLocalHelperVector::erase_tlp(ThreadLocalPtrImpl* ptr) {
std::lock_guard<std::mutex> lg(mutex_);
for (auto* ins : vector_) {
ins->erase(ptr);
}
}
// ThreadLocalHelper
ThreadLocalHelper::ThreadLocalHelper() {
getAllThreadLocalHelperVector()->push_back(this);
}
ThreadLocalHelper::~ThreadLocalHelper() {
getAllThreadLocalHelperVector()->erase(this);
}
void ThreadLocalHelper::insert(
ThreadLocalPtrImpl* tl_ptr,
std::shared_ptr<void> ptr) {
std::lock_guard<std::mutex> lg(mutex_);
mapping_.insert(std::make_pair(tl_ptr, std::move(ptr)));
}
void* ThreadLocalHelper::get(ThreadLocalPtrImpl* key) {
/* Grabbing the mutex for the thread local map protecting the case
* when other object exits(~ThreadLocalPtrImpl()), and removes the
* element in the map, which will change the iterator returned
* by find.
*/
std::lock_guard<std::mutex> lg(mutex_);
auto it = mapping_.find(key);
if (it == mapping_.end()) {
return nullptr;
} else {
return it->second.get();
}
}
void ThreadLocalHelper::erase(ThreadLocalPtrImpl* key) {
std::lock_guard<std::mutex> lg(mutex_);
mapping_.erase(key);
}
ThreadLocalPtrImpl::~ThreadLocalPtrImpl() {
getAllThreadLocalHelperVector()->erase_tlp(this);
}
} // namespace caffe2
|
