[k2] fix error with pending updates (#1337)

Author: apolyakov

runtime-light/coroutine/awaitable.h

@@ -76,17 +76,24 @@ class async_stack_watcher_t {
 } // namespace awaitable_impl_
 
 class wait_for_update_t : awaitable_impl_::fork_id_watcher_t, awaitable_impl_::async_stack_watcher_t {
+public:
+  enum class update_kind : uint8_t { read, write };
+
+private:
   uint64_t stream_d;
+  update_kind upd_kind;
   SuspendToken suspend_token;
   awaitable_impl_::state state{awaitable_impl_::state::init};
 
 public:
-  explicit wait_for_update_t(uint64_t stream_d_) noexcept
+  explicit wait_for_update_t(uint64_t stream_d_, update_kind kind_) noexcept
       : stream_d(stream_d_),
+        upd_kind(kind_),
         suspend_token(std::noop_coroutine(), WaitEvent::UpdateOnStream{.stream_d = stream_d}) {}
 
   wait_for_update_t(wait_for_update_t&& other) noexcept
       : stream_d(std::exchange(other.stream_d, k2::INVALID_PLATFORM_DESCRIPTOR)),
+        upd_kind(other.upd_kind),
         suspend_token(std::exchange(other.suspend_token, std::make_pair(std::noop_coroutine(), WaitEvent::Rechedule{}))),
         state(std::exchange(other.state, awaitable_impl_::state::end)) {}
 
@@ -102,7 +109,12 @@ class wait_for_update_t : awaitable_impl_::fork_id_watcher_t, awaitable_impl_::a
 
   bool await_ready() noexcept {
     kphp::log::assertion(state == awaitable_impl_::state::init);
-    state = InstanceState::get().stream_updated(stream_d) ? awaitable_impl_::state::ready : awaitable_impl_::state::init;
+    k2::StreamStatus status{};
+    k2::stream_status(stream_d, std::addressof(status));
+    state = (upd_kind == update_kind::read && status.read_status == k2::IOStatus::IOAvailable) ||
+                    (upd_kind == update_kind::write && status.write_status == k2::IOStatus::IOAvailable)
+                ? awaitable_impl_::state::ready
+                : awaitable_impl_::state::init;
     return state == awaitable_impl_::state::ready;
   }
 

runtime-light/state/instance-state.cpp

@@ -271,8 +271,7 @@ void InstanceState::process_platform_updates() noexcept {
         case ScheduleStatus::Resumed: { // scheduler's resumed a coroutine waiting for update
           break;
         }
-        case ScheduleStatus::Skipped: { // no one is waiting for the event yet, so just save it
-          pending_updates_.insert(stream_d);
+        case ScheduleStatus::Skipped: { // no one is waiting for the event yet
           break;
         }
         case ScheduleStatus::Error: { // something bad's happened, stop execution
@@ -365,7 +364,6 @@ void InstanceState::release_stream(uint64_t stream_d) noexcept {
     standard_stream_ = k2::INVALID_PLATFORM_DESCRIPTOR;
   }
   opened_streams_.erase(stream_d);
-  pending_updates_.erase(stream_d); // also erase pending updates if exists
   k2::free_descriptor(stream_d);
   kphp::log::debug("released a stream {}", stream_d);
 }
@@ -377,6 +375,5 @@ void InstanceState::release_all_streams() noexcept {
     kphp::log::debug("released a stream {}", stream_d);
   }
   opened_streams_.clear();
-  pending_updates_.clear();
   incoming_streams_.clear();
 }

runtime-light/state/instance-state.h

@@ -85,9 +85,6 @@ struct InstanceState final : vk::not_copyable {
 
   void process_platform_updates() noexcept;
 
-  bool stream_updated(uint64_t stream_d) const noexcept {
-    return pending_updates_.contains(stream_d);
-  }
   const unordered_set<uint64_t>& opened_streams() const noexcept {
     return opened_streams_;
   }
@@ -151,7 +148,6 @@ struct InstanceState final : vk::not_copyable {
   uint64_t standard_stream_{k2::INVALID_PLATFORM_DESCRIPTOR};
   deque<uint64_t> incoming_streams_;
   unordered_set<uint64_t> opened_streams_;
-  unordered_set<uint64_t> pending_updates_;
 
   static constexpr auto INIT_INSTANCE_ALLOCATOR_SIZE = static_cast<size_t>(16U * 1024U * 1024U); // 16MB
 };

runtime-light/streams/streams.cpp

@@ -44,7 +44,7 @@ read_all_from_stream(uint64_t stream_d) noexcept {
       }
       buffer_size += k2::read(stream_d, STREAM_BATCH_SIZE, buffer + buffer_size);
     } else if (status.read_status == k2::IOStatus::IOBlocked) {
-      co_await wait_for_update_t{stream_d};
+      co_await wait_for_update_t{stream_d, wait_for_update_t::update_kind::read};
     }
 
   } while (status.read_status != k2::IOStatus::IOClosed);
@@ -96,7 +96,7 @@ kphp::coro::task<int32_t> read_exact_from_stream(uint64_t stream_d, char* buffer
     if (status.read_status == k2::IOStatus::IOAvailable) {
       read += k2::read(stream_d, len - read, buffer + read);
     } else if (status.read_status == k2::IOStatus::IOBlocked) {
-      co_await wait_for_update_t{stream_d};
+      co_await wait_for_update_t{stream_d, wait_for_update_t::update_kind::read};
     } else {
       co_return read;
     }
@@ -122,7 +122,7 @@ kphp::coro::task<int32_t> write_all_to_stream(uint64_t stream_d, const char* buf
     } else if (status.write_status == k2::IOStatus::IOAvailable) {
       written += k2::write(stream_d, len - written, buffer + written);
     } else if (status.write_status == k2::IOStatus::IOBlocked) {
-      co_await wait_for_update_t{stream_d};
+      co_await wait_for_update_t{stream_d, wait_for_update_t::update_kind::write};
     } else {
       kphp::log::warning("stream closed while writing. Wrote {}. Size {}. Stream {}", written, len, stream_d);
       co_return written;
@@ -173,7 +173,7 @@ kphp::coro::task<int32_t> write_exact_to_stream(uint64_t stream_d, const char* b
     } else if (status.write_status == k2::IOStatus::IOAvailable) {
       written += k2::write(stream_d, len - written, buffer + written);
     } else if (status.write_status == k2::IOStatus::IOBlocked) {
-      co_await wait_for_update_t{stream_d};
+      co_await wait_for_update_t{stream_d, wait_for_update_t::update_kind::write};
     } else {
       co_return written;
     }