library/cpp/unified_agent_client/async_joiner.h


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

#pragma once

#include <library/cpp/threading/future/future.h>

namespace NUnifiedAgent {
    class TAsyncJoiner {
    public:
        inline TAsyncJoiner()
            : Promise(NThreading::NewPromise())
            , Refs(1)
        {
        }

        inline i64 Ref(i64 count = 1) noexcept {
            const auto result = Refs.fetch_add(count);
            Y_ABORT_UNLESS(result >= 1, "already joined");
            return result;
        }

        inline i64 UnRef() noexcept {
            const auto prev = Refs.fetch_sub(1);
            Y_ABORT_UNLESS(prev >= 1);
            if (prev == 1) {
                auto p = Promise;
                p.SetValue();
            }
            return prev;
        }

        inline NThreading::TFuture<void> Join() noexcept {
            auto result = Promise;
            UnRef();
            return result;
        }

    private:
        NThreading::TPromise<void> Promise;
        std::atomic<i64> Refs;
    };

    using TAsyncJoinerToken = TIntrusivePtr<TAsyncJoiner>;
}