capnproto

byte-stream.c++ (40034B)

---

 // Copyright (c) 2019 Cloudflare, Inc. and contributors
 // Licensed under the MIT License:
 //
 // Permission is hereby granted, free of charge, to any person obtaining a copy
 // of this software and associated documentation files (the "Software"), to deal
 // in the Software without restriction, including without limitation the rights
 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 // copies of the Software, and to permit persons to whom the Software is
 // furnished to do so, subject to the following conditions:
 //
 // The above copyright notice and this permission notice shall be included in
 // all copies or substantial portions of the Software.
 //
 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
 #include "byte-stream.h"
 #include <kj/one-of.h>
 #include <kj/debug.h>
 
 namespace capnp {
 
 const uint MAX_BYTES_PER_WRITE = 1 << 16;
 
 class ByteStreamFactory::StreamServerBase: public capnp::ByteStream::Server {
 public:
   virtual void returnStream(uint64_t written) = 0;
   // Called after the StreamServerBase's internal kj::AsyncOutputStream has been borrowed, to
   // indicate that the borrower is done.
   //
   // A stream becomes borrowed either when getShortestPath() returns a BorrowedStream, or when
   // a SubstreamImpl is constructed wrapping an existing stream.
 
   struct BorrowedStream {
     // Represents permission to use the StreamServerBase's inner AsyncOutputStream directly, up
     // to some limit of bytes written.
 
     StreamServerBase& lender;
     kj::AsyncOutputStream& stream;
     uint64_t limit;
   };
 
   typedef kj::OneOf<kj::Promise<void>, capnp::ByteStream::Client*, BorrowedStream> ShortestPath;
 
   virtual ShortestPath getShortestPath() = 0;
   // Called by KjToCapnpStreamAdapter when it has determined that its inner ByteStream::Client
   // actually points back to a StreamServerBase in the same process created by the same
   // ByteStreamFactory. Returns the best shortened path to use, or a promise that resolves when the
   // shortest path is known.
 
   virtual void directEnd() = 0;
   // Called by KjToCapnpStreamAdapter's destructor when it has determined that its inner
   // ByteStream::Client actually points back to a StreamServerBase in the same process created by
   // the same ByteStreamFactory. Since destruction of a KJ stream signals EOF, we need to propagate
   // that by destroying our underlying stream.
   // TODO(cleanup): When KJ streams evolve an end() method, this can go away.
 };
 
 class ByteStreamFactory::SubstreamImpl final: public StreamServerBase {
 public:
   SubstreamImpl(ByteStreamFactory& factory,
                 StreamServerBase& parent,
                 capnp::ByteStream::Client ownParent,
                 kj::AsyncOutputStream& stream,
                 capnp::ByteStream::SubstreamCallback::Client callback,
                 uint64_t limit,
                 kj::PromiseFulfillerPair<void> paf = kj::newPromiseAndFulfiller<void>())
       : factory(factory),
         state(Streaming {parent, kj::mv(ownParent), stream, kj::mv(callback)}),
         limit(limit),
         resolveFulfiller(kj::mv(paf.fulfiller)),
         resolvePromise(paf.promise.fork()) {}
 
   // ---------------------------------------------------------------------------
   // implements StreamServerBase
 
   void returnStream(uint64_t written) override {
     completed += written;
     KJ_ASSERT(completed <= limit);
     auto borrowed = kj::mv(state.get<Borrowed>());
     state = kj::mv(borrowed.originalState);
 
     if (completed == limit) {
       limitReached();
     }
   }
 
   ShortestPath getShortestPath() override {
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(redirected, Redirected) {
         return &redirected.replacement;
       }
       KJ_CASE_ONEOF(e, Ended) {
         KJ_FAIL_REQUIRE("already called end()");
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("can't call other methods while substream is active");
       }
       KJ_CASE_ONEOF(streaming, Streaming) {
         auto& stream = streaming.stream;
         auto oldState = kj::mv(streaming);
         state = Borrowed { kj::mv(oldState) };
         return BorrowedStream { *this, stream, limit - completed };
       }
     }
     KJ_UNREACHABLE;
   }
 
   void directEnd() override {
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(redirected, Redirected) {
         // Ugh I guess we need to send a real end() request here.
         redirected.replacement.endRequest(MessageSize {2, 0}).send().detach([](kj::Exception&&){});
       }
       KJ_CASE_ONEOF(e, Ended) {
         // whatever
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         // ... whatever.
       }
       KJ_CASE_ONEOF(streaming, Streaming) {
         auto req = streaming.callback.endedRequest(MessageSize {4, 0});
         req.setByteCount(completed);
         req.send().detach([](kj::Exception&&){});
         streaming.parent.returnStream(completed);
         state = Ended();
       }
     }
   }
 
   // ---------------------------------------------------------------------------
   // implements ByteStream::Server RPC interface
 
   kj::Maybe<kj::Promise<Capability::Client>> shortenPath() override {
     return resolvePromise.addBranch()
         .then([this]() -> Capability::Client {
       return state.get<Redirected>().replacement;
     });
   }
 
   kj::Promise<void> write(WriteContext context) override {
     auto params = context.getParams();
     auto data = params.getBytes();
 
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(redirected, Redirected) {
         auto req = redirected.replacement.writeRequest(params.totalSize());
         req.setBytes(data);
         return req.send();
       }
       KJ_CASE_ONEOF(e, Ended) {
         KJ_FAIL_REQUIRE("already called end()");
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("can't call other methods while stream is borrowed");
       }
       KJ_CASE_ONEOF(streaming, Streaming) {
         if (completed + data.size() < limit) {
           completed += data.size();
           return streaming.stream.write(data.begin(), data.size());
         } else {
           // This write passes the limit.
           uint64_t remainder = limit - completed;
           auto leftover = data.slice(remainder, data.size());
           return streaming.stream.write(data.begin(), remainder)
               .then([this, leftover]() -> kj::Promise<void> {
             completed = limit;
             limitReached();
 
             if (leftover.size() > 0) {
               // Need to forward the leftover bytes to the next stream.
               auto req = state.get<Redirected>().replacement.writeRequest(
                   MessageSize { 4 + leftover.size() / sizeof(capnp::word), 0 });
               req.setBytes(leftover);
               return req.send();
             } else {
               return kj::READY_NOW;
             }
           });
         }
       }
     }
     KJ_UNREACHABLE;
   }
 
   kj::Promise<void> end(EndContext context) override {
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(redirected, Redirected) {
         return context.tailCall(redirected.replacement.endRequest(MessageSize {2,0}));
       }
       KJ_CASE_ONEOF(e, Ended) {
         KJ_FAIL_REQUIRE("already called end()");
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("can't call other methods while stream is borrowed");
       }
       KJ_CASE_ONEOF(streaming, Streaming) {
         auto req = streaming.callback.endedRequest(MessageSize {4, 0});
         req.setByteCount(completed);
         auto result = req.send().ignoreResult();
         streaming.parent.returnStream(completed);
         state = Ended();
         return result;
       }
     }
     KJ_UNREACHABLE;
   }
 
   kj::Promise<void> getSubstream(GetSubstreamContext context) override {
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(redirected, Redirected) {
         auto params = context.getParams();
         auto req = redirected.replacement.getSubstreamRequest(params.totalSize());
         req.setCallback(params.getCallback());
         req.setLimit(params.getLimit());
         return context.tailCall(kj::mv(req));
       }
       KJ_CASE_ONEOF(e, Ended) {
         KJ_FAIL_REQUIRE("already called end()");
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("can't call other methods while stream is borrowed");
       }
       KJ_CASE_ONEOF(streaming, Streaming) {
         auto params = context.getParams();
         auto callback = params.getCallback();
         auto limit = params.getLimit();
         context.releaseParams();
         auto results = context.getResults(MessageSize { 2, 1 });
         results.setSubstream(factory.streamSet.add(kj::heap<SubstreamImpl>(
             factory, *this, thisCap(), streaming.stream, kj::mv(callback), kj::mv(limit))));
         state = Borrowed { kj::mv(streaming) };
         return kj::READY_NOW;
       }
     }
     KJ_UNREACHABLE;
   }
 
 private:
   ByteStreamFactory& factory;
 
   struct Streaming {
     StreamServerBase& parent;
     capnp::ByteStream::Client ownParent;
     kj::AsyncOutputStream& stream;
     capnp::ByteStream::SubstreamCallback::Client callback;
   };
   struct Borrowed {
     Streaming originalState;
   };
   struct Redirected {
     capnp::ByteStream::Client replacement;
   };
   struct Ended {};
 
   kj::OneOf<Streaming, Borrowed, Redirected, Ended> state;
 
   uint64_t limit;
   uint64_t completed = 0;
 
   kj::Own<kj::PromiseFulfiller<void>> resolveFulfiller;
   kj::ForkedPromise<void> resolvePromise;
 
   void limitReached() {
     auto& streaming = state.get<Streaming>();
     auto next = streaming.callback.reachedLimitRequest(capnp::MessageSize {2,0})
         .send().getNext();
 
     // Set the next stream as our replacement.
     streaming.parent.returnStream(limit);
     state = Redirected { kj::mv(next) };
     resolveFulfiller->fulfill();
   }
 };
 
 // =======================================================================================
 
 class ByteStreamFactory::CapnpToKjStreamAdapter final: public StreamServerBase {
   // Implements Cap'n Proto ByteStream as a wrapper around a KJ stream.
 
   class SubstreamCallbackImpl;
 
 public:
   class PathProber;
 
   CapnpToKjStreamAdapter(ByteStreamFactory& factory,
                          kj::Own<kj::AsyncOutputStream> inner)
       : factory(factory),
         state(kj::heap<PathProber>(*this, kj::mv(inner))) {
     state.get<kj::Own<PathProber>>()->startProbing();
   }
 
   CapnpToKjStreamAdapter(ByteStreamFactory& factory,
                          kj::Own<PathProber> pathProber)
       : factory(factory),
         state(kj::mv(pathProber)) {
     state.get<kj::Own<PathProber>>()->setNewParent(*this);
   }
 
   // ---------------------------------------------------------------------------
   // implements StreamServerBase
 
   void returnStream(uint64_t written) override {
     auto stream = kj::mv(state.get<Borrowed>().stream);
     state = kj::mv(stream);
   }
 
   ShortestPath getShortestPath() override {
     // Called by KjToCapnpStreamAdapter when it has determined that its inner ByteStream::Client
     // actually points back to a CapnpToKjStreamAdapter in the same process. Returns the best
     // shortened path to use, or a promise that resolves when the shortest path is known.
 
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(prober, kj::Own<PathProber>) {
         return prober->whenReady();
       }
       KJ_CASE_ONEOF(kjStream, kj::Own<kj::AsyncOutputStream>) {
         auto& streamRef = *kjStream;
         state = Borrowed { kj::mv(kjStream) };
         return StreamServerBase::BorrowedStream { *this, streamRef, kj::maxValue };
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client) {
         return &capnpStream;
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("concurrent streaming calls disallowed") { break; }
         return kj::Promise<void>(kj::READY_NOW);
       }
       KJ_CASE_ONEOF(e, Ended) {
         KJ_FAIL_REQUIRE("already ended") { break; }
         return kj::Promise<void>(kj::READY_NOW);
       }
     }
     KJ_UNREACHABLE;
   }
 
   void directEnd() override {
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(prober, kj::Own<PathProber>) {
         state = Ended();
       }
       KJ_CASE_ONEOF(kjStream, kj::Own<kj::AsyncOutputStream>) {
         state = Ended();
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client) {
         // Ugh I guess we need to send a real end() request here.
         capnpStream.endRequest(MessageSize {2, 0}).send().detach([](kj::Exception&&){});
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         // Fine, ignore.
       }
       KJ_CASE_ONEOF(e, Ended) {
         // Fine, ignore.
       }
     }
   }
 
   // ---------------------------------------------------------------------------
   // PathProber
 
   class PathProber final: public kj::AsyncInputStream {
   public:
     PathProber(CapnpToKjStreamAdapter& parent, kj::Own<kj::AsyncOutputStream> inner,
                kj::PromiseFulfillerPair<void> paf = kj::newPromiseAndFulfiller<void>())
         : parent(parent), inner(kj::mv(inner)),
           readyPromise(paf.promise.fork()),
           readyFulfiller(kj::mv(paf.fulfiller)),
           task(nullptr) {}
 
     void startProbing() {
       task = probeForShorterPath();
     }
 
     void setNewParent(CapnpToKjStreamAdapter& newParent) {
       KJ_ASSERT(parent == nullptr);
       parent = newParent;
       auto paf = kj::newPromiseAndFulfiller<void>();
       readyPromise = paf.promise.fork();
       readyFulfiller = kj::mv(paf.fulfiller);
     }
 
     kj::Promise<void> whenReady() {
       return readyPromise.addBranch();
     }
 
     kj::Promise<uint64_t> pumpToShorterPath(capnp::ByteStream::Client target, uint64_t limit) {
       // If our probe succeeds in finding a KjToCapnpStreamAdapter somewhere down the stack, that
       // will call this method to provide the shortened path.
 
       KJ_IF_MAYBE(currentParent, parent) {
         parent = nullptr;
 
         auto self = kj::mv(currentParent->state.get<kj::Own<PathProber>>());
         currentParent->state = Ended();  // temporary, we'll set this properly below
         KJ_ASSERT(self.get() == this);
 
         // Open a substream on the target stream.
         auto req = target.getSubstreamRequest();
         req.setLimit(limit);
         auto paf = kj::newPromiseAndFulfiller<uint64_t>();
         req.setCallback(kj::heap<SubstreamCallbackImpl>(currentParent->factory,
             kj::mv(self), kj::mv(paf.fulfiller), limit));
 
         // Now we hook up the incoming stream adapter to point directly to this substream, yay.
         currentParent->state = req.send().getSubstream();
 
         // Let the original CapnpToKjStreamAdapter know that it's safe to handle incoming requests.
         readyFulfiller->fulfill();
 
         // It's now up to the SubstreamCallbackImpl to signal when the pump is done.
         return kj::mv(paf.promise);
       } else {
         // We already completed a path-shortening. Probably SubstreamCallbackImpl::ended() was
         // eventually called, meaning the substream was ended without redirecting back to us. So,
         // we're at EOF.
         return uint64_t(0);
       }
     }
 
     kj::Promise<size_t> tryRead(void* buffer, size_t minBytes, size_t maxBytes) override {
       // If this is called, it means the tryPumpFrom() in probeForShorterPath() eventually invoked
       // code that tries to read manually from the source. We don't know what this code is doing
       // exactly, but we do know for sure that the endpoint is not a KjToCapnpStreamAdapter, so
       // we can't optimize. Instead, we pretend that we immediately hit EOF, ending the pump. This
       // works because pumps do not propagate EOF -- the destination can still receive further
       // writes and pumps. Basically our probing pump becomes a no-op, and then we revert to having
       // each write() RPC directly call write() on the inner stream.
       return size_t(0);
     }
 
     kj::Promise<uint64_t> pumpTo(kj::AsyncOutputStream& output, uint64_t amount) override {
       // Call the stream's `tryPumpFrom()` as a way to discover where the data will eventually go,
       // in hopes that we find we can shorten the path.
       KJ_IF_MAYBE(promise, output.tryPumpFrom(*this, amount)) {
         // tryPumpFrom() returned non-null. Either it called `tryRead()` or `pumpTo()` (see
         // below), or it plans to do so in the future.
         return kj::mv(*promise);
       } else {
         // There is no shorter path. As with tryRead(), we pretend we get immediate EOF.
         return uint64_t(0);
       }
     }
 
   private:
     kj::Maybe<CapnpToKjStreamAdapter&> parent;
     kj::Own<kj::AsyncOutputStream> inner;
     kj::ForkedPromise<void> readyPromise;
     kj::Own<kj::PromiseFulfiller<void>> readyFulfiller;
     kj::Promise<void> task;
 
     friend class SubstreamCallbackImpl;
 
     kj::Promise<void> probeForShorterPath() {
       return kj::evalNow([&]() -> kj::Promise<uint64_t> {
         return pumpTo(*inner, kj::maxValue);
       }).then([this](uint64_t actual) {
         KJ_IF_MAYBE(currentParent, parent) {
           KJ_IF_MAYBE(prober, currentParent->state.tryGet<kj::Own<PathProber>>()) {
             // Either we didn't find any shorter path at all during probing and faked an EOF
             // to get out of the probe (see comments in tryRead(), or we DID find a shorter path,
             // completed a pumpTo() using a substream, and that substream redirected back to us,
             // and THEN we couldn't find any further shorter paths for subsequent pumps.
 
             // HACK: If we overwrite the Probing state now, we'll delete ourselves and delete
             //   this task promise, which is an error... let the event loop do it later by
             //   detaching.
             task.attach(kj::mv(*prober)).detach([](kj::Exception&&){});
             parent = nullptr;
 
             // OK, now we can change the parent state and signal it to proceed.
             currentParent->state = kj::mv(inner);
             readyFulfiller->fulfill();
           }
         }
       }).eagerlyEvaluate([this](kj::Exception&& exception) mutable {
         // Something threw, so propagate the exception to break the parent.
         readyFulfiller->reject(kj::mv(exception));
       });
     }
   };
 
 protected:
   // ---------------------------------------------------------------------------
   // implements ByteStream::Server RPC interface
 
   kj::Maybe<kj::Promise<Capability::Client>> shortenPath() override {
     return shortenPathImpl();
   }
   kj::Promise<Capability::Client> shortenPathImpl() {
     // Called by RPC implementation to find out if a shorter path presents itself.
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(prober, kj::Own<PathProber>) {
         return prober->whenReady().then([this]() {
           KJ_ASSERT(!state.is<kj::Own<PathProber>>());
           return shortenPathImpl();
         });
       }
       KJ_CASE_ONEOF(kjStream, kj::Own<kj::AsyncOutputStream>) {
         // No shortening possible. Pretend we never resolve so that calls continue to be routed
         // to us forever.
         return kj::NEVER_DONE;
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client) {
         return Capability::Client(capnpStream);
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("concurrent streaming calls disallowed") { break; }
         return kj::NEVER_DONE;
       }
       KJ_CASE_ONEOF(e, Ended) {
         // No shortening possible. Pretend we never resolve so that calls continue to be routed
         // to us forever.
         return kj::NEVER_DONE;
       }
     }
     KJ_UNREACHABLE;
   }
 
   kj::Promise<void> write(WriteContext context) override {
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(prober, kj::Own<PathProber>) {
         return prober->whenReady().then([this, context]() mutable {
           KJ_ASSERT(!state.is<kj::Own<PathProber>>());
           return write(context);
         });
       }
       KJ_CASE_ONEOF(kjStream, kj::Own<kj::AsyncOutputStream>) {
         auto data = context.getParams().getBytes();
         return kjStream->write(data.begin(), data.size());
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client) {
         auto params = context.getParams();
         auto req = capnpStream.writeRequest(params.totalSize());
         req.setBytes(params.getBytes());
         return req.send();
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("concurrent streaming calls disallowed") { break; }
         return kj::READY_NOW;
       }
       KJ_CASE_ONEOF(e, Ended) {
         KJ_FAIL_REQUIRE("already called end()") { break; }
         return kj::READY_NOW;
       }
     }
     KJ_UNREACHABLE;
   }
 
   kj::Promise<void> end(EndContext context) override {
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(prober, kj::Own<PathProber>) {
         return prober->whenReady().then([this, context]() mutable {
           KJ_ASSERT(!state.is<kj::Own<PathProber>>());
           return end(context);
         });
       }
       KJ_CASE_ONEOF(kjStream, kj::Own<kj::AsyncOutputStream>) {
         // TODO(someday): When KJ adds a proper .end() call, use it here. For now, we must
         //   drop the stream to close it.
         state = Ended();
         return kj::READY_NOW;
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client) {
         auto params = context.getParams();
         auto req = capnpStream.endRequest(params.totalSize());
         return context.tailCall(kj::mv(req));
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("concurrent streaming calls disallowed") { break; }
         return kj::READY_NOW;
       }
       KJ_CASE_ONEOF(e, Ended) {
         KJ_FAIL_REQUIRE("already called end()") { break; }
         return kj::READY_NOW;
       }
     }
     KJ_UNREACHABLE;
   }
 
   kj::Promise<void> getSubstream(GetSubstreamContext context) override {
     KJ_SWITCH_ONEOF(state) {
       KJ_CASE_ONEOF(prober, kj::Own<PathProber>) {
         return prober->whenReady().then([this, context]() mutable {
           KJ_ASSERT(!state.is<kj::Own<PathProber>>());
           return getSubstream(context);
         });
       }
       KJ_CASE_ONEOF(kjStream, kj::Own<kj::AsyncOutputStream>) {
         auto params = context.getParams();
         auto callback = params.getCallback();
         uint64_t limit = params.getLimit();
         context.releaseParams();
 
         auto results = context.initResults(MessageSize {2, 1});
         results.setSubstream(factory.streamSet.add(kj::heap<SubstreamImpl>(
             factory, *this, thisCap(), *kjStream, kj::mv(callback), kj::mv(limit))));
         state = Borrowed { kj::mv(kjStream) };
         return kj::READY_NOW;
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client) {
         auto params = context.getParams();
         auto req = capnpStream.getSubstreamRequest(params.totalSize());
         req.setCallback(params.getCallback());
         req.setLimit(params.getLimit());
         return context.tailCall(kj::mv(req));
       }
       KJ_CASE_ONEOF(b, Borrowed) {
         KJ_FAIL_REQUIRE("concurrent streaming calls disallowed") { break; }
         return kj::READY_NOW;
       }
       KJ_CASE_ONEOF(e, Ended) {
         KJ_FAIL_REQUIRE("already called end()") { break; }
         return kj::READY_NOW;
       }
     }
     KJ_UNREACHABLE;
   }
 
 private:
   ByteStreamFactory& factory;
 
   struct Borrowed { kj::Own<kj::AsyncOutputStream> stream; };
   struct Ended {};
 
   kj::OneOf<kj::Own<PathProber>, kj::Own<kj::AsyncOutputStream>,
             capnp::ByteStream::Client, Borrowed, Ended> state;
 
   class SubstreamCallbackImpl final: public capnp::ByteStream::SubstreamCallback::Server {
   public:
     SubstreamCallbackImpl(ByteStreamFactory& factory,
                           kj::Own<PathProber> pathProber,
                           kj::Own<kj::PromiseFulfiller<uint64_t>> originalPumpfulfiller,
                           uint64_t originalPumpLimit)
         : factory(factory),
           pathProber(kj::mv(pathProber)),
           originalPumpfulfiller(kj::mv(originalPumpfulfiller)),
           originalPumpLimit(originalPumpLimit) {}
 
     ~SubstreamCallbackImpl() noexcept(false) {
       if (!done) {
         originalPumpfulfiller->reject(KJ_EXCEPTION(DISCONNECTED,
             "stream disconnected because SubstreamCallbackImpl was never called back"));
       }
     }
 
     kj::Promise<void> ended(EndedContext context) override {
       KJ_REQUIRE(!done);
       uint64_t actual = context.getParams().getByteCount();
       KJ_REQUIRE(actual <= originalPumpLimit);
 
       done = true;
 
       // EOF before pump completed. Signal a short pump.
       originalPumpfulfiller->fulfill(context.getParams().getByteCount());
 
       // Give the original pump task a chance to finish up.
       return pathProber->task.attach(kj::mv(pathProber));
     }
 
     kj::Promise<void> reachedLimit(ReachedLimitContext context) override {
       KJ_REQUIRE(!done);
       done = true;
 
       // Allow the shortened stream to redirect back to our original underlying stream.
       auto results = context.getResults(capnp::MessageSize { 4, 1 });
       results.setNext(factory.streamSet.add(
           kj::heap<CapnpToKjStreamAdapter>(factory, kj::mv(pathProber))));
 
       // The full pump completed. Note that it's important that we fulfill this after the
       // PathProber has been attached to the new CapnpToKjStreamAdapter, which will have happened
       // in CapnpToKjStreamAdapter's constructor, which calls pathProber->setNewParent().
       originalPumpfulfiller->fulfill(kj::cp(originalPumpLimit));
 
       return kj::READY_NOW;
     }
 
   private:
     ByteStreamFactory& factory;
     kj::Own<PathProber> pathProber;
     kj::Own<kj::PromiseFulfiller<uint64_t>> originalPumpfulfiller;
     uint64_t originalPumpLimit;
     bool done = false;
   };
 };
 
 // =======================================================================================
 
 class ByteStreamFactory::KjToCapnpStreamAdapter final: public kj::AsyncOutputStream {
 public:
   KjToCapnpStreamAdapter(ByteStreamFactory& factory, capnp::ByteStream::Client innerParam)
       : factory(factory),
         inner(kj::mv(innerParam)),
         findShorterPathTask(findShorterPath(inner).fork()) {}
 
   ~KjToCapnpStreamAdapter() noexcept(false) {
     // HACK: KJ streams are implicitly ended on destruction, but the RPC stream needs a call. We
     //   use a detached promise for now, which is probably OK since capabilities are refcounted and
     //   asynchronously destroyed anyway.
     // TODO(cleanup): Fix this when KJ streads add an explicit end() method.
     KJ_IF_MAYBE(o, optimized) {
       o->directEnd();
     } else {
       inner.endRequest(MessageSize {2, 0}).send().detach([](kj::Exception&&){});
     }
   }
 
   kj::Promise<void> write(const void* buffer, size_t size) override {
     KJ_SWITCH_ONEOF(getShortestPath()) {
       KJ_CASE_ONEOF(promise, kj::Promise<void>) {
         return promise.then([this,buffer,size]() {
           return write(buffer, size);
         });
       }
       KJ_CASE_ONEOF(kjStream, StreamServerBase::BorrowedStream) {
         auto limit = kj::min(kjStream.limit, MAX_BYTES_PER_WRITE);
         if (size <= limit) {
           auto promise = kjStream.stream.write(buffer, size);
           return promise.then([kjStream,size]() mutable {
             kjStream.lender.returnStream(size);
           });
         } else {
           auto promise = kjStream.stream.write(buffer, limit);
           return promise.then([this,kjStream,buffer,size,limit]() mutable {
             kjStream.lender.returnStream(limit);
             return write(reinterpret_cast<const byte*>(buffer) + limit,
                          size - limit);
           });
         }
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client*) {
         if (size <= MAX_BYTES_PER_WRITE) {
           auto req = capnpStream->writeRequest(MessageSize { 8 + size / sizeof(word), 0 });
           req.setBytes(kj::arrayPtr(reinterpret_cast<const byte*>(buffer), size));
           return req.send();
         } else {
           auto req = capnpStream->writeRequest(
               MessageSize { 8 + MAX_BYTES_PER_WRITE / sizeof(word), 0 });
           req.setBytes(kj::arrayPtr(reinterpret_cast<const byte*>(buffer), MAX_BYTES_PER_WRITE));
           return req.send().then([this,buffer,size]() mutable {
             return write(reinterpret_cast<const byte*>(buffer) + MAX_BYTES_PER_WRITE,
                          size - MAX_BYTES_PER_WRITE);
           });
         }
       }
     }
     KJ_UNREACHABLE;
   }
 
   kj::Promise<void> write(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) override {
     KJ_SWITCH_ONEOF(getShortestPath()) {
       KJ_CASE_ONEOF(promise, kj::Promise<void>) {
         return promise.then([this,pieces]() {
           return write(pieces);
         });
       }
       KJ_CASE_ONEOF(kjStream, StreamServerBase::BorrowedStream) {
         size_t size = 0;
         for (auto& piece: pieces) { size += piece.size(); }
         auto limit = kj::min(kjStream.limit, MAX_BYTES_PER_WRITE);
         if (size <= limit) {
           auto promise = kjStream.stream.write(pieces);
           return promise.then([kjStream,size]() mutable {
             kjStream.lender.returnStream(size);
           });
         } else {
           // ughhhhhhhhhh, we need to split the pieces.
           return splitAndWrite(pieces, kjStream.limit,
               [kjStream,limit](kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) mutable {
             return kjStream.stream.write(pieces).then([kjStream,limit]() mutable {
               kjStream.lender.returnStream(limit);
             });
           });
         }
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client*) {
         auto writePieces = [capnpStream](kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces) {
           size_t size = 0;
           for (auto& piece: pieces) size += piece.size();
           auto req = capnpStream->writeRequest(MessageSize { 8 + size / sizeof(word), 0 });
           auto out = req.initBytes(size);
           byte* ptr = out.begin();
           for (auto& piece: pieces) {
             memcpy(ptr, piece.begin(), piece.size());
             ptr += piece.size();
           }
           KJ_ASSERT(ptr == out.end());
           return req.send();
         };
 
         size_t size = 0;
         for (auto& piece: pieces) size += piece.size();
         if (size <= MAX_BYTES_PER_WRITE) {
           return writePieces(pieces);
         } else {
           // ughhhhhhhhhh, we need to split the pieces.
           return splitAndWrite(pieces, MAX_BYTES_PER_WRITE, writePieces);
         }
       }
     }
     KJ_UNREACHABLE;
   }
 
   kj::Maybe<kj::Promise<uint64_t>> tryPumpFrom(
       kj::AsyncInputStream& input, uint64_t amount = kj::maxValue) override {
     KJ_IF_MAYBE(rpc, kj::dynamicDowncastIfAvailable<CapnpToKjStreamAdapter::PathProber>(input)) {
       // Oh interesting, it turns we're hosting an incoming ByteStream which is pumping to this
       // outgoing ByteStream. We can let the Cap'n Proto RPC layer know that it can shorten the
       // path from one to the other.
       return rpc->pumpToShorterPath(inner, amount);
     } else {
       return pumpLoop(input, 0, amount);
     }
   }
 
   kj::Promise<void> whenWriteDisconnected() override {
     return findShorterPathTask.addBranch();
   }
 
 private:
   ByteStreamFactory& factory;
   capnp::ByteStream::Client inner;
   kj::Maybe<StreamServerBase&> optimized;
 
   kj::ForkedPromise<void> findShorterPathTask;
   // This serves two purposes:
   // 1. Waits for the capability to resolve (if it is a promise), and then shortens the path if
   //    possible.
   // 2. Implements whenWriteDisconnected().
 
   kj::Promise<void> findShorterPath(capnp::ByteStream::Client& capnpClient) {
     // If the capnp stream turns out to resolve back to this process, shorten the path.
     // Also, implement whenWriteDisconnected() based on this.
     return factory.streamSet.getLocalServer(capnpClient)
         .then([this](kj::Maybe<capnp::ByteStream::Server&> server) -> kj::Promise<void> {
       KJ_IF_MAYBE(s, server) {
         // Yay, we discovered that the ByteStream actually points back to a local KJ stream.
         // We can use this to shorten the path by skipping the RPC machinery.
         return findShorterPath(kj::downcast<StreamServerBase>(*s));
       } else {
         // The capability is fully-resolved. This suggests that the remote implementation is
         // NOT a CapnpToKjStreamAdapter at all, because CapnpToKjStreamAdapter is designed to
         // always look like a promise. It's some other implementation that doesn't present
         // itself as a promise. We have no way to detect when it is disconnected.
         return kj::NEVER_DONE;
       }
     }, [](kj::Exception&& e) -> kj::Promise<void> {
       // getLocalServer() thrown when the capability is a promise cap that rejects. We can
       // use this to implement whenWriteDisconnected().
       //
       // (Note that because this exception handler is passed to the .then(), it does NOT catch
       // eoxceptions thrown by the success handler immediately above it. This handler will ONLY
       // catch exceptions from getLocalServer() itself.)
       return kj::READY_NOW;
     });
   }
 
   kj::Promise<void> findShorterPath(StreamServerBase& capnpServer) {
     // We found a shorter path back to this process. Record it.
     optimized = capnpServer;
 
     KJ_SWITCH_ONEOF(capnpServer.getShortestPath()) {
       KJ_CASE_ONEOF(promise, kj::Promise<void>) {
         return promise.then([this,&capnpServer]() {
           return findShorterPath(capnpServer);
         });
       }
       KJ_CASE_ONEOF(kjStream, StreamServerBase::BorrowedStream) {
         // The ByteStream::Server wraps a regular KJ stream that does not wrap another capnp
         // stream.
         if (kjStream.limit < (uint64_t)kj::maxValue / 2) {
           // But it isn't wrapping that stream forever. Eventually it plans to redirect back to
           // some other stream. So, let's wait for that, and possibly shorten again.
           kjStream.lender.returnStream(0);
           return KJ_ASSERT_NONNULL(capnpServer.shortenPath())
               .then([this, &capnpServer](auto&&) {
             return findShorterPath(capnpServer);
           });
         } else {
           // This KJ stream is (effectively) the permanent endpoint. We can't get any shorter
           // from here. All we want to do now is watch for disconnect.
           auto promise = kjStream.stream.whenWriteDisconnected();
           kjStream.lender.returnStream(0);
           return promise;
         }
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client*) {
         return findShorterPath(*capnpStream);
       }
     }
     KJ_UNREACHABLE;
   }
 
   StreamServerBase::ShortestPath getShortestPath() {
     KJ_IF_MAYBE(o, optimized) {
       return o->getShortestPath();
     } else {
       return &inner;
     }
   }
 
   kj::Promise<uint64_t> pumpLoop(kj::AsyncInputStream& input,
                                  uint64_t completed, uint64_t remaining) {
     if (remaining == 0) return completed;
 
     KJ_SWITCH_ONEOF(getShortestPath()) {
       KJ_CASE_ONEOF(promise, kj::Promise<void>) {
         return promise.then([this,&input,completed,remaining]() {
           return pumpLoop(input,completed,remaining);
         });
       }
       KJ_CASE_ONEOF(kjStream, StreamServerBase::BorrowedStream) {
         // Oh hell yes, this capability actually points back to a stream in our own thread. We can
         // stop sending RPCs and just pump directly.
 
         if (remaining <= kjStream.limit) {
           return input.pumpTo(kjStream.stream, remaining)
               .then([kjStream,completed](uint64_t actual) {
             kjStream.lender.returnStream(actual);
             return actual + completed;
           });
         } else {
           auto promise = input.pumpTo(kjStream.stream, kjStream.limit);
           return promise.then([this,&input,completed,remaining,kjStream]
                               (uint64_t actual) mutable -> kj::Promise<uint64_t> {
             kjStream.lender.returnStream(actual);
             if (actual < kjStream.limit) {
               // EOF reached.
               return completed + actual;
             } else {
               return pumpLoop(input, completed + actual, remaining - actual);
             }
           });
         }
       }
       KJ_CASE_ONEOF(capnpStream, capnp::ByteStream::Client*) {
         // Pumping from some other kind of steram. Optimize the pump by reading from the input
         // directly into outgoing RPC messages.
         size_t size = kj::min(remaining, 8192);
         auto req = capnpStream->writeRequest(MessageSize { 8 + size / sizeof(word) });
 
         auto orphanage = Orphanage::getForMessageContaining(
             capnp::ByteStream::WriteParams::Builder(req));
 
         auto buffer = orphanage.newOrphan<Data>(size);
 
         struct WriteRequestAndBuffer {
           // The order of construction/destruction of lambda captures is unspecified, but we care
           // about ordering between these two things that we want to capture, so... we need a
           // struct.
           StreamingRequest<capnp::ByteStream::WriteParams> request;
           Orphan<Data> buffer;  // points into `request`...
         };
 
         WriteRequestAndBuffer wrab = { kj::mv(req), kj::mv(buffer) };
 
         return input.tryRead(wrab.buffer.get().begin(), 1, size)
             .then([this, &input, completed, remaining, size, wrab = kj::mv(wrab)]
                   (size_t actual) mutable -> kj::Promise<uint64_t> {
           if (actual == 0) {
             return completed;
           } if (actual < size) {
             wrab.buffer.truncate(actual);
           }
 
           wrab.request.adoptBytes(kj::mv(wrab.buffer));
           return wrab.request.send()
               .then([this, &input, completed, remaining, actual]() {
             return pumpLoop(input, completed + actual, remaining - actual);
           });
         });
       }
     }
     KJ_UNREACHABLE;
   }
 
   template <typename WritePieces>
   kj::Promise<void> splitAndWrite(kj::ArrayPtr<const kj::ArrayPtr<const byte>> pieces,
                                   size_t limit, WritePieces&& writeFirstPieces) {
     size_t splitByte = limit;
     size_t splitPiece = 0;
     while (pieces[splitPiece].size() <= splitByte) {
       splitByte -= pieces[splitPiece].size();
       ++splitPiece;
     }
 
     if (splitByte == 0) {
       // Oh thank god, the split is between two pieces.
       auto rest = pieces.slice(splitPiece, pieces.size());
       return writeFirstPieces(pieces.slice(0, splitPiece))
           .then([this,rest]() mutable {
         return write(rest);
       });
     } else {
       // FUUUUUUUU---- we need to split one of the pieces in two.
       auto left = kj::heapArray<kj::ArrayPtr<const byte>>(splitPiece + 1);
       auto right = kj::heapArray<kj::ArrayPtr<const byte>>(pieces.size() - splitPiece);
       for (auto i: kj::zeroTo(splitPiece)) {
         left[i] = pieces[i];
       }
       for (auto i: kj::zeroTo(right.size())) {
         right[i] = pieces[splitPiece + i];
       }
       left.back() = pieces[splitPiece].slice(0, splitByte);
       right.front() = pieces[splitPiece].slice(splitByte, pieces[splitPiece].size());
 
       return writeFirstPieces(left).attach(kj::mv(left))
           .then([this,right=kj::mv(right)]() mutable {
         return write(right).attach(kj::mv(right));
       });
     }
   }
 };
 
 // =======================================================================================
 
 capnp::ByteStream::Client ByteStreamFactory::kjToCapnp(kj::Own<kj::AsyncOutputStream> kjStream) {
   return streamSet.add(kj::heap<CapnpToKjStreamAdapter>(*this, kj::mv(kjStream)));
 }
 
 kj::Own<kj::AsyncOutputStream> ByteStreamFactory::capnpToKj(capnp::ByteStream::Client capnpStream) {
   return kj::heap<KjToCapnpStreamAdapter>(*this, kj::mv(capnpStream));
 }
 
 }  // namespace capnp