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 {
// write() returns false if the read end has been destroyed
// 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
//
// 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>
class AsyncPipe {
public:
~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();
}
AsyncPipe(AsyncPipe&& pipe) noexcept {
queue_ = std::move(pipe.queue_);
onClosed_ = std::move(pipe.onClosed_);
}
AsyncPipe& operator=(AsyncPipe&& 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();
queue_ = std::move(pipe.queue_);
onClosed_ = std::move(pipe.onClosed_);
}
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 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 {
folly::coro::co_invoke([queue]() -> folly::coro::AsyncGenerator<T&&> {
while (true) {
auto val = co_await queue->dequeue();
if (val.hasValue() || val.hasException()) {
co_yield std::move(*val);
} else {
co_return;
}
}
}),
AsyncPipe(queue)};
folly::coro::co_invoke(
[ queue,
guard = std::move(guard) ]() -> folly::coro::AsyncGenerator<T&&> {
while (true) {
auto val = co_await queue->dequeue();
if (val.hasValue() || val.hasException()) {
co_yield std::move(*val);
} else {
co_return;
}
}
}),
AsyncPipe(queue, std::move(onClosedCallback))};
}
template <typename U = T>
......@@ -85,18 +120,56 @@ class AsyncPipe {
queue->enqueue(folly::Try<T>(std::move(ew)));
queue_.reset();
}
if (onClosed_ != nullptr) {
onClosed_->requestInvoke();
onClosed_.reset();
}
}
void close() && {
if (auto queue = queue_.lock()) {
queue->enqueue(folly::Try<T>());
queue_.reset();
}
if (onClosed_ != nullptr) {
onClosed_->requestInvoke();
onClosed_.reset();
}
}
private:
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::unique_ptr<OnClosedCallback> onClosed_;
};
} // namespace coro
......
......@@ -209,4 +209,122 @@ TEST(AsyncPipeTest, DestroyWhileBlocking) {
EXPECT_TRUE(fut.isReady());
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
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