/build/source/nativelink-store/src/redis_utils/ft_aggregate.rs
Line | Count | Source |
1 | | // Copyright 2024-2025 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); |
4 | | // you may not use this file except in compliance with the License. |
5 | | // You may obtain a copy of the License at |
6 | | // |
7 | | // See LICENSE file for details |
8 | | // |
9 | | // Unless required by applicable law or agreed to in writing, software |
10 | | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | | // See the License for the specific language governing permissions and |
13 | | // limitations under the License. |
14 | | |
15 | | use core::fmt::Debug; |
16 | | |
17 | | use futures::Stream; |
18 | | use redis::aio::ConnectionLike; |
19 | | use redis::{Arg, ErrorKind, RedisError, Value}; |
20 | | use tracing::error; |
21 | | |
22 | | use crate::redis_utils::aggregate_types::RedisCursorData; |
23 | | use crate::redis_utils::ft_cursor_read::ft_cursor_read; |
24 | | |
25 | | #[derive(Debug)] |
26 | | pub(crate) struct FtAggregateCursor { |
27 | | pub count: u64, |
28 | | pub max_idle: u64, |
29 | | } |
30 | | |
31 | | #[derive(Debug)] |
32 | | pub(crate) struct FtAggregateOptions { |
33 | | pub load: Vec<String>, |
34 | | pub cursor: FtAggregateCursor, |
35 | | pub sort_by: Vec<String>, |
36 | | } |
37 | | |
38 | | /// Per-query `FT.AGGREGATE` timeout in milliseconds. |
39 | | /// |
40 | | /// `RediSearch`'s module default (≈500 ms) is far too tight for the |
41 | | /// scheduler's awaited-action index under any meaningful load: queries |
42 | | /// time out, `NativeLink` surfaces them as parse errors, and the dedup |
43 | | /// lookup fails. When dedup fails the scheduler creates a duplicate |
44 | | /// operation for an action that is already in flight — observed as |
45 | | /// "two same actions running on different PRs" with each running the |
46 | | /// full `max_action_executing_timeout_s` window before completing. Pass an |
47 | | /// explicit value generous enough to absorb 1M+ document scans on a |
48 | | /// busy `RediSearch` instance. |
49 | | const FT_AGGREGATE_TIMEOUT_MS: u64 = 10_000; |
50 | | |
51 | | /// Calls `FT.AGGREGATE` in redis. redis-rs does not properly support this command |
52 | | /// so we have to manually handle it. |
53 | 36 | pub(crate) async fn ft_aggregate<C>( |
54 | 36 | mut connection_manager: C, |
55 | 36 | index: String, |
56 | 36 | query: String, |
57 | 36 | options: FtAggregateOptions, |
58 | 36 | ) -> Result<impl Stream<Item = Result<Value, RedisError>> + Send, RedisError> |
59 | 36 | where |
60 | 36 | C: ConnectionLike + Send, |
61 | 36 | { |
62 | | struct State<C: ConnectionLike> { |
63 | | connection_manager: C, |
64 | | index: String, |
65 | | data: RedisCursorData, |
66 | | } |
67 | | |
68 | 36 | let mut cmd = redis::cmd("FT.AGGREGATE"); |
69 | 36 | let mut ft_aggregate_cmd = cmd |
70 | 36 | .arg(&index) |
71 | 36 | .arg(&query) |
72 | 36 | .arg("TIMEOUT") |
73 | 36 | .arg(FT_AGGREGATE_TIMEOUT_MS) |
74 | 36 | .arg("LOAD") |
75 | 36 | .arg(options.load.len()) |
76 | 36 | .arg(&options.load) |
77 | 36 | .arg("WITHCURSOR") |
78 | 36 | .arg("COUNT") |
79 | 36 | .arg(options.cursor.count) |
80 | 36 | .arg("MAXIDLE") |
81 | 36 | .arg(options.cursor.max_idle) |
82 | 36 | .arg("SORTBY") |
83 | 36 | .arg(options.sort_by.len() * 2); |
84 | 36 | for key30 in &options.sort_by { |
85 | 30 | ft_aggregate_cmd = ft_aggregate_cmd.arg(key).arg("ASC"); |
86 | 30 | } |
87 | 36 | let res = ft_aggregate_cmd |
88 | 36 | .query_async::<Value>(&mut connection_manager) |
89 | 36 | .await; |
90 | 36 | let data29 = match res { |
91 | 29 | Ok(d) => d, |
92 | 7 | Err(e) => { |
93 | 7 | let all_args: Vec<_> = ft_aggregate_cmd |
94 | 7 | .args_iter() |
95 | 126 | .map7 (|a| match a { |
96 | 126 | Arg::Simple(bytes) => match str::from_utf8(bytes) { |
97 | 126 | Ok(s) => s.to_string(), |
98 | 0 | Err(_) => format!("{bytes:?}"), |
99 | | }, |
100 | 0 | other => { |
101 | 0 | format!("{other:?}") |
102 | | } |
103 | 126 | }) |
104 | 7 | .collect(); |
105 | 7 | error!( |
106 | | ?e, |
107 | | index, |
108 | | ?query, |
109 | | ?options, |
110 | | ?all_args, |
111 | | "Error calling ft.aggregate" |
112 | | ); |
113 | 7 | return Err(e); |
114 | | } |
115 | | }; |
116 | | |
117 | 29 | let state = State { |
118 | 29 | connection_manager, |
119 | 29 | index, |
120 | 29 | data: data.try_into()?0 , |
121 | | }; |
122 | | |
123 | 29 | Ok(futures::stream::unfold( |
124 | 29 | Some(state), |
125 | 46 | move |maybe_state| async move { |
126 | 46 | let mut state = maybe_state?0 ; |
127 | | loop { |
128 | 47 | if let Some(map19 ) = state.data.data.pop_front() { |
129 | 19 | return Some((Ok(map), Some(state))); |
130 | 28 | } |
131 | 28 | if state.data.cursor == 0 { |
132 | 27 | return None; |
133 | 1 | } |
134 | 1 | let data_res = ft_cursor_read( |
135 | 1 | &mut state.connection_manager, |
136 | 1 | state.index.clone(), |
137 | 1 | state.data.cursor, |
138 | 1 | ) |
139 | 1 | .await; |
140 | 1 | state.data = match data_res { |
141 | 1 | Ok(data) => data, |
142 | 0 | Err(err) => return Some((Err(err), None)), |
143 | | }; |
144 | | } |
145 | 92 | }, |
146 | | )) |
147 | 36 | } |
148 | | |
149 | 27 | fn resp2_data_parse( |
150 | 27 | output: &mut RedisCursorData, |
151 | 27 | results_array: &[Value], |
152 | 27 | ) -> Result<(), RedisError> { |
153 | 27 | let mut results_iter = results_array.iter(); |
154 | 27 | match results_iter.next() { |
155 | 27 | Some(Value::Int(t)) => { |
156 | 27 | output.total = *t; |
157 | 27 | } |
158 | 0 | Some(other) => { |
159 | 0 | error!(?other, "Non-int for first value in ft.aggregate"); |
160 | 0 | return Err(RedisError::from(( |
161 | 0 | ErrorKind::Parse, |
162 | 0 | "Non int for aggregate total", |
163 | 0 | format!("{other:?}"), |
164 | 0 | ))); |
165 | | } |
166 | | None => { |
167 | 0 | error!("No items in results array for ft.aggregate!"); |
168 | 0 | return Err(RedisError::from(( |
169 | 0 | ErrorKind::Parse, |
170 | 0 | "No items in results array for ft.aggregate", |
171 | 0 | ))); |
172 | | } |
173 | | } |
174 | | |
175 | 27 | for item14 in results_iter { |
176 | 14 | match item { |
177 | 14 | Value::Array(items) if items.len() % 2 == 0 => {} |
178 | 0 | other => { |
179 | 0 | error!( |
180 | | ?other, |
181 | | "Expected an array with an even number of items, didn't get it for aggregate value" |
182 | | ); |
183 | 0 | return Err(RedisError::from(( |
184 | 0 | ErrorKind::Parse, |
185 | 0 | "Expected an array with an even number of items, didn't get it for aggregate value", |
186 | 0 | format!("{other:?}"), |
187 | 0 | ))); |
188 | | } |
189 | | } |
190 | | |
191 | 14 | output.data.push_back(item.clone()); |
192 | | } |
193 | 27 | Ok(()) |
194 | 27 | } |
195 | | |
196 | 2 | fn resp3_data_parse( |
197 | 2 | output: &mut RedisCursorData, |
198 | 2 | results_map: &Vec<(Value, Value)>, |
199 | 2 | ) -> Result<(), RedisError> { |
200 | 8 | for (raw_key, value) in results_map2 { |
201 | 8 | let Value::SimpleString(key) = raw_key else { |
202 | 0 | return Err(RedisError::from(( |
203 | 0 | ErrorKind::Parse, |
204 | 0 | "Expected SimpleString keys", |
205 | 0 | format!("{raw_key:?}"), |
206 | 0 | ))); |
207 | | }; |
208 | 8 | match key.as_str() { |
209 | 8 | "attributes" => { |
210 | 2 | let Value::Array(attributes) = value else { |
211 | 0 | return Err(RedisError::from(( |
212 | 0 | ErrorKind::Parse, |
213 | 0 | "Expected array for attributes", |
214 | 0 | format!("{value:?}"), |
215 | 0 | ))); |
216 | | }; |
217 | 2 | if !attributes.is_empty() { |
218 | 0 | return Err(RedisError::from(( |
219 | 0 | ErrorKind::Parse, |
220 | 0 | "Expected empty attributes", |
221 | 0 | format!("{attributes:?}"), |
222 | 0 | ))); |
223 | 2 | } |
224 | | } |
225 | 6 | "format" => { |
226 | 2 | let Value::SimpleString(format) = value else { |
227 | 0 | return Err(RedisError::from(( |
228 | 0 | ErrorKind::Parse, |
229 | 0 | "Expected SimpleString for format", |
230 | 0 | format!("{value:?}"), |
231 | 0 | ))); |
232 | | }; |
233 | 2 | if format.as_str() != "STRING" { |
234 | 0 | return Err(RedisError::from(( |
235 | 0 | ErrorKind::Parse, |
236 | 0 | "Expected STRING format", |
237 | 0 | format.clone(), |
238 | 0 | ))); |
239 | 2 | } |
240 | | } |
241 | 4 | "results" => { |
242 | 2 | let Value::Array(values) = value else { |
243 | 0 | return Err(RedisError::from(( |
244 | 0 | ErrorKind::Parse, |
245 | 0 | "Expected Array for results", |
246 | 0 | format!("{value:?}"), |
247 | 0 | ))); |
248 | | }; |
249 | 4 | for raw_value in values2 { |
250 | 4 | let Value::Map(value) = raw_value else { |
251 | 0 | return Err(RedisError::from(( |
252 | 0 | ErrorKind::Parse, |
253 | 0 | "Expected list of maps in result", |
254 | 0 | format!("{raw_value:?}"), |
255 | 0 | ))); |
256 | | }; |
257 | 8 | for (raw_map_key, raw_map_value) in value4 { |
258 | 8 | let Value::SimpleString(map_key) = raw_map_key else { |
259 | 0 | return Err(RedisError::from(( |
260 | 0 | ErrorKind::Parse, |
261 | 0 | "Expected SimpleString keys for result maps", |
262 | 0 | format!("{raw_key:?}"), |
263 | 0 | ))); |
264 | | }; |
265 | 8 | match map_key.as_str() { |
266 | 8 | "extra_attributes" => { |
267 | 4 | let extra_attributes_values3 = match raw_map_value { |
268 | 3 | Value::Map(extra_attributes_values) => extra_attributes_values, |
269 | | // A document that expired or was deleted between |
270 | | // the search phase and the load phase comes back |
271 | | // as a row with Nil attributes. Under load this is |
272 | | // routine — completed awaited-action records expire |
273 | | // constantly — so drop the row instead of failing |
274 | | // the whole aggregate. Failing here surfaced to |
275 | | // clients as `INVALID_ARGUMENT`, which Bazel treats |
276 | | // as permanent, so a single expiry race killed the |
277 | | // build. |
278 | 1 | Value::Nil => continue, |
279 | 0 | other => { |
280 | 0 | return Err(RedisError::from(( |
281 | 0 | ErrorKind::Parse, |
282 | 0 | "Expected Map for extra_attributes", |
283 | 0 | format!("{other:?}"), |
284 | 0 | ))); |
285 | | } |
286 | | }; |
287 | 3 | let mut output_array = vec![]; |
288 | 6 | for (e_key, e_value) in extra_attributes_values3 { |
289 | 6 | output_array.push(e_key.clone()); |
290 | 6 | output_array.push(e_value.clone()); |
291 | 6 | } |
292 | 3 | output.data.push_back(Value::Array(output_array)); |
293 | | } |
294 | 4 | "values" => { |
295 | 4 | let Value::Array(values_values) = raw_map_value else { |
296 | 0 | return Err(RedisError::from(( |
297 | 0 | ErrorKind::Parse, |
298 | 0 | "Expected Array for values", |
299 | 0 | format!("{raw_map_value:?}"), |
300 | 0 | ))); |
301 | | }; |
302 | 4 | if !values_values.is_empty() { |
303 | 0 | return Err(RedisError::from(( |
304 | 0 | ErrorKind::Parse, |
305 | 0 | "Expected empty values (all in extra_attributes)", |
306 | 0 | format!("{values_values:?}"), |
307 | 0 | ))); |
308 | 4 | } |
309 | | } |
310 | | _ => { |
311 | 0 | return Err(RedisError::from(( |
312 | 0 | ErrorKind::Parse, |
313 | 0 | "Unknown result map key", |
314 | 0 | format!("{map_key:?}"), |
315 | 0 | ))); |
316 | | } |
317 | | } |
318 | | } |
319 | | } |
320 | | } |
321 | 2 | "total_results" => { |
322 | 1 | let Value::Int(total) = value else { |
323 | 0 | return Err(RedisError::from(( |
324 | 0 | ErrorKind::Parse, |
325 | 0 | "Expected int for total_results", |
326 | 0 | format!("{value:?}"), |
327 | 0 | ))); |
328 | | }; |
329 | 1 | output.total = *total; |
330 | | } |
331 | 1 | "warning" => { |
332 | 1 | let Value::Array(warnings) = value else { |
333 | 0 | return Err(RedisError::from(( |
334 | 0 | ErrorKind::Parse, |
335 | 0 | "Expected Array for warning", |
336 | 0 | format!("{value:?}"), |
337 | 0 | ))); |
338 | | }; |
339 | 1 | if !warnings.is_empty() { |
340 | 0 | return Err(RedisError::from(( |
341 | 0 | ErrorKind::Parse, |
342 | 0 | "Expected empty warnings", |
343 | 0 | format!("{warnings:?}"), |
344 | 0 | ))); |
345 | 1 | } |
346 | | } |
347 | | _ => { |
348 | 0 | return Err(RedisError::from(( |
349 | 0 | ErrorKind::Parse, |
350 | 0 | "Unexpected key in ft.aggregate", |
351 | 0 | format!("{key} => {value:?}"), |
352 | 0 | ))); |
353 | | } |
354 | | } |
355 | | } |
356 | 2 | Ok(()) |
357 | 2 | } |
358 | | |
359 | | impl TryFrom<Value> for RedisCursorData { |
360 | | type Error = RedisError; |
361 | 29 | fn try_from(raw_value: Value) -> Result<Self, RedisError> { |
362 | 29 | let Value::Array(value) = raw_value else { |
363 | 0 | error!( |
364 | | ?raw_value, |
365 | | "Bad data in ft.aggregate, expected array at top-level" |
366 | | ); |
367 | 0 | return Err(RedisError::from((ErrorKind::Parse, "Expected array"))); |
368 | | }; |
369 | 29 | if value.len() < 2 { |
370 | 0 | return Err(RedisError::from(( |
371 | 0 | ErrorKind::Parse, |
372 | 0 | "Expected at least 2 elements", |
373 | 0 | ))); |
374 | 29 | } |
375 | 29 | let mut output = Self::default(); |
376 | 29 | let mut value = value.into_iter(); |
377 | 29 | match value.next().unwrap() { |
378 | 27 | Value::Array(d) => resp2_data_parse(&mut output, &d)?0 , |
379 | 2 | Value::Map(d) => resp3_data_parse(&mut output, &d)?0 , |
380 | 0 | other => { |
381 | 0 | error!( |
382 | | ?other, |
383 | | "Bad data in ft.aggregate, expected array for results" |
384 | | ); |
385 | 0 | return Err(RedisError::from(( |
386 | 0 | ErrorKind::Parse, |
387 | 0 | "Non map item", |
388 | 0 | format!("{other:?}"), |
389 | 0 | ))); |
390 | | } |
391 | | } |
392 | 29 | let Value::Int(cursor) = value.next().unwrap() else { |
393 | 0 | return Err(RedisError::from(( |
394 | 0 | ErrorKind::Parse, |
395 | 0 | "Expected integer as last element", |
396 | 0 | ))); |
397 | | }; |
398 | 29 | output.cursor = cursor as u64; |
399 | 29 | Ok(output) |
400 | 29 | } |
401 | | } |