高性能C++之无锁队列

高性能C++之无锁队列

lock_free_queue.h

#pragma once

#include <atomic>
#include <vector>
#include <cstddef>

template<typename T, size_t N = 1024>
class LockFreeQueue
{
public:
	LockFreeQueue() {}
	LockFreeQueue(size_t s) : data_(s) {
		read_index_ = 0;
		write_index_ = 0;
	}
	~LockFreeQueue() {}

public:

	bool enqueue(T value) {
		size_t write_index = 0;
		Element* e = nullptr;
		do {
			write_index = write_index_.load(std::memory_order_relaxed);
			if (write_index >= (read_index_.load(std::memory_order_relaxed) + data_.size())) { // 等于时刚好满
				return false;
			}
			size_t index = write_index % N;
			e = &data_[index];
			if (e->full_.load(std::memory_order_relaxed)) {
				return false;
			}
		} while (!write_index_.compare_exchange_weak(write_index, write_index + 1), std::memory_order_release, std::memory_order_relaxed);
		
		e->data_ = std::move(value);
		e->full_.store(true, std::memory_order_release);
		return true;
	}

	bool dequeue(T value) {
		size_t read_index = 0;
		Element* e = nullptr;
		do {
			read_index = read_index_.load(std::memory_order_relaxed);
			if (read_index >= write_index_.load(std::memory_order_relaxed)) {
				return false;
			}
			size_t index = read_index % N;
			e = &data_[index];
			if (!e->full_.load(std::memory_order_relaxed)) {
				return false;
			}
		} while (!read_index_.compare_exchange_weak(read_index, read_index + 1), std::memory_order_release, std::memory_order_relaxed);
		
		value = std::move(e->data_);
		e->full_.store(false, std::memory_order_release);
		return true;
	}

public:
	struct Element {
		std::atomic<bool> full_;
		T data_;
	};

private:
	std::vector<Element> data_;
	std::atomic<size_t> read_index_;
	std::atomic<size_t> write_index_;
};

main.cpp

#include "lock_free_queue.h"
#include <iostream>
#include <thread>
#include <atomic>

int main()
{
	LockFreeQueue<uint16_t> q;
	std::atomic<uint16_t> seq;
	seq = 0;
	static const int PRODUCE_N = 4;
	static const int CONSUME_N = 4;
	static const int MULTIPLIER = 4;
	std::atomic<int> finished_producer;
	finished_producer = 0;
	auto producer = [&q, &seq, &finished_producer]() {
		for (int i = 0; i < 65535 * MULTIPLIER; i++) {
			uint16_t s = seq++;
			while (!q.enqueue(s));
		}
		finished_producer++;
	};

	std::atomic<uint32_t> counter[65535];
	for (int i = 0; i < 65535; i++) {
		counter[i] = 0;
	}
	auto consumer = [&q, &counter, &finished_producer]() {
		uint16_t s = 0;
		while (finished_producer < PRODUCE_N) {
			if (!q.dequeue(s)) {
				counter[s]++;
			}
		}
		while (q.dequeue(s)) {
			counter[s]++;
		}
	};

	std::unique_ptr<std::thread> produce_threads[PRODUCE_N];
	std::unique_ptr<std::thread> consumer_threads[CONSUME_N];
	
	for (int i = 0; i < CONSUME_N; i++) {
		consumer_threads[i].reset(new std::thread(consumer));
	}
	for (int i = 0; i < PRODUCE_N; i++) {
		produce_threads[i].reset(new std::thread(producer));
	}
	for (int i = 0; i < PRODUCE_N; i++) {
		produce_threads[i]->join();
	}
	for (int i = 0; i < CONSUME_N; i++) {
		consumer_threads[i]->join();
	}

	bool has_race = false;
	for (int i = 0; i < 65535; i++) {
		if (counter[i] != MULTIPLIER * PRODUCE_N) {
			std::cout << "found face condition\t" << i << "\t" << counter[i] << std::endl;
			has_race = true;
			break;
		}
	}
	if (has_race) {
		std::cout << "found face condition" << std::endl;
	}
	else {
		std::cout << "no face condition" << std::endl;
	}

	return 0;
}