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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
|
#include "mkql_spiller_adapter.h"
#include "mock_spiller_factory_ut.h"
#include <library/cpp/testing/unittest/registar.h>
#include <yql/essentials/minikql/mkql_type_builder.h>
#include <yql/essentials/minikql/mkql_node.h>
#include <yql/essentials/minikql/mkql_alloc.h>
#include <yql/essentials/minikql/mkql_string_util.h>
namespace NKikimr::NMiniKQL {
namespace {
struct TMemoryCounters {
ui64 Allocated = 0;
ui64 Freed = 0;
};
TMultiType* CreateTestMultiType(TTypeEnvironment& typeEnv) {
TTypeBuilder builder(typeEnv);
std::vector<TType*> types = {
builder.NewDataType(NUdf::TDataType<ui32>::Id),
builder.NewDataType(NUdf::TDataType<ui64>::Id),
builder.NewDataType(NUdf::TDataType<char*>::Id)
};
return TMultiType::Create(types.size(), types.data(), typeEnv);
}
std::string MakeStringOfSize(ui64 size) {
return std::string(size, 'x');
}
std::vector<NUdf::TUnboxedValue> CreateTestWideItem(ui32 val1, ui64 val2, ui64 stringLen) {
return {
NUdf::TUnboxedValuePod(val1),
NUdf::TUnboxedValuePod(val2),
NMiniKQL::MakeString(MakeStringOfSize(stringLen))
};
}
void VerifyWideItem(const TArrayRef<NUdf::TUnboxedValue>& wideItem, ui32 expectedVal1, ui64 expectedVal2, const TString& expectedStr) {
UNIT_ASSERT_VALUES_EQUAL(wideItem.size(), 3);
UNIT_ASSERT_VALUES_EQUAL(wideItem[0].Get<ui32>(), expectedVal1);
UNIT_ASSERT_VALUES_EQUAL(wideItem[1].Get<ui64>(), expectedVal2);
UNIT_ASSERT_VALUES_EQUAL(TStringBuf(wideItem[2].AsStringRef()), expectedStr);
}
void VerifyCounters(TMemoryCounters& counters, ui64 spilledData) {
UNIT_ASSERT_VALUES_EQUAL(counters.Allocated, counters.Freed);
// estimated packer size is reported. This is enough for spilling purposes.
// So, the size of really spilled data may be slightly different from the reported value.
ui64 delta = 1000;
ui64 lower = 0;
if (spilledData > delta) {
lower = spilledData - delta;
}
ui64 upper = spilledData + delta;
UNIT_ASSERT_GE(counters.Allocated, lower);
UNIT_ASSERT_LE(counters.Allocated, upper);
}
std::shared_ptr<TMockSpillerFactory> MakeFactoryWithCounters(TMemoryCounters& counters) {
auto factory = std::make_shared<TMockSpillerFactory>();
factory->SetMemoryReportingCallbacks(
[&counters](ui64 sz) { counters.Allocated += sz; },
[&counters](ui64 sz) { counters.Freed += sz; }
);
return factory;
}
ui64 RunScenario(
TTypeEnvironment& typeEnv,
THolderFactory& holderFactory,
ui64 flushThreshold,
ui64 reportMemoryThreshold,
ui64 numItems,
TMemoryCounters& counters
) {
auto factory = MakeFactoryWithCounters(counters);
auto spiller = factory->CreateSpiller();
UNIT_ASSERT(spiller != nullptr);
auto multiType = CreateTestMultiType(typeEnv);
TWideUnboxedValuesSpillerAdapter adapter(spiller, multiType, flushThreshold, reportMemoryThreshold);
for (ui64 i = 0; i < numItems; ++i) {
auto future = adapter.WriteWideItem(CreateTestWideItem(static_cast<ui32>(i), 100 * i, i * 100));
if (future.has_value()) {
UNIT_ASSERT(future->HasValue());
adapter.AsyncWriteCompleted(future->GetValue());
}
}
auto finishFuture = adapter.FinishWriting();
if (finishFuture.has_value()) {
adapter.AsyncWriteCompleted(finishFuture->ExtractValue());
}
ui64 item = 0;
while (!adapter.Empty()) {
std::vector<NUdf::TUnboxedValue> readItem(multiType->GetElementsCount());
auto readFuture = adapter.ExtractWideItem(readItem);
if (readFuture.has_value()) {
adapter.AsyncReadCompleted(*readFuture->ExtractValue(), holderFactory);
continue;
}
VerifyWideItem(readItem, static_cast<ui32>(item), 100 * item, MakeStringOfSize(item * 100));
++item;
}
UNIT_ASSERT_VALUES_EQUAL(item, numItems);
const auto mock_spiller = std::dynamic_pointer_cast<TMockSpiller>(spiller);
return mock_spiller->GetTotalSpilled();
}
} // namespace
Y_UNIT_TEST_SUITE(TWideUnboxedValuesSpillerAdapterTest) {
Y_UNIT_TEST(TestBasicReadWrite) {
TScopedAlloc alloc(__LOCATION__);
TMemoryUsageInfo memInfo("test");
THolderFactory holderFactory(alloc.Ref(), memInfo);
TTypeEnvironment env(alloc);
TMemoryCounters counters;
const ui64 flushThreshold = 100;
const ui64 reportMemoryThreshold = 10;
const ui64 numItems = 10;
const ui64 totalSpilled = RunScenario(
env, holderFactory,
flushThreshold, reportMemoryThreshold,
numItems, counters
);
VerifyCounters(counters, totalSpilled);
}
Y_UNIT_TEST(TestLargeLimits) {
TScopedAlloc Alloc(__LOCATION__);
TTypeEnvironment TypeEnv(Alloc);
TMemoryUsageInfo memInfo("test");
auto holderFactory = THolderFactory(Alloc.Ref(), memInfo);
TMemoryCounters counters;
const ui64 flushThreshold = 1_MB;
const ui64 reportMemoryThreshold = 64_KB;
const ui64 numItems = 20;
const ui64 totalSpilled = RunScenario(
TypeEnv, holderFactory,
flushThreshold, reportMemoryThreshold,
numItems, counters
);
VerifyCounters(counters, totalSpilled);
}
Y_UNIT_TEST(TestFrequentMemoryReports) {
TScopedAlloc Alloc(__LOCATION__);
TTypeEnvironment TypeEnv(Alloc);
TMemoryUsageInfo memInfo("test");
auto holderFactory = THolderFactory(Alloc.Ref(), memInfo);
TMemoryCounters counters;
const ui64 flushThreshold = 256_KB;
const ui64 reportMemoryThreshold = 1;
const ui64 numItems = 200;
const ui64 totalSpilled = RunScenario(
TypeEnv, holderFactory,
flushThreshold, reportMemoryThreshold,
numItems, counters
);
VerifyCounters(counters, totalSpilled);
}
Y_UNIT_TEST(TestZeroItems) {
TScopedAlloc Alloc(__LOCATION__);
TTypeEnvironment TypeEnv(Alloc);
TMemoryUsageInfo memInfo("test");
auto holderFactory = THolderFactory(Alloc.Ref(), memInfo);
TMemoryCounters counters;
const ui64 flushThreshold = 128;
const ui64 reportMemoryThreshold = 64;
const ui64 numItems = 0;
const ui64 estimatedPackedSize = RunScenario(
TypeEnv, holderFactory,
flushThreshold, reportMemoryThreshold,
numItems, counters
);
UNIT_ASSERT_EQUAL(estimatedPackedSize, 0);
UNIT_ASSERT_EQUAL(counters.Allocated, 0);
UNIT_ASSERT_EQUAL(counters.Freed, 0);
}
} // Y_UNIT_TEST_SUITE
} // namespace NKikimr::NMiniKQL
|
