diff --git a/cyber/io/example/tcp_echo_client.cc b/cyber/io/example/tcp_echo_client.cc
index 263c76748fcfa7b8de46c5fc2e478a5ebef18e92..91dd456c187e967f565d1257371bc50fdd793327 100644
--- a/cyber/io/example/tcp_echo_client.cc
+++ b/cyber/io/example/tcp_echo_client.cc
@@ -64,6 +64,9 @@ int main(int argc, char* argv[]) {
           std::cout << "please enter a message (enter Ctrl+C to exit):"
                     << std::endl;
           std::getline(std::cin, user_input);
+          if (!apollo::cyber::OK()) {
+            break;
+          }
           if (user_input.empty()) {
             continue;
           }
@@ -103,4 +106,5 @@ int main(int argc, char* argv[]) {
       "echo_client");
 
   apollo::cyber::WaitForShutdown();
+  return 0;
 }
diff --git a/cyber/io/example/udp_echo_client.cc b/cyber/io/example/udp_echo_client.cc
index 8e5c949dea3312076af1e6bccabfa651eabe7ee0..baf2455111f831d590e53cba4610a973764c5d89 100644
--- a/cyber/io/example/udp_echo_client.cc
+++ b/cyber/io/example/udp_echo_client.cc
@@ -64,6 +64,9 @@ int main(int argc, char* argv[]) {
           std::cout << "please enter a message (enter Ctrl+C to exit):"
                     << std::endl;
           std::getline(std::cin, user_input);
+          if (!apollo::cyber::OK()) {
+            break;
+          }
           if (user_input.empty()) {
             continue;
           }
diff --git a/cyber/io/poll_handler.cc b/cyber/io/poll_handler.cc
index bde91857c00dc5e4eba43e12cdef8882a5f89187..cba7e9cd010ab7e47f5d615287e5d61728694568 100644
--- a/cyber/io/poll_handler.cc
+++ b/cyber/io/poll_handler.cc
@@ -17,6 +17,7 @@
 #include "cyber/io/poll_handler.h"
 #include "cyber/common/log.h"
 #include "cyber/io/poller.h"
+#include "cyber/scheduler/scheduler_factory.h"
 
 namespace apollo {
 namespace cyber {
@@ -105,7 +106,7 @@ void PollHandler::ResponseCallback(const PollResponse& rsp) {
   response_ = rsp;
 
   if (routine_->state() == RoutineState::IO_WAIT) {
-    routine_->SetUpdateFlag();
+    scheduler::Instance()->NotifyTask(routine_->id());
   }
 }
 
diff --git a/cyber/scheduler/policy/scheduler_choreography.cc b/cyber/scheduler/policy/scheduler_choreography.cc
index 6936aa158136df953adc2facce2e41c5ebf24808..c0e93410b75d66e7a616675cf51037688dbebdd5 100644
--- a/cyber/scheduler/policy/scheduler_choreography.cc
+++ b/cyber/scheduler/policy/scheduler_choreography.cc
@@ -258,7 +258,8 @@ bool SchedulerChoreography::NotifyProcessor(uint64_t crid) {
     auto it = id_cr_.find(crid);
     if (it != id_cr_.end()) {
       cr = it->second;
-      if (cr->state() == RoutineState::DATA_WAIT) {
+      if (cr->state() == RoutineState::DATA_WAIT ||
+          cr->state() == RoutineState::IO_WAIT) {
         cr->SetUpdateFlag();
       }
     } else {
diff --git a/cyber/scheduler/policy/scheduler_classic.cc b/cyber/scheduler/policy/scheduler_classic.cc
index a1fca03f2bf82c14cafc47fdf628b27ee4441f8c..93672888d0b8387045167d5fa30710e59bd9f9a4 100644
--- a/cyber/scheduler/policy/scheduler_classic.cc
+++ b/cyber/scheduler/policy/scheduler_classic.cc
@@ -173,7 +173,8 @@ bool SchedulerClassic::NotifyProcessor(uint64_t crid) {
     ReadLockGuard<AtomicRWLock> lk(id_cr_lock_);
     if (id_cr_.find(crid) != id_cr_.end()) {
       auto cr = id_cr_[crid];
-      if (cr->state() == RoutineState::DATA_WAIT) {
+      if (cr->state() == RoutineState::DATA_WAIT ||
+          cr->state() == RoutineState::IO_WAIT) {
         cr->SetUpdateFlag();
       }