folly/fibers/AddTasks-inl.h

   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 #include <memory>
  17 #include <vector>
  18
  19 namespace folly {
  20 namespace fibers {
  21
  22 template <typename T>
  23 TaskIterator<T>::TaskIterator(TaskIterator&& other) noexcept
  24     : context_(std::move(other.context_)), id_(other.id_), fm_(other.fm_) {}
  25
  26 template <typename T>
  27 inline bool TaskIterator<T>::hasCompleted() const {
  28   return context_->tasksConsumed < context_->results.size();
  29 }
  30
  31 template <typename T>
  32 inline bool TaskIterator<T>::hasPending() const {
  33   return !context_.unique();
  34 }
  35
  36 template <typename T>
  37 inline bool TaskIterator<T>::hasNext() const {
  38   return hasPending() || hasCompleted();
  39 }
  40
  41 template <typename T>
  42 folly::Try<T> TaskIterator<T>::awaitNextResult() {
  43   assert(hasCompleted() || hasPending());
  44   reserve(1);
  45
  46   size_t i = context_->tasksConsumed++;
  47   id_ = context_->results[i].first;
  48   return std::move(context_->results[i].second);
  49 }
  50
  51 template <typename T>
  52 inline T TaskIterator<T>::awaitNext() {
  53   return std::move(awaitNextResult().value());
  54 }
  55
  56 template <>
  57 inline void TaskIterator<void>::awaitNext() {
  58   awaitNextResult().value();
  59 }
  60
  61 template <typename T>
  62 inline void TaskIterator<T>::reserve(size_t n) {
  63   size_t tasksReady = context_->results.size() - context_->tasksConsumed;
  64
  65   // we don't need to do anything if there are already n or more tasks complete
  66   // or if we have no tasks left to execute.
  67   if (!hasPending() || tasksReady >= n) {
  68     return;
  69   }
  70
  71   n -= tasksReady;
  72   size_t tasksLeft = context_->totalTasks - context_->results.size();
  73   n = std::min(n, tasksLeft);
  74
  75   await([this, n](Promise<void> promise) {
  76     context_->tasksToFulfillPromise = n;
  77     context_->promise.assign(std::move(promise));
  78   });
  79 }
  80
  81 template <typename T>
  82 inline size_t TaskIterator<T>::getTaskID() const {
  83   assert(id_ != static_cast<size_t>(-1));
  84   return id_;
  85 }
  86
  87 template <typename T>
  88 template <typename F>
  89 void TaskIterator<T>::addTask(F&& func) {
  90   static_assert(
  91       std::is_convertible<typename std::result_of<F()>::type, T>::value,
  92       "TaskIterator<T>: T must be convertible from func()'s return type");
  93
  94   auto taskId = context_->totalTasks++;
  95
  96   fm_.addTask(
  97       [ taskId, context = context_, func = std::forward<F>(func) ]() mutable {
  98         context->results.emplace_back(
  99             taskId, folly::makeTryWith(std::move(func)));
 100
 101         // Check for awaiting iterator.
 102         if (context->promise.hasValue()) {
 103           if (--context->tasksToFulfillPromise == 0) {
 104             context->promise->setValue();
 105             context->promise.clear();
 106           }
 107         }
 108       });
 109 }
 110
 111 template <class InputIterator>
 112 TaskIterator<typename std::result_of<
 113     typename std::iterator_traits<InputIterator>::value_type()>::type>
 114 addTasks(InputIterator first, InputIterator last) {
 115   typedef typename std::result_of<
 116       typename std::iterator_traits<InputIterator>::value_type()>::type
 117       ResultType;
 118   typedef TaskIterator<ResultType> IteratorType;
 119
 120   IteratorType iterator;
 121
 122   for (; first != last; ++first) {
 123     iterator.addTask(std::move(*first));
 124   }
 125
 126   iterator.context_->results.reserve(iterator.context_->totalTasks);
 127
 128   return std::move(iterator);
 129 }
 130 }
 131 }