Skip to main content

reifydb_store_multi/store/
worker.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2026 ReifyDB
3
4use std::{collections::HashMap, time::Duration};
5
6use reifydb_core::{
7	common::CommitVersion,
8	encoded::key::EncodedKey,
9	event::{
10		EventBus,
11		metric::{MultiCommittedEvent, MultiDrop},
12	},
13	interface::store::EntryKind,
14};
15use reifydb_runtime::{
16	actor::{
17		context::Context,
18		mailbox::ActorRef,
19		system::{ActorConfig, ActorSystem},
20		timers::TimerHandle,
21		traits::{Actor, Directive},
22	},
23	context::clock::{Clock, Instant},
24};
25use tracing::{Span, debug, error, instrument};
26
27use super::drop::find_keys_to_drop;
28use crate::tier::{TierStorage, commit::buffer::MultiCommitBufferTier};
29
30#[derive(Debug, Clone)]
31pub struct DropWorkerConfig {
32	pub batch_size: usize,
33
34	pub flush_interval: Duration,
35}
36
37impl Default for DropWorkerConfig {
38	fn default() -> Self {
39		Self {
40			batch_size: 100,
41			flush_interval: Duration::from_millis(50),
42		}
43	}
44}
45
46use reifydb_core::actors::drop::{DropMessage, DropRequest};
47
48pub struct DropActor {
49	storage: MultiCommitBufferTier,
50	event_bus: EventBus,
51	config: DropWorkerConfig,
52	clock: Clock,
53}
54
55pub struct DropActorState {
56	pending_requests: Vec<DropRequest>,
57
58	last_flush: Instant,
59
60	_timer_handle: Option<TimerHandle>,
61
62	flush_count: u64,
63}
64
65impl DropActor {
66	pub fn new(
67		config: DropWorkerConfig,
68		storage: MultiCommitBufferTier,
69		event_bus: EventBus,
70		clock: Clock,
71	) -> Self {
72		Self {
73			storage,
74			event_bus,
75			config,
76			clock,
77		}
78	}
79
80	pub fn spawn(
81		system: &ActorSystem,
82		config: DropWorkerConfig,
83		storage: MultiCommitBufferTier,
84		event_bus: EventBus,
85		clock: Clock,
86	) -> ActorRef<DropMessage> {
87		let actor = Self::new(config, storage, event_bus, clock);
88		system.spawn_system("drop-worker", actor).actor_ref().clone()
89	}
90
91	fn maybe_flush(&self, state: &mut DropActorState) {
92		if state.pending_requests.len() >= self.config.batch_size {
93			self.flush(state);
94		}
95	}
96
97	fn flush(&self, state: &mut DropActorState) {
98		if state.pending_requests.is_empty() {
99			return;
100		}
101
102		Self::process_batch(&self.storage, &mut state.pending_requests, &self.event_bus);
103		state.last_flush = self.clock.instant();
104
105		state.flush_count += 1;
106		if state.flush_count.is_multiple_of(100) {
107			self.storage.maintenance();
108		}
109	}
110
111	#[instrument(name = "drop::process_batch", level = "debug", skip_all, fields(num_requests = requests.len(), total_dropped))]
112	fn process_batch(storage: &MultiCommitBufferTier, requests: &mut Vec<DropRequest>, event_bus: &EventBus) {
113		let mut batches: HashMap<EntryKind, Vec<(EncodedKey, CommitVersion)>> = HashMap::new();
114
115		let mut drops_with_stats = Vec::new();
116		let mut max_pending_version = CommitVersion(0);
117
118		for request in requests.drain(..) {
119			let version_for_event = request.pending_version.unwrap_or(request.commit_version);
120			if version_for_event > max_pending_version {
121				max_pending_version = version_for_event;
122			}
123
124			match find_keys_to_drop(storage, request.table, request.key.as_ref(), request.pending_version) {
125				Ok(entries_to_drop) => {
126					for entry in entries_to_drop {
127						drops_with_stats.push(MultiDrop {
128							key: request.key.clone(),
129							value_bytes: entry.value_bytes,
130						});
131
132						batches.entry(request.table)
133							.or_default()
134							.push((entry.key, entry.version));
135					}
136				}
137				Err(e) => {
138					error!("Drop actor failed to find keys to drop: {}", e);
139				}
140			}
141		}
142
143		if !batches.is_empty()
144			&& let Err(e) = storage.drop(batches)
145		{
146			error!("Drop actor failed to execute drops: {}", e);
147		}
148
149		let total_dropped = drops_with_stats.len();
150		Span::current().record("total_dropped", total_dropped);
151
152		event_bus.emit(MultiCommittedEvent::new(vec![], vec![], drops_with_stats, max_pending_version));
153	}
154}
155
156impl Actor for DropActor {
157	type State = DropActorState;
158	type Message = DropMessage;
159
160	fn init(&self, ctx: &Context<Self::Message>) -> Self::State {
161		debug!("Drop actor started");
162
163		let timer_handle = ctx.schedule_repeat(Duration::from_millis(10), DropMessage::Tick);
164
165		DropActorState {
166			pending_requests: Vec::with_capacity(self.config.batch_size),
167			last_flush: self.clock.instant(),
168			_timer_handle: Some(timer_handle),
169			flush_count: 0,
170		}
171	}
172
173	fn handle(&self, state: &mut Self::State, msg: Self::Message, ctx: &Context<Self::Message>) -> Directive {
174		if ctx.is_cancelled() {
175			self.flush(state);
176			return Directive::Stop;
177		}
178
179		match msg {
180			DropMessage::Request(request) => {
181				state.pending_requests.push(request);
182				self.maybe_flush(state);
183			}
184			DropMessage::Batch(requests) => {
185				state.pending_requests.extend(requests);
186				self.maybe_flush(state);
187			}
188			DropMessage::Tick => {
189				if !state.pending_requests.is_empty()
190					&& state.last_flush.elapsed() >= self.config.flush_interval
191				{
192					self.flush(state);
193				}
194			}
195			DropMessage::Shutdown => {
196				debug!("Drop actor received shutdown signal");
197
198				self.flush(state);
199				return Directive::Stop;
200			}
201		}
202
203		Directive::Continue
204	}
205
206	fn post_stop(&self) {
207		debug!("Drop actor stopped");
208	}
209
210	fn config(&self) -> ActorConfig {
211		ActorConfig::new().mailbox_capacity(4096 * 16)
212	}
213}