aboutsummaryrefslogtreecommitdiffstats
path: root/library/cpp/yt/memory/memory_usage_tracker.cpp
blob: 4055bf62c5eae1f77769c2ba17864b3b8ad0d18a (plain) (blame)
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
#include "memory_usage_tracker.h"

#include "leaky_ref_counted_singleton.h"

namespace NYT {

////////////////////////////////////////////////////////////////////////////////

class TNulSimplelMemoryUsageTracker
    : public ISimpleMemoryUsageTracker
{
public:
    bool Acquire(i64 /*size*/) override
    {
        return false;
    }

    void Release(i64 /*size*/) override
    { }

    TSharedRef Track(TSharedRef reference, bool /*keepExistingTracking*/) override
    {
        return reference;
    }
};

////////////////////////////////////////////////////////////////////////////////

TSimpleMemoryUsageTrackerGuard::TSimpleMemoryUsageTrackerGuard(TSimpleMemoryUsageTrackerGuard&& other)
{
    MoveFrom(std::move(other));
}

TSimpleMemoryUsageTrackerGuard::~TSimpleMemoryUsageTrackerGuard()
{
    Release();
}

TSimpleMemoryUsageTrackerGuard& TSimpleMemoryUsageTrackerGuard::operator=(TSimpleMemoryUsageTrackerGuard&& other)
{
    if (this != &other) {
        Release();
        MoveFrom(std::move(other));
    }
    return *this;
}

void TSimpleMemoryUsageTrackerGuard::MoveFrom(TSimpleMemoryUsageTrackerGuard&& other)
{
    Tracker_ = other.Tracker_;
    AcquiredSize_ = other.AcquiredSize_;

    other.Tracker_ = nullptr;
    other.AcquiredSize_ = 0;
}

TSimpleMemoryUsageTrackerGuard TSimpleMemoryUsageTrackerGuard::Build(ISimpleMemoryUsageTrackerPtr tracker)
{
    if (!tracker) {
        return {};
    }

    TSimpleMemoryUsageTrackerGuard guard;
    guard.Tracker_ = tracker;
    return guard;
}

void TSimpleMemoryUsageTrackerGuard::Release()
{
    if (Tracker_) {
        if (AcquiredSize_) {
            Tracker_->Release(AcquiredSize_);
        }

        Tracker_.Reset();
        AcquiredSize_ = 0;
    }
}

void TSimpleMemoryUsageTrackerGuard::SetSize(i64 size)
{
    if (!Tracker_) {
        return;
    }

    YT_VERIFY(size >= 0);

    if (size > AcquiredSize_) {
        Tracker_->Acquire(size - AcquiredSize_);
    } else if (size < AcquiredSize_) {
        Tracker_->Release(AcquiredSize_ - size);
    }

    AcquiredSize_ = size;
}

////////////////////////////////////////////////////////////////////////////////

ISimpleMemoryUsageTrackerPtr GetNullSimpleMemoryUsageTracker()
{
    return LeakyRefCountedSingleton<TNulSimplelMemoryUsageTracker>();
}

TSharedRef TrackMemory(
    const ISimpleMemoryUsageTrackerPtr& tracker,
    TSharedRef reference,
    bool keepExistingTracking)
{
    if (!tracker || !reference) {
        return reference;
    }
    return tracker->Track(std::move(reference), keepExistingTracking);
}

////////////////////////////////////////////////////////////////////////////////

} // namespace NYT