| @@ -238,6 +238,9 @@ Status TFReaderOp::CalculateNumRowsPerShard() { | |||||
| Status TFReaderOp::operator()() { | Status TFReaderOp::operator()() { | ||||
| RETURN_IF_NOT_OK(CalculateNumRowsPerShard()); | RETURN_IF_NOT_OK(CalculateNumRowsPerShard()); | ||||
| // Put here to avoid register failed when Worker_Entry thread exits unexpected | |||||
| RETURN_IF_NOT_OK(io_block_queue_wait_post_.Register(tree_->AllTasks())); | |||||
| // launch one thread, responsible for filling mIOBlockQueue | // launch one thread, responsible for filling mIOBlockQueue | ||||
| RETURN_IF_NOT_OK(tree_->LaunchWorkers(1, std::bind(&TFReaderOp::WaitToFillIOBlockQueue, this))); | RETURN_IF_NOT_OK(tree_->LaunchWorkers(1, std::bind(&TFReaderOp::WaitToFillIOBlockQueue, this))); | ||||
| @@ -250,8 +253,6 @@ Status TFReaderOp::operator()() { | |||||
| // so workers have to be kept alive until the end of the program | // so workers have to be kept alive until the end of the program | ||||
| TaskManager::FindMe()->Post(); | TaskManager::FindMe()->Post(); | ||||
| RETURN_IF_NOT_OK(io_block_queue_wait_post_.Register(tree_->AllTasks())); | |||||
| NotifyToFillIOBlockQueue(); | NotifyToFillIOBlockQueue(); | ||||
| while (!finished_reading_dataset_) { | while (!finished_reading_dataset_) { | ||||
| int64_t buffer_id = 0; | int64_t buffer_id = 0; | ||||