Commit a6de480b authored by Andrew Smith's avatar Andrew Smith Committed by Facebook GitHub Bot

Change AsyncPipe to accept an optional onClosed callback

Summary: This diff allows the user to specify an optional onClosed callback for AsyncPipe, analogous to the callback accepted by thrift's ServerStreamPublisher.

Reviewed By: andriigrynenko

Differential Revision: D23024437

fbshipit-source-id: 546cc7e0ea1099cadfa085f5dd2fb9085b88f695
parent 5e4f665c
...@@ -37,38 +37,73 @@ namespace coro { ...@@ -37,38 +37,73 @@ namespace coro {
// write() returns false if the read end has been destroyed // write() returns false if the read end has been destroyed
// The generator is completed when the write end is destroyed or on close() // The generator is completed when the write end is destroyed or on close()
// close() can also be passed an exception, which is thrown when read // close() can also be passed an exception, which is thrown when read
//
// An optional onClosed callback can be passed to create(). This callback will
// be called either when the generator is destroyed by the consumer, or when
// the pipe is closed by the publisher (whichever comes first). The onClosed
// callback may destroy the AsyncPipe object inline, and must not call close()
// on the AsyncPipe object inline. If an onClosed callback is specified and the
// publisher would like to destroy the pipe outside of the callback, it must
// first close the pipe.
template <typename T> template <typename T>
class AsyncPipe { class AsyncPipe {
public: public:
~AsyncPipe() { ~AsyncPipe() {
CHECK(!onClosed_ || onClosed_->wasInvokeRequested())
<< "If an onClosed callback is specified and the generator still "
<< "exists, the publisher must explicitly close the pipe prior to "
<< "destruction.";
std::move(*this).close(); std::move(*this).close();
} }
AsyncPipe(AsyncPipe&& pipe) noexcept { AsyncPipe(AsyncPipe&& pipe) noexcept {
queue_ = std::move(pipe.queue_); queue_ = std::move(pipe.queue_);
onClosed_ = std::move(pipe.onClosed_);
} }
AsyncPipe& operator=(AsyncPipe&& pipe) { AsyncPipe& operator=(AsyncPipe&& pipe) {
if (this != &pipe) { if (this != &pipe) {
CHECK(!onClosed_ || onClosed_->wasInvokeRequested())
<< "If an onClosed callback is specified and the generator still "
<< "exists, the publisher must explicitly close the pipe prior to "
<< "destruction.";
std::move(*this).close(); std::move(*this).close();
queue_ = std::move(pipe.queue_); queue_ = std::move(pipe.queue_);
onClosed_ = std::move(pipe.onClosed_);
} }
return *this; return *this;
} }
static std::pair<folly::coro::AsyncGenerator<T&&>, AsyncPipe<T>> create() { static std::pair<folly::coro::AsyncGenerator<T&&>, AsyncPipe<T>> create(
folly::Function<void()> onClosed = nullptr) {
auto queue = std::make_shared<Queue>(); auto queue = std::make_shared<Queue>();
auto cancellationSource = std::shared_ptr<folly::CancellationSource>();
auto onClosedCallback = std::unique_ptr<OnClosedCallback>();
if (onClosed != nullptr) {
cancellationSource = std::make_shared<folly::CancellationSource>();
onClosedCallback = std::make_unique<OnClosedCallback>(
cancellationSource, std::move(onClosed));
}
auto guard = folly::makeGuard([cancellationSource] {
if (cancellationSource != nullptr) {
cancellationSource->requestCancellation();
}
});
return { return {
folly::coro::co_invoke([queue]() -> folly::coro::AsyncGenerator<T&&> { folly::coro::co_invoke(
while (true) { [ queue,
auto val = co_await queue->dequeue(); guard = std::move(guard) ]() -> folly::coro::AsyncGenerator<T&&> {
if (val.hasValue() || val.hasException()) { while (true) {
co_yield std::move(*val); auto val = co_await queue->dequeue();
} else { if (val.hasValue() || val.hasException()) {
co_return; co_yield std::move(*val);
} } else {
} co_return;
}), }
AsyncPipe(queue)}; }
}),
AsyncPipe(queue, std::move(onClosedCallback))};
} }
template <typename U = T> template <typename U = T>
...@@ -85,18 +120,56 @@ class AsyncPipe { ...@@ -85,18 +120,56 @@ class AsyncPipe {
queue->enqueue(folly::Try<T>(std::move(ew))); queue->enqueue(folly::Try<T>(std::move(ew)));
queue_.reset(); queue_.reset();
} }
if (onClosed_ != nullptr) {
onClosed_->requestInvoke();
onClosed_.reset();
}
} }
void close() && { void close() && {
if (auto queue = queue_.lock()) { if (auto queue = queue_.lock()) {
queue->enqueue(folly::Try<T>()); queue->enqueue(folly::Try<T>());
queue_.reset(); queue_.reset();
} }
if (onClosed_ != nullptr) {
onClosed_->requestInvoke();
onClosed_.reset();
}
} }
private: private:
using Queue = folly::coro::UnboundedQueue<folly::Try<T>, true, true>; using Queue = folly::coro::UnboundedQueue<folly::Try<T>, true, true>;
explicit AsyncPipe(std::weak_ptr<Queue> queue) : queue_(std::move(queue)) {}
class OnClosedCallback {
public:
OnClosedCallback(
std::shared_ptr<folly::CancellationSource> cancellationSource,
folly::Function<void()> onClosedFunc)
: cancellationSource_(std::move(cancellationSource)),
cancellationCallback_(
cancellationSource_->getToken(),
std::move(onClosedFunc)) {}
void requestInvoke() {
cancellationSource_->requestCancellation();
}
bool wasInvokeRequested() {
return cancellationSource_->isCancellationRequested();
}
private:
std::shared_ptr<folly::CancellationSource> cancellationSource_;
folly::CancellationCallback cancellationCallback_;
};
explicit AsyncPipe(
std::weak_ptr<Queue> queue,
std::unique_ptr<OnClosedCallback> onClosed)
: queue_(std::move(queue)), onClosed_(std::move(onClosed)) {}
std::weak_ptr<Queue> queue_; std::weak_ptr<Queue> queue_;
std::unique_ptr<OnClosedCallback> onClosed_;
}; };
} // namespace coro } // namespace coro
......
...@@ -209,4 +209,122 @@ TEST(AsyncPipeTest, DestroyWhileBlocking) { ...@@ -209,4 +209,122 @@ TEST(AsyncPipeTest, DestroyWhileBlocking) {
EXPECT_TRUE(fut.isReady()); EXPECT_TRUE(fut.isReady());
EXPECT_FALSE(std::move(fut).get()); EXPECT_FALSE(std::move(fut).get());
} }
TEST(AsyncPipeTest, OnClosedCallbackCalledWhenGeneratorDestroyed) {
auto onCloseCallbackBaton = folly::Baton<>();
auto pipe = folly::coro::AsyncPipe<int>::create(
[&]() { onCloseCallbackBaton.post(); } /* onClosed */);
auto ex = folly::ManualExecutor();
auto cancellationSource = folly::CancellationSource();
auto fut = folly::coro::co_withCancellation(
cancellationSource.getToken(),
folly::coro::co_invoke(
[gen = std::move(pipe.first)]() mutable
-> folly::coro::Task<
folly::coro::AsyncGenerator<int&&>::NextResult> {
co_return co_await gen.next();
}))
.scheduleOn(&ex)
.start();
ex.drain();
EXPECT_FALSE(fut.isReady());
EXPECT_FALSE(onCloseCallbackBaton.ready());
cancellationSource.requestCancellation();
ex.drain();
EXPECT_TRUE(onCloseCallbackBaton.ready());
}
TEST(AsyncPipeTest, OnClosedCallbackCalledWhenPublisherClosesPipe) {
auto onCloseCallbackBaton = folly::Baton<>();
auto pipe = folly::coro::AsyncPipe<int>::create(
[&]() { onCloseCallbackBaton.post(); } /* onClosed */);
EXPECT_FALSE(onCloseCallbackBaton.ready());
std::move(pipe.second).close();
EXPECT_TRUE(onCloseCallbackBaton.ready());
}
TEST(AsyncPipeTest, OnClosedCallbackCalledWhenPublisherClosesPipeWithError) {
auto onCloseCallbackExecuted = folly::Promise<folly::Unit>();
auto pipe = folly::coro::AsyncPipe<int>::create(
[&]() { onCloseCallbackExecuted.setValue(); } /* onClosed */);
EXPECT_FALSE(onCloseCallbackExecuted.isFulfilled());
std::move(pipe.second).close(std::runtime_error("error"));
EXPECT_TRUE(onCloseCallbackExecuted.isFulfilled());
}
TEST(
AsyncPipeTest,
OnClosedCallbackJoinedWhenPublisherClosesPipeWhileGeneratorDestructing) {
auto onCloseCallbackStartedBaton = folly::Baton<>();
auto onCloseCallbackCompletionBaton = folly::Baton<>();
auto pipe = folly::coro::AsyncPipe<int>::create([&]() {
onCloseCallbackStartedBaton.post();
onCloseCallbackCompletionBaton.wait();
} /* onClosed */);
auto ex = folly::ManualExecutor();
auto cancellationSource = folly::CancellationSource();
auto fut = folly::coro::co_withCancellation(
cancellationSource.getToken(),
folly::coro::co_invoke(
[gen = std::move(pipe.first)]() mutable
-> folly::coro::Task<
folly::coro::AsyncGenerator<int&&>::NextResult> {
co_return co_await gen.next();
}))
.scheduleOn(&ex)
.start();
ex.drain();
EXPECT_FALSE(fut.isReady());
EXPECT_FALSE(onCloseCallbackStartedBaton.ready());
auto cancelThread = std::thread([&]() {
cancellationSource.requestCancellation();
ex.drain();
});
onCloseCallbackStartedBaton.wait();
auto pipeClosedBaton = folly::Baton<>();
auto closePipeThread = std::thread([&]() {
std::move(pipe.second).close();
pipeClosedBaton.post();
});
/* sleep override */
std::this_thread::sleep_for(std::chrono::milliseconds(100));
EXPECT_FALSE(pipeClosedBaton.ready());
onCloseCallbackCompletionBaton.post();
pipeClosedBaton.wait();
cancelThread.join();
closePipeThread.join();
}
TEST(AsyncPipeTest, PublisherMustCloseIfCallbackSetAndGeneratorAlive) {
EXPECT_DEATH(
([&]() {
auto pipe = folly::coro::AsyncPipe<int>::create([]() {} /* onClosed */);
})(),
"If an onClosed callback is specified and the generator still exists, "
"the publisher must explicitly close the pipe prior to destruction.");
EXPECT_DEATH(
([&]() {
auto pipe1 =
folly::coro::AsyncPipe<int>::create([]() {} /* onClosed */);
auto pipe2 = folly::coro::AsyncPipe<int>::create();
pipe1.second = std::move(pipe2.second);
})(),
"If an onClosed callback is specified and the generator still exists, "
"the publisher must explicitly close the pipe prior to destruction.");
}
#endif #endif
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment