| 236 | } |
| 237 | |
| 238 | Status AsyncFeatureStoreMgr::AddTask(SparseTask* t) { |
| 239 | // NOTE(jiankebg.pt): No need atomic add here, maybe |
| 240 | // more then one Op will call the same thread(queue). |
| 241 | // TODO: Need excetly balance here ? |
| 242 | // |
| 243 | // uint64_t index = active_thread_index_.fetch_add(1); |
| 244 | uint64_t index = active_thread_index_++; |
| 245 | index %= thread_num_; |
| 246 | bool ret = false; |
| 247 | { |
| 248 | static std::mutex mu; |
| 249 | std::lock_guard<std::mutex> lock(mu); |
| 250 | ret = task_queues_[index].enqueue(t); |
| 251 | } |
| 252 | // TODO: should retry ? |
| 253 | if (!ret) { |
| 254 | return tensorflow::errors::Internal( |
| 255 | "can not enqueue task into the task_queues, index is ", |
| 256 | std::to_string(index)); |
| 257 | } |
| 258 | |
| 259 | if (sleeping_[index]) { |
| 260 | { // TODO: Need lock to promise the cv->wait(...) |
| 261 | // behavior in ThreadRun function. |
| 262 | std::lock_guard<std::mutex> lock(*mutex_[index]); |
| 263 | ready_[index] = true; |
| 264 | } |
| 265 | |
| 266 | cv_[index]->notify_all(); |
| 267 | } |
| 268 | |
| 269 | return Status::OK(); |
| 270 | } |
| 271 | |
| 272 | Status AsyncFeatureStoreMgr::AddUpdateTask(SparseTask* t) { |
| 273 | // TODO: Need excetly balance here ? |