1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

#pragma once
/// @file

#include "lix/libutil/result.hh"
#include <concepts>
#include <kj/async.h>
#include <kj/common.h>
#include <kj/vector.h>
#include <list>
#include <optional>
#include <type_traits>

namespace nix {

template<typename K, typename V>
class AsyncCollect
{
public:
    using Item = std::conditional_t<std::is_void_v<V>, K, std::pair<K, V>>;

private:
    kj::ForkedPromise<void> allPromises;
    std::list<Item> results;
    size_t remaining;

    kj::ForkedPromise<void> signal;
    kj::Maybe<kj::Own<kj::PromiseFulfiller<void>>> notify;

    void oneDone(Item item)
    {
        results.emplace_back(std::move(item));
        remaining -= 1;
        KJ_IF_MAYBE (n, notify) {
            (*n)->fulfill();
            notify = nullptr;
        }
    }

    kj::Promise<void> collectorFor(K key, kj::Promise<V> promise)
    {
        if constexpr (std::is_void_v<V>) {
            return promise.then([this, key{std::move(key)}] { oneDone(std::move(key)); });
        } else {
            return promise.then([this, key{std::move(key)}](V v) {
                oneDone(Item{std::move(key), std::move(v)});
            });
        }
    }

    kj::ForkedPromise<void> waitForAll(kj::Array<std::pair<K, kj::Promise<V>>> & promises)
    {
        kj::Vector<kj::Promise<void>> wrappers;
        for (auto & [key, promise] : promises) {
            wrappers.add(collectorFor(std::move(key), std::move(promise)));
        }

        return kj::joinPromisesFailFast(wrappers.releaseAsArray()).fork();
    }

public:
    AsyncCollect(kj::Array<std::pair<K, kj::Promise<V>>> && promises)
        : allPromises(waitForAll(promises))
        , remaining(promises.size())
        , signal{nullptr}
    {
    }

    // oneDone promises capture `this`
    KJ_DISALLOW_COPY_AND_MOVE(AsyncCollect);

    kj::Promise<std::optional<Item>> next()
    {
        if (remaining == 0 && results.empty()) {
            return {std::nullopt};
        }

        if (!results.empty()) {
            auto result = std::move(results.front());
            results.pop_front();
            return {{std::move(result)}};
        }

        if (notify == nullptr) {
            auto pair = kj::newPromiseAndFulfiller<void>();
            notify = std::move(pair.fulfiller);
            signal = pair.promise.fork();
        }

        return signal.addBranch().exclusiveJoin(allPromises.addBranch()).then([this] {
            return next();
        });
    }
};

/**
  * Collect the results of a list of promises, in order of completion.
  * Once any input promise is rejected all promises that have not been
  * resolved or rejected will be cancelled and the exception rethrown.
  */
template<typename K, typename V>
AsyncCollect<K, V> asyncCollect(kj::Array<std::pair<K, kj::Promise<V>>> promises)
{
    return AsyncCollect<K, V>(std::move(promises));
}

/**
 * Run `fn` for every item in the `input` range asynchronously, using
 * the same fail-fast semantics as `asyncCollect`. `asyncSpread` is a
 * shorthand for calling `asyncCollect` with `std::tuple()` values as
 * keys, awaiting that to completion, and propagating all exceptions.
 */
template<typename Input, typename Fn>
kj::Promise<Result<void>> asyncSpread(Input && input, Fn fn)
    requires requires {
        {
            fn(*begin(input))
        } -> std::same_as<kj::Promise<Result<void>>>;
    }
{
    kj::Vector<std::pair<std::tuple<>, kj::Promise<Result<void>>>> children;
    if constexpr (requires { input.size(); }) {
        children.reserve(input.size());
    }

    for (auto & i : input) {
        children.add(std::tuple(), fn(i));
    }

    auto collect = asyncCollect(children.releaseAsArray());
    while (auto r = co_await collect.next()) {
        if (!r->second.has_value()) {
            co_return std::move(r->second);
        }
    }

    co_return result::success();
}

/**
 * Run `promises` concurrently until they've all succeeded, or pass on
 * the error returned by the first failing promise. This is similar to
 * `kj::joinPromisesFailFast()`, but not limited to thrown exceptions.
 */
template<typename... Promises>
kj::Promise<Result<void>> asyncJoin(Promises &&... promises)
    requires(std::same_as<Promises, kj::Promise<Result<void>>> && ...)
{
    auto collect =
        asyncCollect(kj::arr(std::pair{std::tuple{}, std::forward<Promises>(promises)}...));
    while (auto r = co_await collect.next()) {
        if (!r->second.has_value()) {
            co_return std::move(r->second);
        }
    }

    co_return result::success();
}
}