fix lenet hang problem on windows

This commit is contained in:
xiefangqi 2020-05-29 07:14:06 +08:00
parent 6599cc1aca
commit 2c42665e90
2 changed files with 24 additions and 2 deletions

View File

@ -24,7 +24,9 @@ namespace dataset {
thread_local Task *gMyTask = nullptr;
void Task::operator()() {
#if !defined(_WIN32) && !defined(_WIN64)
gMyTask = this;
#endif
id_ = this_thread::get_id();
std::stringstream ss;
ss << id_;

View File

@ -87,7 +87,27 @@ void TaskManager::interrupt_all() noexcept {
(void)master_->Interrupt();
}
Task *TaskManager::FindMe() { return gMyTask; }
Task *TaskManager::FindMe() {
#if !defined(_WIN32) && !defined(_WIN64)
return gMyTask;
#else
TaskManager &tm = TaskManager::GetInstance();
SharedLock lock(&tm.lru_lock_);
auto id = this_thread::get_id();
auto tk = std::find_if(tm.lru_.begin(), tm.lru_.end(), [id](const Task &tk) { return tk.id_ == id; });
if (tk != tm.lru_.end()) {
return &(*tk);
}
// If we get here, either I am the watchdog or the master thread.
if (tm.master_->id_ == id) {
return tm.master_.get();
} else if (tm.watchdog_ != nullptr && tm.watchdog_->id_ == id) {
return tm.watchdog_;
}
MS_LOG(ERROR) << "Task not found.";
return nullptr;
#endif
}
TaskManager::TaskManager() try : global_interrupt_(0),
lru_(&Task::node),
@ -101,8 +121,8 @@ TaskManager::TaskManager() try : global_interrupt_(0),
master_->id_ = this_thread::get_id();
master_->running_ = true;
master_->is_master_ = true;
gMyTask = master_.get();
#if !defined(_WIN32) && !defined(_WIN64)
gMyTask = master_.get();
// Initialize the semaphore for the watchdog
errno_t rc = sem_init(&sem_, 0, 0);
if (rc == -1) {