tdlight-team/tdlight
tdlight-team/tdlight
1
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add limit on total TQueue events size.

Browse Source 
GitOrigin-RevId: cd0d77e4b0e2d66d7c5253a6d36bd543bca95bd1
This commit is contained in:
levlam 2020-07-13 23:48:15 +03:00
parent e16952121b
commit 32f9e4924f
1 changed files with 23 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
tddb/td/db/TQueue.cpp
View File

@ -84,8 +84,9 @@ bool EventId::is_valid_id(int32 id) {
} }
class TQueueImpl : public TQueue { class TQueueImpl : public TQueue {
static constexpr size_t MAX_EVENT_LEN = 65536 * 8; static constexpr size_t MAX_EVENT_LENGTH = 65536 * 8;
static constexpr size_t MAX_QUEUE_EVENTS = 1000000; static constexpr size_t MAX_QUEUE_EVENTS = 1000000;
static constexpr size_t MAX_TOTAL_EVENT_LENGTH = 1 << 30;
public: public:
void set_callback(unique_ptr<StorageCallback> callback) override { void set_callback(unique_ptr<StorageCallback> callback) override {
@ -103,6 +104,7 @@ class TQueueImpl : public TQueue {
raw_event.logevent_id = callback_->push(queue_id, raw_event); raw_event.logevent_id = callback_->push(queue_id, raw_event);
} }
q.tail_id = raw_event.event_id.next().move_as_ok(); q.tail_id = raw_event.event_id.next().move_as_ok();
q.total_event_length += raw_event.data.size();
q.events.push(std::move(raw_event)); q.events.push(std::move(raw_event));
return true; return true;
} }
@ -117,9 +119,12 @@ class TQueueImpl : public TQueue {
if (data.empty()) { if (data.empty()) {
return Status::Error("Data is empty"); return Status::Error("Data is empty");
} }
if (data.size() > MAX_EVENT_LEN) { if (data.size() > MAX_EVENT_LENGTH) {
return Status::Error("Data is too big"); return Status::Error("Data is too big");
} }
if (q.total_event_length > MAX_TOTAL_EVENT_LENGTH - data.size()) {
return Status::Error("Queue size is too big");
}
EventId event_id; EventId event_id;
while (true) { while (true) {
if (q.tail_id.empty()) { if (q.tail_id.empty()) {
@ -135,7 +140,7 @@ class TQueueImpl : public TQueue {
break; break;
} }
for (auto &event : q.events.as_mutable_span()) { for (auto &event : q.events.as_mutable_span()) {
pop(queue_id, event, {}); pop(q, queue_id, event, {});
} }
q.tail_id = EventId(); q.tail_id = EventId();
q.events = {}; q.events = {};
@ -185,7 +190,7 @@ class TQueueImpl : public TQueue {
if (it == from_events.end() || !(it->event_id == event_id)) { if (it == from_events.end() || !(it->event_id == event_id)) {
return; return;
} }
pop(queue_id, *it, q.tail_id); pop(q, queue_id, *it, q.tail_id);
} }
Result<size_t> get(QueueId queue_id, EventId from_id, bool forget_previous, double now, Result<size_t> get(QueueId queue_id, EventId from_id, bool forget_previous, double now,
@ -210,7 +215,7 @@ class TQueueImpl : public TQueue {
void run_gc(double now) override { void run_gc(double now) override {
for (auto &it : queues_) { for (auto &it : queues_) {
for (auto &e : it.second.events.as_mutable_span()) { for (auto &e : it.second.events.as_mutable_span()) {
try_pop(it.first, e, EventId(), it.second.tail_id, now); try_pop(it.second, it.first, e, EventId(), now);
} }
} }
} }
@ -233,6 +238,7 @@ class TQueueImpl : public TQueue {
struct Queue { struct Queue {
EventId tail_id; EventId tail_id;
VectorQueue<RawEvent> events; VectorQueue<RawEvent> events;
size_t total_event_length = 0;
}; };
std::unordered_map<QueueId, Queue> queues_; std::unordered_map<QueueId, Queue> queues_;
@ -249,31 +255,36 @@ class TQueueImpl : public TQueue {
events.pop_n(removed_n); events.pop_n(removed_n);
} }
void try_pop(QueueId queue_id, RawEvent &event, EventId from_id, EventId tail_id, double now) { void try_pop(Queue &q, QueueId queue_id, RawEvent &event, EventId from_id, double now) {
if (event.expires_at < now || event.event_id < from_id || event.data.empty()) { if (event.expires_at < now || event.event_id < from_id || event.data.empty()) {
pop(queue_id, event, tail_id); pop(q, queue_id, event, q.tail_id);
} }
} }
void pop(QueueId queue_id, RawEvent &event, EventId tail_id) { void pop(Queue &q, QueueId queue_id, RawEvent &event, EventId tail_id) {
if (callback_ == nullptr || event.logevent_id == 0) { if (callback_ == nullptr || event.logevent_id == 0) {
event.logevent_id = 0; event.logevent_id = 0;
event.data = {}; clear_event_data(q, event);
return; return;
} }
if (event.event_id.next().ok() == tail_id) { if (event.event_id.next().ok() == tail_id) {
if (!event.data.empty()) { if (!event.data.empty()) {
event.data = {}; clear_event_data(q, event);
callback_->push(queue_id, event); callback_->push(queue_id, event);
} }
} else { } else {
callback_->pop(event.logevent_id); callback_->pop(event.logevent_id);
event.logevent_id = 0; event.logevent_id = 0;
event.data = {}; clear_event_data(q, event);
} }
} }
static void clear_event_data(Queue &q, RawEvent &event) {
q.total_event_length -= event.data.size();
event.data = {};
}
size_t do_get(QueueId queue_id, Queue &q, EventId from_id, bool forget_previous, double now, size_t do_get(QueueId queue_id, Queue &q, EventId from_id, bool forget_previous, double now,
MutableSpan<Event> &result_events) { MutableSpan<Event> &result_events) {
MutableSpan<RawEvent> from_events; MutableSpan<RawEvent> from_events;
@ -291,7 +302,7 @@ class TQueueImpl : public TQueue {
} }
for (i = first_i; i < from_events.size(); i++) { for (i = first_i; i < from_events.size(); i++) {
auto &from = from_events[i]; auto &from = from_events[i];
try_pop(queue_id, from, forget_previous ? from_id : EventId{}, q.tail_id, now); try_pop(q, queue_id, from, forget_previous ? from_id : EventId{}, now);
if (from.data.empty()) { if (from.data.empty()) {
continue; continue;
} }