foundationdb/flow/ThreadSafeQueue.h

// ThreadSafeQueue<T> is a multi-producer, single-consumer queue.

// It is almost but not quite lock-free (the exception is that if a thread is
// stopped in a very narrow window in push(), pop() will "block" in the sense
// of returning Optional<T>() to the consumer until the former thread makes progress)

// It has an extra facility for event-loop integration: calling canSleep() before blocking
// the consumer thread and waking the consumer thread whenever push() returns true will permit
// the consumer thread to block when the queue is empty with a minimum of overhead.  A caller
// can ignore this facility by simply not calling canSleep and ignoring the return value of push().

// Based in part on the queue at
// http://www.1024cores.net/home/lock-free-algorithms/queues/intrusive-mpsc-node-based-queue which is covered by this
// BSD license:
/* Copyright (c) 2010-2011 Dmitry Vyukov. All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
following conditions are met:

   1. Redistributions of source code must retain the above copyright notice, this list of
      conditions and the following disclaimer.

   2. Redistributions in binary form must reproduce the above copyright notice, this list
      of conditions and the following disclaimer in the documentation and/or other materials
      provided with the distribution.

THIS SOFTWARE IS PROVIDED BY DMITRY VYUKOV "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL DMITRY
VYUKOV OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
DAMAGE.

The views and conclusions contained in the software and documentation are those of the authors and should not be
interpreted as representing official policies, either expressed or implied, of Dmitry Vyukov.*/

#include <atomic>

#if VALGRIND
#include <drd.h>
#endif

template <class T>
class ThreadSafeQueue : NonCopyable {
	struct BaseNode {
		std::atomic<BaseNode*> next;
		BaseNode() : next(nullptr) {}
	};

	struct Node : BaseNode, FastAllocated<Node> {
		T data;
		Node(T const& data) : data(data) {}
	};
	std::atomic<BaseNode*> head;
	BaseNode* tail;
	BaseNode stub, sleeping;
	bool sleepy;

	BaseNode* popNode() {
		BaseNode* tail = this->tail;
		BaseNode* next = tail->next.load();
#if VALGRIND
		ANNOTATE_HAPPENS_BEFORE(&tail->next);
#endif
		if (tail == &this->stub) {
			if (!next) {
				return nullptr;
			}
			this->tail = next;
			tail = next;
			next = next->next.load();
#if VALGRIND
			ANNOTATE_HAPPENS_BEFORE(&next->next);
#endif
		}
		if (next) {
			this->tail = next;
			return tail;
		}
		BaseNode* head = this->head.load();
#if VALGRIND
		ANNOTATE_HAPPENS_BEFORE(&this->head);
#endif
		if (tail != head) {
			return nullptr;
		}
#if VALGRIND
		ANNOTATE_HAPPENS_AFTER(&this->stub.next);
#endif
		this->stub.next.store(nullptr);
		pushNode(&this->stub);
		next = tail->next.load();
#if VALGRIND
		ANNOTATE_HAPPENS_BEFORE(&tail->next);
#endif
		if (next) {
			this->tail = next;
			return tail;
		}
		return nullptr;
	}

	// Pushes n at the end of the queue and returns the node immediately before it
	BaseNode* pushNode(BaseNode* n) {
#if VALGRIND
		ANNOTATE_HAPPENS_AFTER(&head);
#endif
		BaseNode* prev = head.exchange(n);
#if VALGRIND
		ANNOTATE_HAPPENS_BEFORE(&head);
		ANNOTATE_HAPPENS_AFTER(&prev->next);
#endif
		prev->next.store(n);
		return prev;
	}

public:
	ThreadSafeQueue() {
#if VALGRIND
		ANNOTATE_HAPPENS_AFTER(&this->head);
#endif
		this->head.store(&this->stub);
		this->tail = &this->stub;
		this->sleepy = false;
	}
	~ThreadSafeQueue() {
		while (pop().present())
			;
	}

	// If push() returns true, the consumer may be sleeping and should be woken
	bool push(T const& data) {
		Node* n = new Node(data);
		n->data = data;
		return pushNode(n) == &sleeping;
	}

	///////////// The below functions may only be called by a single, consumer thread //////////////////

	// If canSleep returns true, then the queue is empty and the next push() will return true
	bool canSleep() {
		if (sleepy) {
			return false; // We already have sleeping in the queue from a previous call to canSleep.  Pop it and then
			              // maybe you can sleep!
		}
		if (this->tail != &stub || this->tail->next.load()) {
#if VALGRIND
			ANNOTATE_HAPPENS_BEFORE(&this->tail->next);
#endif
			return false; // There is definitely something in the queue.  This is a rejection test not needed for
			              // correctness, but avoids calls to pushNode
		}
#if VALGRIND
		ANNOTATE_HAPPENS_BEFORE(&this->tail->next);
#endif
		// sleeping is definitely not in the queue, so we can safely try...
		sleeping.next.store(nullptr);
		bool ok = pushNode(&sleeping) == &stub;
		sleepy = true; // sleeping is in the queue, regardless of whether it's the first thing; we need to pop it before
		               // sleeping again
		return ok;
	}

	Optional<T> pop() {
		BaseNode* b = popNode();
		if (b == &sleeping) {
			sleepy = false;
			b = popNode();
		}
		if (b == &sleeping)
			ASSERT(false);
		if (b == &stub)
			ASSERT(false);
		if (!b)
			return Optional<T>();

		Node* n = (Node*)b;
		T data = std::move(n->data);
		delete n;
		return Optional<T>(std::move(data));
	}
};