Commit 49b1f90b authored by Mike Kolupaev's avatar Mike Kolupaev Committed by Facebook GitHub Bot

Fix coro::UnboundedQueue::try_dequeue

Summary: try_dequeue() used to consume from the queue but not from the semaphore, causing the two to diverge and breaking subsequent dequeue() calls.

Reviewed By: yfeldblum

Differential Revision: D23160324

fbshipit-source-id: f709733d984f51abc53595d38dbef3daffbfa6d9
parent 481a93de
...@@ -49,11 +49,11 @@ class UnboundedQueue { ...@@ -49,11 +49,11 @@ class UnboundedQueue {
} }
folly::Optional<T> try_dequeue() { folly::Optional<T> try_dequeue() {
return queue_.try_dequeue(); return sem_.try_wait() ? queue_.try_dequeue() : folly::none;
} }
bool try_dequeue(T& out) { bool try_dequeue(T& out) {
return queue_.try_dequeue(out); return sem_.try_wait() ? queue_.try_dequeue(out) : false;
} }
bool empty() { bool empty() {
......
...@@ -190,4 +190,22 @@ TEST(UnboundedQueueTest, CancelledDequeueCompletesNormallyIfAnItemIsAvailable) { ...@@ -190,4 +190,22 @@ TEST(UnboundedQueueTest, CancelledDequeueCompletesNormallyIfAnItemIsAvailable) {
}()); }());
} }
TEST(UnboundedQueueTest, TryDequeue) {
folly::coro::UnboundedQueue<int> queue;
queue.enqueue(42);
EXPECT_EQ(42, queue.try_dequeue());
folly::ManualExecutor ex;
auto fut = queue.dequeue().scheduleOn(&ex).start();
ex.drain();
EXPECT_FALSE(fut.isReady());
queue.enqueue(13);
ex.drain();
EXPECT_TRUE(fut.isReady());
EXPECT_EQ(std::move(fut).get(), 13);
}
#endif #endif
...@@ -119,6 +119,20 @@ bool Semaphore::try_wait(Waiter& waiter) { ...@@ -119,6 +119,20 @@ bool Semaphore::try_wait(Waiter& waiter) {
return true; return true;
} }
bool Semaphore::try_wait() {
auto oldVal = tokens_.load(std::memory_order_acquire);
do {
if (oldVal == 0) {
return false;
}
} while (!tokens_.compare_exchange_weak(
oldVal,
oldVal - 1,
std::memory_order_release,
std::memory_order_acquire));
return true;
}
#if FOLLY_HAS_COROUTINES #if FOLLY_HAS_COROUTINES
coro::Task<void> Semaphore::co_wait() { coro::Task<void> Semaphore::co_wait() {
......
...@@ -72,6 +72,12 @@ class Semaphore { ...@@ -72,6 +72,12 @@ class Semaphore {
*/ */
bool try_wait(Waiter& waiter); bool try_wait(Waiter& waiter);
/**
* If the semaphore has capacity, removes a token and returns true. Otherwise
* returns false and leaves the semaphore unchanged.
*/
bool try_wait();
#if FOLLY_HAS_COROUTINES #if FOLLY_HAS_COROUTINES
/* /*
......
...@@ -1732,7 +1732,7 @@ TEST(FiberManager, semaphore) { ...@@ -1732,7 +1732,7 @@ TEST(FiberManager, semaphore) {
for (size_t i = 0; i < kTasks; ++i) { for (size_t i = 0; i < kTasks; ++i) {
manager.addTask([&, completionCounter]() { manager.addTask([&, completionCounter]() {
for (size_t j = 0; j < kIterations; ++j) { for (size_t j = 0; j < kIterations; ++j) {
switch (j % 3) { switch (j % 4) {
case 0: case 0:
sem.wait(); sem.wait();
break; break;
...@@ -1751,6 +1751,11 @@ TEST(FiberManager, semaphore) { ...@@ -1751,6 +1751,11 @@ TEST(FiberManager, semaphore) {
} }
break; break;
} }
case 3:
if (!sem.try_wait()) {
sem.wait();
}
break;
} }
++counter; ++counter;
sem.signal(); sem.signal();
......
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