folly/futures/Barrier.cpp

   1 /*
   2  * Copyright 2017 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 #include <folly/futures/Barrier.h>
  18
  19 namespace folly { namespace futures {
  20
  21 Barrier::Barrier(uint32_t n)
  22   : size_(n),
  23     controlBlock_(allocateControlBlock()) { }
  24
  25 Barrier::~Barrier() {
  26   auto block = controlBlock_.load(std::memory_order_relaxed);
  27   auto prev = block->valueAndReaderCount.load(std::memory_order_relaxed);
  28   DCHECK_EQ(prev >> kReaderShift, 0u);
  29   auto val = prev & kValueMask;
  30   auto p = promises(block);
  31
  32   for (uint32_t i = 0; i < val; ++i) {
  33     p[i].setException(
  34         folly::make_exception_wrapper<std::runtime_error>("Barrier destroyed"));
  35   }
  36
  37   freeControlBlock(controlBlock_);
  38 }
  39
  40 auto Barrier::allocateControlBlock() -> ControlBlock* {
  41   auto block = static_cast<ControlBlock*>(malloc(controlBlockSize(size_)));
  42   if (!block) {
  43     throw std::bad_alloc();
  44   }
  45   block->valueAndReaderCount = 0;
  46
  47   auto p = promises(block);
  48   uint32_t i = 0;
  49   try {
  50     for (i = 0; i < size_; ++i) {
  51       new (p + i) BoolPromise();
  52     }
  53   } catch (...) {
  54     for (; i != 0; --i) {
  55       p[i - 1].~BoolPromise();
  56     }
  57     throw;
  58   }
  59
  60   return block;
  61 }
  62
  63 void Barrier::freeControlBlock(ControlBlock* block) {
  64   auto p = promises(block);
  65   for (uint32_t i = size_; i != 0; --i) {
  66     p[i - 1].~BoolPromise();
  67   }
  68   free(block);
  69 }
  70
  71 folly::Future<bool> Barrier::wait() {
  72   // Load the current control block first. As we know there is at least
  73   // one thread in the current epoch (us), this means that the value is
  74   // < size_, so controlBlock_ can't change until we bump the value below.
  75   auto block = controlBlock_.load(std::memory_order_acquire);
  76   auto p = promises(block);
  77
  78   // Bump the value and record ourselves as reader.
  79   // This ensures that block stays allocated, as the reader count is > 0.
  80   auto prev = block->valueAndReaderCount.fetch_add(kReader + 1,
  81                                                    std::memory_order_acquire);
  82
  83   auto prevValue = static_cast<uint32_t>(prev & kValueMask);
  84   DCHECK_LT(prevValue, size_);
  85   auto future = p[prevValue].getFuture();
  86
  87   if (prevValue + 1 == size_) {
  88     // Need to reset the barrier before fulfilling any futures. This is
  89     // when the epoch is flipped to the next.
  90     controlBlock_.store(allocateControlBlock(), std::memory_order_release);
  91
  92     p[0].setValue(true);
  93     for (uint32_t i = 1; i < size_; ++i) {
  94       p[i].setValue(false);
  95     }
  96   }
  97
  98   // Free the control block if we're the last reader at max value.
  99   prev = block->valueAndReaderCount.fetch_sub(kReader,
 100                                               std::memory_order_acq_rel);
 101   if (prev == (kReader | uint64_t(size_))) {
 102     freeControlBlock(block);
 103   }
 104
 105   return future;
 106 }
 107
 108 }}  // namespaces