1 // Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 
6 #include "monitoring/thread_status_util.h"
7 
8 #include "monitoring/thread_status_updater.h"
9 #include "rocksdb/env.h"
10 
11 namespace ROCKSDB_NAMESPACE {
12 
13 #ifdef ROCKSDB_USING_THREAD_STATUS
14 __thread ThreadStatusUpdater* ThreadStatusUtil::thread_updater_local_cache_ =
15     nullptr;
16 __thread bool ThreadStatusUtil::thread_updater_initialized_ = false;
17 
RegisterThread(const Env * env,ThreadStatus::ThreadType thread_type)18 void ThreadStatusUtil::RegisterThread(const Env* env,
19                                       ThreadStatus::ThreadType thread_type) {
20   if (!MaybeInitThreadLocalUpdater(env)) {
21     return;
22   }
23   assert(thread_updater_local_cache_);
24   thread_updater_local_cache_->RegisterThread(thread_type, env->GetThreadID());
25 }
26 
UnregisterThread()27 void ThreadStatusUtil::UnregisterThread() {
28   thread_updater_initialized_ = false;
29   if (thread_updater_local_cache_ != nullptr) {
30     thread_updater_local_cache_->UnregisterThread();
31     thread_updater_local_cache_ = nullptr;
32   }
33 }
34 
SetColumnFamily(const ColumnFamilyData * cfd,const Env * env,bool enable_thread_tracking)35 void ThreadStatusUtil::SetColumnFamily(const ColumnFamilyData* cfd,
36                                        const Env* env,
37                                        bool enable_thread_tracking) {
38   if (!MaybeInitThreadLocalUpdater(env)) {
39     return;
40   }
41   assert(thread_updater_local_cache_);
42   if (cfd != nullptr && enable_thread_tracking) {
43     thread_updater_local_cache_->SetColumnFamilyInfoKey(cfd);
44   } else {
45     // When cfd == nullptr or enable_thread_tracking == false, we set
46     // ColumnFamilyInfoKey to nullptr, which makes SetThreadOperation
47     // and SetThreadState become no-op.
48     thread_updater_local_cache_->SetColumnFamilyInfoKey(nullptr);
49   }
50 }
51 
SetThreadOperation(ThreadStatus::OperationType op)52 void ThreadStatusUtil::SetThreadOperation(ThreadStatus::OperationType op) {
53   if (thread_updater_local_cache_ == nullptr) {
54     // thread_updater_local_cache_ must be set in SetColumnFamily
55     // or other ThreadStatusUtil functions.
56     return;
57   }
58 
59   if (op != ThreadStatus::OP_UNKNOWN) {
60     uint64_t current_time = Env::Default()->NowMicros();
61     thread_updater_local_cache_->SetOperationStartTime(current_time);
62   } else {
63     // TDOO(yhchiang): we could report the time when we set operation to
64     // OP_UNKNOWN once the whole instrumentation has been done.
65     thread_updater_local_cache_->SetOperationStartTime(0);
66   }
67   thread_updater_local_cache_->SetThreadOperation(op);
68 }
69 
SetThreadOperationStage(ThreadStatus::OperationStage stage)70 ThreadStatus::OperationStage ThreadStatusUtil::SetThreadOperationStage(
71     ThreadStatus::OperationStage stage) {
72   if (thread_updater_local_cache_ == nullptr) {
73     // thread_updater_local_cache_ must be set in SetColumnFamily
74     // or other ThreadStatusUtil functions.
75     return ThreadStatus::STAGE_UNKNOWN;
76   }
77 
78   return thread_updater_local_cache_->SetThreadOperationStage(stage);
79 }
80 
SetThreadOperationProperty(int code,uint64_t value)81 void ThreadStatusUtil::SetThreadOperationProperty(int code, uint64_t value) {
82   if (thread_updater_local_cache_ == nullptr) {
83     // thread_updater_local_cache_ must be set in SetColumnFamily
84     // or other ThreadStatusUtil functions.
85     return;
86   }
87 
88   thread_updater_local_cache_->SetThreadOperationProperty(code, value);
89 }
90 
IncreaseThreadOperationProperty(int code,uint64_t delta)91 void ThreadStatusUtil::IncreaseThreadOperationProperty(int code,
92                                                        uint64_t delta) {
93   if (thread_updater_local_cache_ == nullptr) {
94     // thread_updater_local_cache_ must be set in SetColumnFamily
95     // or other ThreadStatusUtil functions.
96     return;
97   }
98 
99   thread_updater_local_cache_->IncreaseThreadOperationProperty(code, delta);
100 }
101 
SetThreadState(ThreadStatus::StateType state)102 void ThreadStatusUtil::SetThreadState(ThreadStatus::StateType state) {
103   if (thread_updater_local_cache_ == nullptr) {
104     // thread_updater_local_cache_ must be set in SetColumnFamily
105     // or other ThreadStatusUtil functions.
106     return;
107   }
108 
109   thread_updater_local_cache_->SetThreadState(state);
110 }
111 
ResetThreadStatus()112 void ThreadStatusUtil::ResetThreadStatus() {
113   if (thread_updater_local_cache_ == nullptr) {
114     return;
115   }
116   thread_updater_local_cache_->ResetThreadStatus();
117 }
118 
NewColumnFamilyInfo(const DB * db,const ColumnFamilyData * cfd,const std::string & cf_name,const Env * env)119 void ThreadStatusUtil::NewColumnFamilyInfo(const DB* db,
120                                            const ColumnFamilyData* cfd,
121                                            const std::string& cf_name,
122                                            const Env* env) {
123   if (!MaybeInitThreadLocalUpdater(env)) {
124     return;
125   }
126   assert(thread_updater_local_cache_);
127   if (thread_updater_local_cache_) {
128     thread_updater_local_cache_->NewColumnFamilyInfo(db, db->GetName(), cfd,
129                                                      cf_name);
130   }
131 }
132 
EraseColumnFamilyInfo(const ColumnFamilyData * cfd)133 void ThreadStatusUtil::EraseColumnFamilyInfo(const ColumnFamilyData* cfd) {
134   if (thread_updater_local_cache_ == nullptr) {
135     return;
136   }
137   thread_updater_local_cache_->EraseColumnFamilyInfo(cfd);
138 }
139 
EraseDatabaseInfo(const DB * db)140 void ThreadStatusUtil::EraseDatabaseInfo(const DB* db) {
141   ThreadStatusUpdater* thread_updater = db->GetEnv()->GetThreadStatusUpdater();
142   if (thread_updater == nullptr) {
143     return;
144   }
145   thread_updater->EraseDatabaseInfo(db);
146 }
147 
MaybeInitThreadLocalUpdater(const Env * env)148 bool ThreadStatusUtil::MaybeInitThreadLocalUpdater(const Env* env) {
149   if (!thread_updater_initialized_ && env != nullptr) {
150     thread_updater_initialized_ = true;
151     thread_updater_local_cache_ = env->GetThreadStatusUpdater();
152   }
153   return (thread_updater_local_cache_ != nullptr);
154 }
155 
AutoThreadOperationStageUpdater(ThreadStatus::OperationStage stage)156 AutoThreadOperationStageUpdater::AutoThreadOperationStageUpdater(
157     ThreadStatus::OperationStage stage) {
158   prev_stage_ = ThreadStatusUtil::SetThreadOperationStage(stage);
159 }
160 
~AutoThreadOperationStageUpdater()161 AutoThreadOperationStageUpdater::~AutoThreadOperationStageUpdater() {
162   ThreadStatusUtil::SetThreadOperationStage(prev_stage_);
163 }
164 
165 #else
166 
167 ThreadStatusUpdater* ThreadStatusUtil::thread_updater_local_cache_ = nullptr;
168 bool ThreadStatusUtil::thread_updater_initialized_ = false;
169 
170 bool ThreadStatusUtil::MaybeInitThreadLocalUpdater(const Env* /*env*/) {
171   return false;
172 }
173 
174 void ThreadStatusUtil::SetColumnFamily(const ColumnFamilyData* /*cfd*/,
175                                        const Env* /*env*/,
176                                        bool /*enable_thread_tracking*/) {}
177 
178 void ThreadStatusUtil::SetThreadOperation(ThreadStatus::OperationType /*op*/) {}
179 
180 void ThreadStatusUtil::SetThreadOperationProperty(int /*code*/,
181                                                   uint64_t /*value*/) {}
182 
183 void ThreadStatusUtil::IncreaseThreadOperationProperty(int /*code*/,
184                                                        uint64_t /*delta*/) {}
185 
186 void ThreadStatusUtil::SetThreadState(ThreadStatus::StateType /*state*/) {}
187 
188 void ThreadStatusUtil::NewColumnFamilyInfo(const DB* /*db*/,
189                                            const ColumnFamilyData* /*cfd*/,
190                                            const std::string& /*cf_name*/,
191                                            const Env* /*env*/) {}
192 
193 void ThreadStatusUtil::EraseColumnFamilyInfo(const ColumnFamilyData* /*cfd*/) {}
194 
195 void ThreadStatusUtil::EraseDatabaseInfo(const DB* /*db*/) {}
196 
197 void ThreadStatusUtil::ResetThreadStatus() {}
198 
199 AutoThreadOperationStageUpdater::AutoThreadOperationStageUpdater(
200     ThreadStatus::OperationStage /*stage*/) {}
201 
202 AutoThreadOperationStageUpdater::~AutoThreadOperationStageUpdater() {}
203 
204 #endif  // ROCKSDB_USING_THREAD_STATUS
205 
206 }  // namespace ROCKSDB_NAMESPACE
207