Coverage Report

Created: 2026-08-13 05:32

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/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
}