forked from ThePhD/sol2
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathallocation_limit.cpp
More file actions
135 lines (115 loc) · 3.95 KB
/
Copy pathallocation_limit.cpp
File metadata and controls
135 lines (115 loc) · 3.95 KB
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
#define SOL_ALL_SAFETIES_ON 1
#include <sol/sol.hpp>
#include <cstddef>
#include <cstdlib>
class memory_tracker {
public:
// 10 MB or something?
// idk whatever
inline static constexpr std::size_t arbitrary_default_limit = 1024 * 1024 * 10;
memory_tracker() : memory_tracker(arbitrary_default_limit) {
}
memory_tracker(std::size_t maximum_memory) : used(0), limit(maximum_memory), n_threads(0), n_tables(0), n_functions(0), n_userdata(0), n_strings(0) {
}
std::size_t currently_used() const {
return used;
}
std::size_t memory_limit() const {
return limit;
}
static void* allocate(void* memory_tracker_ud, void* ptr, size_t object_code, size_t nsize) {
memory_tracker& self = (*static_cast<memory_tracker*>(memory_tracker_ud));
return self.alloc(ptr, object_code, nsize);
}
private:
void* alloc(void* ptr, size_t original_block_size_or_code, size_t new_block_size) {
std::size_t original_block_size = original_block_size_or_code;
if (ptr == nullptr) {
// object code!
sol::type object_type = static_cast<sol::type>(original_block_size_or_code);
switch (object_type) {
case sol::type::function:
++n_functions;
break;
case sol::type::string:
++n_strings;
break;
case sol::type::userdata:
++n_userdata;
break;
case sol::type::table:
++n_tables;
break;
case sol::type::thread:
++n_threads;
break;
default:
// not a clue, fam
break;
}
// because it is an object code,
// it tells us literally nothing about
// the old block size,
// so set that to 0
original_block_size = 0;
}
if (new_block_size == 0) {
// Lua expects us to act like a "free"
// when the new block size is 0
std::free(ptr);
used -= original_block_size;
return nullptr;
}
// did we hit the limit?
std::size_t memory_differntial = new_block_size - original_block_size;
std::size_t desired_use = used + memory_differntial;
if (desired_use > limit) {
// tell the Lua Virtual Machine
// to toss off (by returning nullptr)
return nullptr;
}
// alright now we have to expand this shit
// guess we use C's realloc
ptr = std::realloc(ptr, new_block_size);
if (ptr != nullptr) {
// alright, we successfully allocated some space
// track it
used = desired_use;
}
// even if we're null, let Lua crash and burn for us
return ptr;
}
std::size_t used;
std::size_t limit;
std::size_t n_threads;
std::size_t n_tables;
std::size_t n_functions;
std::size_t n_userdata;
std::size_t n_strings;
};
struct my_type {
int a = 24;
bool b = false;
double d = 3.5;
};
#include <iostream>
int main() {
std::cout << "=== memory tracker ===" << std::endl;
#if defined(SOL_LUAJIT) && SOL_LUA_VERSION < 20100 && (UINTPTR_MAX > 0xFFFFFFFF)
std::cout << "LuaJIT in x64 mode on LuaJIT 2.0.X versions does not support using a custom allocator!" << std::endl;
#else
memory_tracker box;
std::cout << "memory at start: " << box.currently_used() << " bytes / " << box.memory_limit() << " bytes" << std::endl;
sol::state lua(&sol::default_at_panic, &memory_tracker::allocate, &box);
lua.open_libraries(sol::lib::base);
std::cout << "memory after state creation: " << box.currently_used() << " bytes / " << box.memory_limit() << " bytes" << std::endl;
// trigger some allocations
lua.new_usertype<my_type>("my_type", "a", &my_type::a, "b", &my_type::b, "d", &my_type::d);
lua["f"] = [](std::string s) { return s + " woof!"; };
std::cout << "memory after function and usertype setup: " << box.currently_used() << " bytes / " << box.memory_limit() << " bytes" << std::endl;
lua.safe_script("print(f('bark'))");
lua.safe_script("local obj = my_type.new() print(obj.a, obj.b, obj.c) obj.b = true print(obj.a, obj.b, obj.c)");
std::cout << "memory at end: " << box.currently_used() << " bytes / " << box.memory_limit() << " bytes" << std::endl;
#endif
return 0;
}