Prefer returning read data on cancel

Summary:
It's possible that AsyncTransport read completed successfully in one loop, then the read coro gets a cancellation request.  In this case, prefer to return successfully with the read data.

Read cancellation need not be fatal, eg: an application may just want to interrupt a read and come back to it later.  In this case, the successfully read data could be lost.

In particular, if an EOF was lost, the application might start another read, AsyncSocket::setReadCB can assert in invalidState(ReadCallback*) because the DestructorGuard count is 0.

Reviewed By: yairgott

Differential Revision: D30964504

fbshipit-source-id: ea963e6572f2bb552899d3b1e6274bbd0bcdc265

folly/io/coro/Transport.cpp

@@ -305,13 +305,15 @@ Task<size_t> Transport::read(
   transport_->setReadCB(&cb);
   auto waitRet =
       co_await co_awaitTry(cb.wait(co_await co_current_cancellation_token));
-
-  if (waitRet.hasException()) {
-    co_yield co_error(std::move(waitRet.exception()));
-  }
   if (cb.error()) {
     co_yield co_error(std::move(cb.error()));
   }
+  if (waitRet.hasException() &&
+      (!waitRet.tryGetExceptionObject<OperationCancelled>() ||
+       (!cb.eof && cb.length == 0))) {
+    // Got a non-cancel exception, or cancel with nothing read
+    co_yield co_error(std::move(waitRet.exception()));
+  }
   transport_->setReadCB(nullptr);
   deferredReadEOF_ = (cb.eof && cb.length > 0);
   co_return cb.length;
@@ -338,12 +340,15 @@ Task<size_t> Transport::read(
   transport_->setReadCB(&cb);
   auto waitRet =
       co_await co_awaitTry(cb.wait(co_await co_current_cancellation_token));
-  if (waitRet.hasException()) {
-    co_yield co_error(std::move(waitRet.exception()));
-  }
   if (cb.error()) {
     co_yield co_error(std::move(cb.error()));
   }
+  if (waitRet.hasException() &&
+      (!waitRet.tryGetExceptionObject<OperationCancelled>() ||
+       (!cb.eof && cb.length == 0))) {
+    // Got a non-cancel exception, or cancel with nothing read
+    co_yield co_error(std::move(waitRet.exception()));
+  }
   transport_->setReadCB(nullptr);
   deferredReadEOF_ = (cb.eof && cb.length > 0);
   co_return cb.length;

folly/io/coro/test/TransportTest.cpp

@@ -19,6 +19,7 @@
 #include <folly/experimental/coro/BlockingWait.h>
 #include <folly/experimental/coro/Collect.h>
 #include <folly/io/async/test/AsyncSocketTest.h>
+#include <folly/io/async/test/MockAsyncTransport.h>
 #include <folly/io/async/test/ScopedBoundPort.h>
 #include <folly/io/coro/ServerSocket.h>
 #include <folly/io/coro/Transport.h>
@@ -327,4 +328,41 @@ TEST_F(TransportTest, AsyncClientAndServer) {
   });
 }
 
+class MockTransportTest : public TransportTest {
+ public:
+  folly::coro::Task<Transport> connect() {
+    mockTransport = new testing::NiceMock<test::MockAsyncTransport>();
+    folly::AsyncTransport::UniquePtr transport(mockTransport);
+    co_return Transport(&evb, std::move(transport));
+  }
+  test::MockAsyncTransport* mockTransport;
+};
+
+TEST_F(MockTransportTest, readSuccessCanceled) {
+  run([&]() -> Task<> {
+    auto cs = co_await connect();
+    constexpr auto kBufSize = 65536;
+    std::array<uint8_t, kBufSize> rcvBuf;
+    EXPECT_CALL(*mockTransport, setReadCB(testing::_))
+        .WillOnce(testing::Invoke(
+            [](AsyncReader::ReadCallback* rcb) { rcb->readEOF(); }));
+    EXPECT_CALL(*mockTransport, setReadCB(nullptr)).Times(2);
+    folly::CancellationSource cancellationSource;
+    auto readFut =
+        co_withCancellation(
+            cancellationSource.getToken(),
+            cs.read(
+                MutableByteRange(rcvBuf.data(), rcvBuf.data() + rcvBuf.size()),
+                100ms))
+            .scheduleOn(&evb)
+            .start();
+    // Let the read coro start and get the EOF
+    co_await co_reschedule_on_current_executor;
+    // cancel
+    cancellationSource.requestCancellation();
+    // read succeeds with nRead == 0
+    auto nRead = co_await std::move(readFut);
+    EXPECT_EQ(nRead, 0);
+  });
+}
 #endif // FOLLY_HAS_COROUTINES