Small improvements to the WASIp2 implementation of HttpOutgoingBody::…

…write (#162) Besides some simplifications, this turns some checks that we know must pass into asserts instead. Signed-off-by: Till Schneidereit <[email protected]>
bytecodealliance · Oct 18, 2024 · ab33b5e · ab33b5e
1 parent 114208c
commit ab33b5e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 46 deletions.
diff --git a/host-apis/wasi-0.2.0/host_api.cpp b/host-apis/wasi-0.2.0/host_api.cpp
@@ -566,22 +566,12 @@ Result<uint64_t> HttpOutgoingBody::capacity() {
   return Result<uint64_t>::ok(capacity);
 }
 
-Result<uint32_t> HttpOutgoingBody::write(const uint8_t *bytes, size_t len) {
-  auto res = capacity();
-  if (res.is_err()) {
-    // TODO: proper error handling for all 154 error codes.
-    return Result<uint32_t>::err(154);
-  }
-  auto capacity = res.unwrap();
-  auto bytes_to_write = std::min(len, static_cast<size_t>(capacity));
+void HttpOutgoingBody::write(const uint8_t *bytes, size_t len) {
+  MOZ_ASSERT(capacity().unwrap() >= len);
 
   auto *state = static_cast<OutgoingBodyHandle *>(this->handle_state_.get());
   Borrow<OutputStream> borrow(state->stream_handle_);
-  if (!write_to_outgoing_body(borrow, bytes, bytes_to_write)) {
-    return Result<uint32_t>::err(154);
-  }
-
-  return Result<uint32_t>::ok(bytes_to_write);
+  MOZ_RELEASE_ASSERT(write_to_outgoing_body(borrow, bytes, len));
 }
 
 Result<Void> HttpOutgoingBody::write_all(const uint8_t *bytes, size_t len) {
@@ -646,18 +636,15 @@ class BodyAppendTask final : public api::AsyncTask {
                           HttpOutgoingBody *outgoing_body,
                           api::TaskCompletionCallback completion_callback,
                           HandleObject callback_receiver)
-      : incoming_body_(incoming_body), outgoing_body_(outgoing_body), cb_(completion_callback) {
+      : incoming_body_(incoming_body), outgoing_body_(outgoing_body), cb_(completion_callback),
+        cb_receiver_(callback_receiver), state_(State::BlockedOnBoth) {
     auto res = incoming_body_->subscribe();
     MOZ_ASSERT(!res.is_err());
     incoming_pollable_ = res.unwrap();
 
     res = outgoing_body_->subscribe();
     MOZ_ASSERT(!res.is_err());
     outgoing_pollable_ = res.unwrap();
-
-    cb_receiver_ = callback_receiver;
-
-    set_state(engine->cx(), State::BlockedOnBoth);
   }
 
   [[nodiscard]] bool run(api::Engine *engine) override {
@@ -674,19 +661,17 @@ class BodyAppendTask final : public api::AsyncTask {
       set_state(engine->cx(), State::BlockedOnOutgoing);
     }
 
-    uint64_t capacity = 0;
-    if (state_ == State::BlockedOnOutgoing) {
-      auto res = outgoing_body_->capacity();
-      if (res.is_err()) {
-        return false;
-      }
-      capacity = res.unwrap();
-      if (capacity > 0) {
-        set_state(engine->cx(), State::Ready);
-      } else {
-        engine->queue_async_task(this);
-        return true;
-      }
+    MOZ_ASSERT(state_ == State::BlockedOnOutgoing);
+    auto res = outgoing_body_->capacity();
+    if (res.is_err()) {
+      return false;
+    }
+    uint64_t capacity = res.unwrap();
+    if (capacity > 0) {
+      set_state(engine->cx(), State::Ready);
+    } else {
+      engine->queue_async_task(this);
+      return true;
     }
 
     MOZ_ASSERT(state_ == State::Ready);
@@ -705,17 +690,8 @@ class BodyAppendTask final : public api::AsyncTask {
         return true;
       }
 
-      unsigned offset = 0;
-      while (bytes.len - offset > 0) {
-        // TODO: remove double checking of write-readiness
-        // TODO: make this async by storing the remaining chunk in the task and marking it as
-        // being blocked on write
-        auto write_res = outgoing_body_->write(bytes.ptr.get() + offset, bytes.len - offset);
-        if (write_res.is_err()) {
-          // TODO: proper error handling.
-          return false;
-        }
-        offset += write_res.unwrap();
+      if (bytes.len > 0) {
+        outgoing_body_->write(bytes.ptr.get(), bytes.len);
       }
 
       if (done) {

diff --git a/include/host_api.h b/include/host_api.h
@@ -284,10 +284,10 @@ class HttpOutgoingBody final : public Pollable {
 
   /// Write a chunk to this handle.
   ///
-  /// Doesn't necessarily write the entire chunk, and doesn't take ownership of `bytes`.
-  ///
-  /// @return the number of bytes written.
-  Result<uint32_t> write(const uint8_t *bytes, size_t len);
+  /// Asserts that the receiver is ready to write the entire chunk, which the caller must ensure
+  /// by calling `receiver->capacity()` first and not attempting to write more than the returned
+  /// value.
+  void write(const uint8_t *bytes, size_t len);
 
   /// Writes the given number of bytes from the given buffer to the given handle.
   ///