133 lines
3.6 KiB
C++
133 lines
3.6 KiB
C++
/*
|
|
* Copyright (c) Facebook, Inc. and its affiliates.
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
namespace folly {
|
|
namespace fibers {
|
|
|
|
template <typename T>
|
|
TaskIterator<T>::TaskIterator(TaskIterator&& other) noexcept
|
|
: context_(std::move(other.context_)), id_(other.id_), fm_(other.fm_) {}
|
|
|
|
template <typename T>
|
|
inline bool TaskIterator<T>::hasCompleted() const {
|
|
return context_->tasksConsumed < context_->results.size();
|
|
}
|
|
|
|
template <typename T>
|
|
inline bool TaskIterator<T>::hasPending() const {
|
|
return !context_.unique();
|
|
}
|
|
|
|
template <typename T>
|
|
inline bool TaskIterator<T>::hasNext() const {
|
|
return hasPending() || hasCompleted();
|
|
}
|
|
|
|
template <typename T>
|
|
folly::Try<T> TaskIterator<T>::awaitNextResult() {
|
|
assert(hasCompleted() || hasPending());
|
|
reserve(1);
|
|
|
|
size_t i = context_->tasksConsumed++;
|
|
id_ = context_->results[i].first;
|
|
return std::move(context_->results[i].second);
|
|
}
|
|
|
|
template <typename T>
|
|
inline T TaskIterator<T>::awaitNext() {
|
|
return std::move(awaitNextResult().value());
|
|
}
|
|
|
|
template <>
|
|
inline void TaskIterator<void>::awaitNext() {
|
|
awaitNextResult().value();
|
|
}
|
|
|
|
template <typename T>
|
|
inline void TaskIterator<T>::reserve(size_t n) {
|
|
size_t tasksReady = context_->results.size() - context_->tasksConsumed;
|
|
|
|
// we don't need to do anything if there are already n or more tasks complete
|
|
// or if we have no tasks left to execute.
|
|
if (!hasPending() || tasksReady >= n) {
|
|
return;
|
|
}
|
|
|
|
n -= tasksReady;
|
|
size_t tasksLeft = context_->totalTasks - context_->results.size();
|
|
n = std::min(n, tasksLeft);
|
|
|
|
await_async([this, n](Promise<void> promise) {
|
|
context_->tasksToFulfillPromise = n;
|
|
context_->promise.assign(std::move(promise));
|
|
});
|
|
}
|
|
|
|
template <typename T>
|
|
inline size_t TaskIterator<T>::getTaskID() const {
|
|
assert(id_ != static_cast<size_t>(-1));
|
|
return id_;
|
|
}
|
|
|
|
template <typename T>
|
|
template <typename F>
|
|
void TaskIterator<T>::addTask(F&& func) {
|
|
static_assert(
|
|
std::is_convertible<invoke_result_t<F>, T>::value,
|
|
"TaskIterator<T>: T must be convertible from func()'s return type");
|
|
|
|
auto taskId = context_->totalTasks++;
|
|
|
|
fm_.addTask(
|
|
[taskId, context = context_, func = std::forward<F>(func)]() mutable {
|
|
context->results.emplace_back(
|
|
taskId, folly::makeTryWith(std::move(func)));
|
|
|
|
// Check for awaiting iterator.
|
|
if (context->promise.hasValue()) {
|
|
if (--context->tasksToFulfillPromise == 0) {
|
|
context->promise->setValue();
|
|
context->promise.clear();
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
template <class InputIterator>
|
|
TaskIterator<
|
|
invoke_result_t<typename std::iterator_traits<InputIterator>::value_type>>
|
|
addTasks(InputIterator first, InputIterator last) {
|
|
typedef invoke_result_t<
|
|
typename std::iterator_traits<InputIterator>::value_type>
|
|
ResultType;
|
|
typedef TaskIterator<ResultType> IteratorType;
|
|
|
|
IteratorType iterator;
|
|
|
|
for (; first != last; ++first) {
|
|
iterator.addTask(std::move(*first));
|
|
}
|
|
|
|
iterator.context_->results.reserve(iterator.context_->totalTasks);
|
|
|
|
return iterator;
|
|
}
|
|
} // namespace fibers
|
|
} // namespace folly
|