142 lines
4.9 KiB
C++
142 lines
4.9 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.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include <exception>
|
|
#include <memory>
|
|
#include <stdexcept>
|
|
#include <vector>
|
|
|
|
#include <folly/Function.h>
|
|
#include <folly/futures/Future.h>
|
|
#include <folly/futures/Promise.h>
|
|
|
|
namespace folly {
|
|
namespace fibers {
|
|
|
|
/**
|
|
* BatchDispatcher is useful for batching values while doing I/O.
|
|
* For example, if you are launching multiple tasks which take a
|
|
* single id and each task fetches from database, you can use BatchDispatcher
|
|
* to batch those ids and do a single query requesting all those ids.
|
|
*
|
|
* To use this, create a BatchDispatcher with a dispatch function
|
|
* which consumes a vector of values and returns a vector of results
|
|
* in the same order. Add values to BatchDispatcher using add function,
|
|
* which returns a future to the result set in your dispatch function.
|
|
*
|
|
* Implementation Logic:
|
|
* - using FiberManager as executor example, user creates a
|
|
* thread_local BatchDispatcher, on which user calls add(value).
|
|
* - add(value) adds the value in a vector and also schedules a new
|
|
* task(BatchDispatchFunction) which will read the vector of values and call
|
|
* user's DispatchFunction() on it.
|
|
* - assuming the executor queues all the task and runs them in order of their
|
|
* creation time, then BatchDispatcher will run later than all the tasks
|
|
* already created. Depending on this, all the values were added in these
|
|
* tasks would be picked up by BatchDispatchFunction()
|
|
*
|
|
* Example:
|
|
* - User schedules Task1, Task2, Task3 each of them calls BatchDispatch.add()
|
|
* with id1, id2, id3 respectively.
|
|
* - Executor's state {Task1, Task2, Task3}, BatchDispatchers state {}
|
|
* - After Task1 calls BatchDispatcher.add():
|
|
* Executor's state {Task2, Task3, BatchDispatchFunction},
|
|
* BatchDispatcher's state {id1}
|
|
* - After Task2 calls BatchDispatcher.add():
|
|
* Executor's state {Task3, BatchDispatchFunction},
|
|
* BatchDispatcher's state {id1, id2}
|
|
* - After Task3 calls BatchDispatcher.add():
|
|
* Executor's state {BatchDispatchFunction},
|
|
* BatchDispatcher's state {id1, id2, id3}
|
|
* - Now BatchDispatcher calls user's Dispatch function with {id1, id2, id3}
|
|
*
|
|
* Note:
|
|
* - This only works with executors which runs
|
|
* the tasks in order of their schedule time.
|
|
* - BatchDispatcher is not thread safe.
|
|
*/
|
|
template <typename ValueT, typename ResultT, typename ExecutorT>
|
|
class BatchDispatcher {
|
|
public:
|
|
using ValueBatchT = std::vector<ValueT>;
|
|
using ResultBatchT = std::vector<ResultT>;
|
|
using PromiseBatchT = std::vector<folly::Promise<ResultT>>;
|
|
using DispatchFunctionT = folly::Function<ResultBatchT(ValueBatchT&&)>;
|
|
|
|
BatchDispatcher(ExecutorT& executor, DispatchFunctionT dispatchFunc)
|
|
: executor_(executor),
|
|
state_(new DispatchState(std::move(dispatchFunc))) {}
|
|
|
|
Future<ResultT> add(ValueT value) {
|
|
if (state_->values.empty()) {
|
|
executor_.add([state = state_]() { dispatchFunctionWrapper(*state); });
|
|
}
|
|
|
|
folly::Promise<ResultT> resultPromise;
|
|
auto resultFuture = resultPromise.getFuture();
|
|
|
|
state_->values.emplace_back(std::move(value));
|
|
state_->promises.emplace_back(std::move(resultPromise));
|
|
|
|
return resultFuture;
|
|
}
|
|
|
|
private:
|
|
struct DispatchState {
|
|
explicit DispatchState(DispatchFunctionT&& dispatchFunction)
|
|
: dispatchFunc(std::move(dispatchFunction)) {}
|
|
|
|
DispatchFunctionT dispatchFunc;
|
|
ValueBatchT values;
|
|
PromiseBatchT promises;
|
|
};
|
|
|
|
static void dispatchFunctionWrapper(DispatchState& state) {
|
|
ValueBatchT values;
|
|
PromiseBatchT promises;
|
|
state.values.swap(values);
|
|
state.promises.swap(promises);
|
|
|
|
try {
|
|
auto results = state.dispatchFunc(std::move(values));
|
|
if (results.size() != promises.size()) {
|
|
throw std::logic_error(
|
|
"Unexpected number of results returned from dispatch function");
|
|
}
|
|
|
|
for (size_t i = 0; i < promises.size(); i++) {
|
|
promises[i].setValue(std::move(results[i]));
|
|
}
|
|
} catch (const std::exception& ex) {
|
|
for (size_t i = 0; i < promises.size(); i++) {
|
|
promises[i].setException(
|
|
exception_wrapper(std::current_exception(), ex));
|
|
}
|
|
} catch (...) {
|
|
for (size_t i = 0; i < promises.size(); i++) {
|
|
promises[i].setException(exception_wrapper(std::current_exception()));
|
|
}
|
|
}
|
|
}
|
|
|
|
ExecutorT& executor_;
|
|
std::shared_ptr<DispatchState> state_;
|
|
};
|
|
} // namespace fibers
|
|
} // namespace folly
|