concurrency-in-action/ch7/queue_mpmc.cpp

#include <atomic>
#include <cstdio>
#include <iostream>
#include <memory>
#include <thread>

// NOTE: if atomic operations on std::shared_ptr<> are lock-free,
// then compare/exchange would be enough to implement the lock-free
// mpmc queue

template <typename T>
class queue {
private:
  struct node;
  struct counted_node_ptr {
    int external_count;
    node* ptr;
  };

  std::atomic<counted_node_ptr> head;
  std::atomic<counted_node_ptr> tail;

  struct node_counter {
    unsigned internal_count : 30;
    unsigned external_counters : 2;
    // you need only 2 bits because there are at most two such
    // counters (next and tail)
  };

  struct node {
    std::atomic<T*> data;
    std::atomic<node_counter> count;
    std::atomic<counted_node_ptr> next;

    node()
    {
      node_counter new_count;
      new_count.internal_count    = 0;
      new_count.external_counters = 2;
      // because every node starts out referenced from tail and
      // from the next pointer of the previous node once you've
      // actually aadded it to the queue
      count.store(new_count);

      counted_node_ptr new_next;
      new_next.ptr            = nullptr;
      new_next.external_count = 0;
      next.store(new_next);
    }

    void release_ref()
    {
      node_counter old_counter = count.load(std::memory_order_relaxed);
      node_counter new_counter;
      do {
        new_counter = old_counter;
        --new_counter.internal_count;
      } while (!count.compare_exchange_strong(old_counter,
                                              new_counter,
                                              std::memory_order_acquire,
                                              std::memory_order_relaxed));
      if (!new_counter.internal_count && !new_counter.external_counters) {
        delete this;
      }
    }
  };

  void set_new_tail(counted_node_ptr& old_tail,
                    counted_node_ptr const& new_tail)
  {
    node* const current_tail_ptr = old_tail.ptr;
    while (!tail.compare_exchange_weak(old_tail, new_tail) &&
           old_tail.ptr == current_tail_ptr)
      ;
    if (old_tail.ptr == current_tail_ptr)
      free_external_counter(old_tail);
    else
      current_tail_ptr->release_ref();
  }

  static void increase_external_count(std::atomic<counted_node_ptr>& counter,
                                      counted_node_ptr& old_counter)
  {
    counted_node_ptr new_counter;
    do {
      new_counter = old_counter;
      ++new_counter.external_count;
    } while (!counter.compare_exchange_strong(old_counter,
                                              new_counter,
                                              std::memory_order_acquire,
                                              std::memory_order_relaxed));
    old_counter.external_count = new_counter.external_count;
  }

  static void free_external_counter(counted_node_ptr& old_node_ptr)
  {
    node* const ptr          = old_node_ptr.ptr;
    int const count_increase = old_node_ptr.external_count - 2;
    node_counter old_counter = ptr->count.load(std::memory_order_relaxed);
    node_counter new_counter;
    do {
      new_counter = old_counter;
      --new_counter.external_counters;
      new_counter.internal_count += count_increase;
    } while (!ptr->count.compare_exchange_strong(old_counter,
                                                 new_counter,
                                                 std::memory_order_acquire,
                                                 std::memory_order_relaxed));
    if (!new_counter.internal_count && !new_counter.external_counters) {
      delete ptr;
    }
  }

public:
  queue() : head(), tail() {}

  queue(const queue& other) = delete;
  queue& operator=(const queue& other) = delete;

  ~queue()
  {
    /*
    while (node* const old_head = head.load()) {
        head.store(old_head->next);
        delete old_head;
    }
    */
  }

  void push(T new_value)
  {
    std::unique_ptr<T> new_data(new T(new_value));
    counted_node_ptr new_next;
    new_next.ptr              = new node;
    new_next.external_count   = 1;
    counted_node_ptr old_tail = tail.load();
    for (;;) {
      increase_external_count(tail, old_tail);
      T* old_data = nullptr;
      if (old_tail.ptr->data.compare_exchange_strong(old_data,
                                                     new_data.get())) {
        counted_node_ptr old_next = {0};
        if (!old_tail.ptr->next.compare_exchange_strong(old_next, new_next)) {
          delete new_next.ptr;
          new_next = old_next;
        }
        set_new_tail(old_tail, new_next);
        new_data.release();
        break;
      }
      else {
        counted_node_ptr old_next = {0};
        if (old_tail.ptr->next.compare_exchange_strong(old_next, new_next)) {
          old_next     = new_next;
          new_next.ptr = new node;
        }
        set_new_tail(old_tail, old_next);
      }
    }
  }

  std::unique_ptr<T> pop()
  {
    counted_node_ptr old_head = head.load(std::memory_order_relaxed);

    for (;;) {
      increase_external_count(head, old_head);

      node* const ptr = old_head.ptr;
      if (ptr == tail.load().ptr) {
        ptr->release_ref();
        return std::unique_ptr<T>();
      }

      counted_node_ptr next = ptr->next.load();

      if (head.compare_exchange_strong(old_head, next)) {
        T* const res = ptr->data.exchange(nullptr);
        free_external_counter(old_head);
        return std::unique_ptr<T>(res);
      }
      ptr->release_ref();
    }
  }
};

void
push(queue<int>* q)
{
  for (int i = 0; i < 10; ++i) {
    printf("pushing %d\n", i);
    q->push(i);
  }
}

void
pop(queue<int>* q)
{
  int i = 0;
  while (i < 10) {
    std::shared_ptr<int> p = q->pop();
    if (p) {
      printf("poping %d\n", *p);
      ++i;
    }
  }
}

int
main()
{
  queue<int> q;
  std::thread t1(push, &q);
  std::thread t2(pop, &q);
  t1.join();
  t2.join();
  return 0;
}