folly/synchronization/detail/ThreadCachedLists.h

   1 /*
   2  * Copyright 2017-present Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #pragma once
  18
  19 #include <atomic>
  20
  21 #include <folly/Function.h>
  22 #include <folly/ThreadLocal.h>
  23 #include <glog/logging.h>
  24
  25 namespace folly {
  26
  27 namespace detail {
  28
  29 // This is a thread-local cached, multi-producer single-consumer
  30 // queue, similar to a concurrent version of std::list.
  31 //
  32 class ThreadCachedListsBase {
  33  public:
  34   struct Node {
  35     folly::Function<void()> cb_;
  36     Node* next_{nullptr};
  37   };
  38 };
  39
  40 template <typename Tag>
  41 class ThreadCachedLists : public ThreadCachedListsBase {
  42  public:
  43   struct AtomicListHead {
  44     std::atomic<Node*> tail_{nullptr};
  45     std::atomic<Node*> head_{nullptr};
  46   };
  47
  48   // Non-concurrent list, similar to std::list.
  49   struct ListHead {
  50     Node* head_{nullptr};
  51     Node* tail_{nullptr};
  52
  53     // Run func on each list node.
  54     template <typename Func>
  55     void forEach(Func func) {
  56       auto node = tail_;
  57       while (node != nullptr) {
  58         auto next = node->next_;
  59         func(node);
  60         node = next;
  61       }
  62     }
  63
  64     // Splice other in to this list.
  65     // Afterwards, other is a valid empty listhead.
  66     void splice(ListHead& other);
  67
  68     void splice(AtomicListHead& other);
  69   };
  70
  71   // Push a node on a thread-local list.  Returns true if local list
  72   // was pushed global.
  73   void push(Node* node);
  74
  75   // Collect all thread local lists to a single local list.
  76   // This function is threadsafe with concurrent push()es,
  77   // but only a single thread may call collect() at a time.
  78   void collect(ListHead& list);
  79
  80  private:
  81   // Push list to the global list.
  82   void pushGlobal(ListHead& list);
  83
  84   ListHead ghead_;
  85
  86   struct TLHead : public AtomicListHead {
  87     ThreadCachedLists* parent_;
  88
  89    public:
  90     TLHead(ThreadCachedLists* parent) : parent_(parent) {}
  91
  92     ~TLHead() {
  93       parent_->ghead_.splice(*this);
  94     }
  95   };
  96
  97   folly::ThreadLocalPtr<TLHead, Tag> lhead_;
  98 };
  99
 100 // push() and splice() are optimistic w.r.t setting the list head: The
 101 // first pusher cas's the list head, which functions as a lock until
 102 // tail != null.  The first pusher then sets tail_ = head_.
 103 //
 104 // splice() does the opposite: steals the tail_ via exchange, then
 105 // unlocks the list again by setting head_ to null.
 106 template <typename Tag>
 107 void ThreadCachedLists<Tag>::push(Node* node) {
 108   DCHECK(node->next_ == nullptr);
 109   static thread_local TLHead* cache_{nullptr};
 110
 111   if (!cache_) {
 112     auto l = lhead_.get();
 113     if (!l) {
 114       lhead_.reset(new TLHead(this));
 115       l = lhead_.get();
 116       DCHECK(l);
 117     }
 118     cache_ = l;
 119   }
 120
 121   while (true) {
 122     auto head = cache_->head_.load(std::memory_order_relaxed);
 123     if (!head) {
 124       node->next_ = nullptr;
 125       if (cache_->head_.compare_exchange_weak(head, node)) {
 126         cache_->tail_.store(node);
 127         break;
 128       }
 129     } else {
 130       auto tail = cache_->tail_.load(std::memory_order_relaxed);
 131       if (tail) {
 132         node->next_ = tail;
 133         if (cache_->tail_.compare_exchange_weak(node->next_, node)) {
 134           break;
 135         }
 136       }
 137     }
 138   }
 139 }
 140
 141 template <typename Tag>
 142 void ThreadCachedLists<Tag>::collect(ListHead& list) {
 143   auto acc = lhead_.accessAllThreads();
 144
 145   for (auto& thr : acc) {
 146     list.splice(thr);
 147   }
 148
 149   list.splice(ghead_);
 150 }
 151
 152 template <typename Tag>
 153 void ThreadCachedLists<Tag>::ListHead::splice(ListHead& other) {
 154   if (other.head_ != nullptr) {
 155     DCHECK(other.tail_ != nullptr);
 156   } else {
 157     DCHECK(other.tail_ == nullptr);
 158     return;
 159   }
 160
 161   if (head_) {
 162     DCHECK(tail_ != nullptr);
 163     DCHECK(head_->next_ == nullptr);
 164     head_->next_ = other.tail_;
 165     head_ = other.head_;
 166   } else {
 167     DCHECK(head_ == nullptr);
 168     head_ = other.head_;
 169     tail_ = other.tail_;
 170   }
 171
 172   other.head_ = nullptr;
 173   other.tail_ = nullptr;
 174 }
 175
 176 template <typename Tag>
 177 void ThreadCachedLists<Tag>::ListHead::splice(AtomicListHead& list) {
 178   ListHead local;
 179
 180   auto tail = list.tail_.load();
 181   if (tail) {
 182     local.tail_ = list.tail_.exchange(nullptr);
 183     local.head_ = list.head_.exchange(nullptr);
 184     splice(local);
 185   }
 186 }
 187
 188 } // namespace detail
 189 } // namespace folly