nonblocking_notifier.hpp

#pragma once
 
#include <iostream>
#include <vector>
#include <cstdlib>
#include <cstdio>
#include <atomic>
#include <memory>
#include <deque>
#include <mutex>
#include <condition_variable>
#include <thread>
#include <algorithm>
#include <numeric>
#include <cassert>
#include "../utility/os.hpp"

 
namespace tf {

class NonblockingNotifier {
 
  friend class Executor;
  
  struct Waiter {
    alignas (TF_CACHELINE_SIZE) std::atomic<Waiter*> next;
    uint64_t epoch;
 
    enum : unsigned {
      kNotSignaled = 0,
      kWaiting,
      kSignaled,
    };
    std::atomic<unsigned> state {0};
 
    //mutable std::mutex mu;
    //std::condition_variable cv;
    //unsigned state;
  };
 
  public:
  
  // The state variable consists of the following three parts:
  // - low STACK_BITS is a stack of waiters committed wait.
  // - next PREWAITER_BITS is count of waiters in prewait state.
  // - next EPOCH_BITS is modification counter.
  // [ 32-bit epoch | 16-bit pre-waiter count | 16-bit pre-waiter stack]

  static const uint64_t STACK_BITS = 16;

  static const uint64_t STACK_MASK = (1ull << STACK_BITS) - 1;

  static const uint64_t PREWAITER_BITS = 16;

  static const uint64_t PREWAITER_SHIFT = 16;

  static const uint64_t PREWAITER_MASK = ((1ull << PREWAITER_BITS) - 1) << PREWAITER_SHIFT;

  static const uint64_t PREWAITER_INC = 1ull << PREWAITER_BITS;

  static const uint64_t EPOCH_BITS = 32;

  static const uint64_t EPOCH_SHIFT = 32;

  static const uint64_t EPOCH_MASK = ((1ull << EPOCH_BITS) - 1) << EPOCH_SHIFT;

  static const uint64_t EPOCH_INC = 1ull << EPOCH_SHIFT;

  explicit NonblockingNotifier(size_t N) : _state(STACK_MASK), _waiters(N) {
    if(_waiters.size() >= ((1 << PREWAITER_BITS) - 1)) {
      TF_THROW("nonblocking waiter supports only up to ", (1<<PREWAITER_BITS)-1, " waiters");
    }
    //assert(_waiters.size() < (1 << PREWAITER_BITS) - 1);
    // Initialize epoch to something close to overflow to test overflow.
    //_state = STACK_MASK | (EPOCH_MASK - EPOCH_INC * _waiters.size() * 2);
  }
  
  ~NonblockingNotifier() {
    // Ensure there are no waiters.
    assert((_state.load() & (STACK_MASK | PREWAITER_MASK)) == STACK_MASK);
  }
  
  size_t num_waiters() const {
    size_t n = 0;
    for(auto& w : _waiters) {
      n += (w.state.load(std::memory_order_relaxed) == Waiter::kWaiting);
      //std::scoped_lock lock(w.mu);
      //n += (w.state == Waiter::kWaiting);
    }
    return n;
  }
  
  size_t capacity() const {
    return 1 << STACK_BITS;
  }

  void prepare_wait(size_t wid) {
    _waiters[wid].epoch = _state.fetch_add(PREWAITER_INC, std::memory_order_relaxed);
    std::atomic_thread_fence(std::memory_order_seq_cst);
  }

  void commit_wait(size_t wid) {
 
    auto w = &_waiters[wid];
 
    w->state.store(Waiter::kNotSignaled, std::memory_order_relaxed);
    
    /*
    Epoch and ticket semantics.
    
      `sepoch` = _state & EPOCH_MASK
      `wepoch` = w->epoch & EPOCH_MASK
      `ticket` = w->epoch & PREWAITER_MASK
    
    Each waiter entering the pre-waiting stage is assigned a monotonically
    increasing ticket that determines the processing order (e.g.,
    cancel_wait, commit_wait, notify). Ticket 0 is processed first, followed
    by ticket 1, and so on.
    
    The global epoch `sepoch` is incremented whenever a request is fulfilled.
    Therefore, the difference `sepoch - wepoch` indicates which ticket is
    currently ready to be handled:
    
      - `sepoch - wepoch == ticket` : this waiter's turn
      - `sepoch - wepoch >  ticket` : this waiter's ticket has expired
      - `sepoch - wepoch <  ticket` : this waiter's ticket has not yet reached
    
    Unsigned wraparound does not affect correctness. All epoch arithmetic is
    performed using unsigned integers, which obey modulo-2^N arithmetic.
    Converting the unsigned difference to a signed value yields the correct
    result as long as the true difference lies within the signed range.
    
    In general:
      - Unsigned range: [0, 2^N − 1]
      - Signed range  : [−2^(N−1), 2^(N−1) − 1]
    
    When overflow occurs, unsigned subtraction computes:
    
      (sepoch − wepoch) mod 2^N
    
    If the true value of `sepoch − wepoch` is within the signed range
    [−2^(N−1), 2^(N−1) − 1], reinterpreting this result as a signed integer
    produces the correct mathematical difference.
    
    Example (3-bit arithmetic):
    
      a  b | true a−b | unsigned (bin / dec) | signed (dec)
      ----------------------------------------------------
      1  0 |   1      | 001 / 1              | +1
      1  1 |   0      | 000 / 0              |  0
      1  2 |  -1      | 111 / 7              | -1
      1  3 |  -2      | 110 / 6              | -2
      1  4 |  -3      | 101 / 5              | -3
      1  5 |  -4      | 100 / 4              | -4
      1  6 |  -5      | 011 / 3              | +3 (wrap around)
      1  7 |  -6      | 010 / 2              | +2 (wrap around)
    
    Signed interpretation is correct only when the true difference lies
    within [−4, +3].
    
    In this implementation, `sepoch − wepoch` is guaranteed not to exceed
    2^16 in magnitude, which is far smaller than 2^(EPOCH_BITS − 1).
    Consequently, the expression:
    
      int64_t((state & EPOCH_MASK) - epoch)
    
    remains correct even if `sepoch` and `wepoch` individually overflow.
    */
    uint64_t epoch =
        (w->epoch & EPOCH_MASK) +
        (((w->epoch & PREWAITER_MASK) >> PREWAITER_SHIFT) << EPOCH_SHIFT);
    uint64_t state = _state.load(std::memory_order_seq_cst);
    for (;;) {
      if (int64_t((state & EPOCH_MASK) - epoch) < 0) {
        // The preceding waiter has not decided on its fate. Wait until it
        // calls either cancel_wait or commit_wait, or is notified.
        std::this_thread::yield();
        state = _state.load(std::memory_order_seq_cst);
        continue;
      }
      // We've already been notified.
      if (int64_t((state & EPOCH_MASK) - epoch) > 0) {
        return;
      }
      // Remove this thread from prewait counter and add it to the waiter stack.
      assert((state & PREWAITER_MASK) != 0);
      uint64_t newstate = state - PREWAITER_INC + EPOCH_INC;
      newstate = (newstate & ~STACK_MASK) | wid;
 
      // stack is empty -> this waiter is at the top of the stack, pointing to nothing
      if ((state & STACK_MASK) == STACK_MASK) {
        w->next.store(nullptr, std::memory_order_relaxed);
      }
      // stack is non-empty -> this waiter is at the top of the stack, pointing to the origin top
      else {
        w->next.store(&_waiters[state & STACK_MASK], std::memory_order_relaxed);
      }
      if (_state.compare_exchange_weak(state, newstate, std::memory_order_release)) {
        break;
      }
    }
    _park(w);
  }

  void cancel_wait(size_t wid) {
    uint64_t epoch =
      (_waiters[wid].epoch & EPOCH_MASK) +
      (((_waiters[wid].epoch & PREWAITER_MASK) >> PREWAITER_SHIFT) << EPOCH_SHIFT);
    uint64_t state = _state.load(std::memory_order_relaxed);
    for (;;) {
      if (int64_t((state & EPOCH_MASK) - epoch) < 0) {
        // The preceding waiter has not decided on its fate. Wait until it
        // calls either cancel_wait or commit_wait, or is notified.
        std::this_thread::yield();
        state = _state.load(std::memory_order_relaxed);
        continue;
      }
      // We've already been notified.
      if (int64_t((state & EPOCH_MASK) - epoch) > 0) {
        return;
      }
      // Remove this thread from prewait counter.
      assert((state & PREWAITER_MASK) != 0);
      if (_state.compare_exchange_weak(state, state - PREWAITER_INC + EPOCH_INC,
                                       std::memory_order_relaxed)) {
        return;
      }
    }
  }
  
  void notify_one() {
    std::atomic_thread_fence(std::memory_order_seq_cst);
    uint64_t state = _state.load(std::memory_order_acquire);
    for (;;) {
      // Easy case: no waiters.
      if ((state & STACK_MASK) == STACK_MASK && (state & PREWAITER_MASK) == 0) {
        return;
      }
      uint64_t num_prewaiters = (state & PREWAITER_MASK) >> PREWAITER_SHIFT;
      uint64_t newstate;
      if (num_prewaiters) {
        // There is a thread in pre-wait state, unblock it.
        newstate = state + EPOCH_INC - PREWAITER_INC;
      } 
      else {
        // Pop a waiter from list and unpark it.
        Waiter* w = &_waiters[state & STACK_MASK];
        Waiter* wnext = w->next.load(std::memory_order_relaxed);
        uint64_t next = STACK_MASK;
        //if (wnext != nullptr) next = wnext - &_waiters[0];
        if (wnext != nullptr) {
          next = static_cast<uint64_t>(wnext - &_waiters[0]);
        }
        // Note: we don't add EPOCH_INC here. ABA problem on the lock-free stack
        // can't happen because a waiter is re-pushed onto the stack only after
        // it was in the pre-wait state which inevitably leads to epoch increment.
        newstate = (state & EPOCH_MASK) + next;
      }
      if (_state.compare_exchange_weak(state, newstate, std::memory_order_acquire)) {
        if(num_prewaiters) {
          return; // unblocked pre-wait thread
        }
        // if there is no pre-waiters, the stack must have something 
        Waiter* w = &_waiters[state & STACK_MASK];
        w->next.store(nullptr, std::memory_order_relaxed);
        _unpark(w);
        return;
      }
    }
  }
  
  void notify_all() {
    std::atomic_thread_fence(std::memory_order_seq_cst);
    uint64_t state = _state.load(std::memory_order_acquire);
    for (;;) {
 
      // Easy case: no waiters.
      if ((state & STACK_MASK) == STACK_MASK && (state & PREWAITER_MASK) == 0) {
        return;
      }
      uint64_t num_prewaiters = (state & PREWAITER_MASK) >> PREWAITER_SHIFT;
 
      // Reset prewait counter and empty wait list.
      uint64_t newstate = (state & EPOCH_MASK) + (EPOCH_INC * num_prewaiters) + STACK_MASK;
 
      if (_state.compare_exchange_weak(state, newstate, std::memory_order_acquire)) {
        if ((state & STACK_MASK) == STACK_MASK) {
          return;
        }
        Waiter* w = &_waiters[state & STACK_MASK];
        _unpark(w);
        return;
      }
    }
  }
  
  void notify_n(size_t N) {
 
    // trivial case
    if(N == 0) {
      return;
    }
    
    // if the target N is bigger than the waiter size, notify all waiters
    if(N >= _waiters.size()) {
      notify_all();
      return;
    }
 
    std::atomic_thread_fence(std::memory_order_seq_cst);
    uint64_t state = _state.load(std::memory_order_acquire);
    do {
      // Easy case: no waiters.
      if ((state & STACK_MASK) == STACK_MASK && (state & PREWAITER_MASK) == 0) {
        return;
      }
      uint64_t num_prewaiters = (state & PREWAITER_MASK) >> PREWAITER_SHIFT;
      uint64_t newstate;
      size_t   newN;
      
      // unblock waiters from pre-waiting list first.
      if(num_prewaiters) {
        size_t to_unblock = (N < num_prewaiters) ? N : num_prewaiters;
        newstate = state + (EPOCH_INC * to_unblock) - (PREWAITER_INC * to_unblock);
        newN = N - to_unblock;
      }
      // pop one waiter from the stack
      else {
        Waiter* w = &_waiters[state & STACK_MASK];
        Waiter* wnext = w->next.load(std::memory_order_relaxed);
        uint64_t next = STACK_MASK;
        //if (wnext != nullptr) next = wnext - &_waiters[0];
        if (wnext != nullptr) {
          next = static_cast<uint64_t>(wnext - &_waiters[0]);
        }
        // Note: we don't add EPOCH_INC here. ABA problem on the lock-free stack
        // can't happen because a waiter is re-pushed onto the stack only after
        // it was in the pre-wait state which inevitably leads to epoch increment.
        newstate = (state & EPOCH_MASK) + next;
        newN = N - 1;
      }
 
      if (_state.compare_exchange_weak(state, newstate, std::memory_order_acquire)) {
        N = newN;
        if(num_prewaiters == 0) {
          Waiter* w = &_waiters[state & STACK_MASK];
          w->next.store(nullptr, std::memory_order_relaxed);
          _unpark(w);
        }
      }
    } while(N > 0);
 
    //if(n >= _waiters.size()) {
    //  notify_all();
    //}
    //else {
    //  for(size_t k=0; k<n; ++k) {
    //    notify_one();
    //  }
    //}
  }
  
  size_t size() const {
    return _waiters.size();
  }
 
 private:
 
  std::atomic<uint64_t> _state;
  std::vector<Waiter> _waiters;
  
  // only this waiter can park itself, with the following two possible paths:
  // 1. kNotSignaled (this) -> in-stack -> kWaiting (this) -> wait
  // 2. kNotSignaled (this) -> in-stack -> kSignaled -> unwait
  void _park(Waiter* w) {
    unsigned target = Waiter::kNotSignaled;
    if(w->state.compare_exchange_strong(target, Waiter::kWaiting, std::memory_order_relaxed
                                                                , std::memory_order_relaxed)) {
      w->state.wait(Waiter::kWaiting, std::memory_order_relaxed);
    }
    //std::unique_lock<std::mutex> lock(w->mu);
    //while (w->state != Waiter::kSignaled) {
    //  w->state = Waiter::kWaiting;
    //  w->cv.wait(lock);
    //}
  }
  
  // others can unpark
  void _unpark(Waiter* waiters) {
    Waiter* next = nullptr;
    for (Waiter* w = waiters; w; w = next) {
      next = w->next.load(std::memory_order_relaxed);
      // We only notify if the other is waiting - this is why we use tri-state
      // variable instead of binary-state variable (i.e., atomic_flag)
      // Performance is about 0.1% faster
      if(w->state.exchange(Waiter::kSignaled, std::memory_order_relaxed) == Waiter::kWaiting) {
        w->state.notify_one();
      }
      //unsigned state;
      //{
      //  std::unique_lock<std::mutex> lock(w->mu);
      //  state = w->state;
      //  w->state = Waiter::kSignaled;
      //}
      //if (state == Waiter::kWaiting) w->cv.notify_one();
    }
  }
  
  // notify wakes one or all waiting threads.
  // Must be called after changing the associated wait predicate.
  //void _notify(bool all) {
  //  std::atomic_thread_fence(std::memory_order_seq_cst);
  //  uint64_t state = _state.load(std::memory_order_acquire);
  //  for (;;) {
  //    // Easy case: no waiters.
  //    if ((state & STACK_MASK) == STACK_MASK && (state & PREWAITER_MASK) == 0) {
  //      return;
  //    }
  //    uint64_t num_prewaiters = (state & PREWAITER_MASK) >> PREWAITER_SHIFT;
  //    uint64_t newstate;
  //    if (all) {
  //      // Reset prewait counter and empty wait list.
  //      newstate = (state & EPOCH_MASK) + (EPOCH_INC * num_prewaiters) + STACK_MASK;
  //    } else if (num_prewaiters) {
  //      // There is a thread in pre-wait state, unblock it.
  //      newstate = state + EPOCH_INC - PREWAITER_INC;
  //    } else {
  //      // Pop a waiter from list and unpark it.
  //      Waiter* w = &_waiters[state & STACK_MASK];
  //      Waiter* wnext = w->next.load(std::memory_order_relaxed);
  //      uint64_t next = STACK_MASK;
  //      //if (wnext != nullptr) next = wnext - &_waiters[0];
  //      if (wnext != nullptr) {
  //        next = static_cast<uint64_t>(wnext - &_waiters[0]);
  //      }
  //      // Note: we don't add EPOCH_INC here. ABA problem on the lock-free stack
  //      // can't happen because a waiter is re-pushed onto the stack only after
  //      // it was in the pre-wait state which inevitably leads to epoch increment.
  //      newstate = (state & EPOCH_MASK) + next;
  //    }
  //    if (_state.compare_exchange_weak(state, newstate, std::memory_order_acquire)) {
  //      if(!all && num_prewaiters) return; // unblocked pre-wait thread
  //      if ((state & STACK_MASK) == STACK_MASK) return;
  //      Waiter* w = &_waiters[state & STACK_MASK];
  //      if(!all) {
  //        w->next.store(nullptr, std::memory_order_relaxed);
  //      }
  //      _unpark(w);
  //      return;
  //    }
  //  }
  //}
};
 
 
}  // namespace tf ------------------------------------------------------------