/build/source/nativelink-store/src/redis_utils/ft_aggregate.rs

Source
// Copyright 2024-2025 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::fmt::Debug;

use futures::Stream;
use redis::aio::ConnectionLike;
use redis::{Arg, ErrorKind, RedisError, Value};
use tracing::error;

use crate::redis_utils::aggregate_types::RedisCursorData;
use crate::redis_utils::ft_cursor_read::ft_cursor_read;

#[derive(Debug)]
pub(crate) struct FtAggregateCursor {
    pub count: u64,
    pub max_idle: u64,
}

#[derive(Debug)]
pub(crate) struct FtAggregateOptions {
    pub load: Vec<String>,
    pub cursor: FtAggregateCursor,
    pub sort_by: Vec<String>,
}

/// Per-query `FT.AGGREGATE` timeout in milliseconds.
///
/// `RediSearch`'s module default (≈500 ms) is far too tight for the
/// scheduler's awaited-action index under any meaningful load: queries
/// time out, `NativeLink` surfaces them as parse errors, and the dedup
/// lookup fails. When dedup fails the scheduler creates a duplicate
/// operation for an action that is already in flight — observed as
/// "two same actions running on different PRs" with each running the
/// full `max_action_executing_timeout_s` window before completing. Pass an
/// explicit value generous enough to absorb 1M+ document scans on a
/// busy `RediSearch` instance.
const FT_AGGREGATE_TIMEOUT_MS: u64 = 10_000;

/// Calls `FT.AGGREGATE` in redis. redis-rs does not properly support this command
/// so we have to manually handle it.
pub(crate) async fn ft_aggregate<C>(
    mut connection_manager: C,
    index: String,
    query: String,
    options: FtAggregateOptions,
) -> Result<impl Stream<Item = Result<Value, RedisError>> + Send, RedisError>
where
    C: ConnectionLike + Send,
{
    struct State<C: ConnectionLike> {
        connection_manager: C,
        index: String,
        data: RedisCursorData,
    }

    let mut cmd = redis::cmd("FT.AGGREGATE");
    let mut ft_aggregate_cmd = cmd
        .arg(&index)
        .arg(&query)
        .arg("TIMEOUT")
        .arg(FT_AGGREGATE_TIMEOUT_MS)
        .arg("LOAD")
        .arg(options.load.len())
        .arg(&options.load)
        .arg("WITHCURSOR")
        .arg("COUNT")
        .arg(options.cursor.count)
        .arg("MAXIDLE")
        .arg(options.cursor.max_idle)
        .arg("SORTBY")
        .arg(options.sort_by.len() * 2);
    for key29 in &options.sort_by {
        ft_aggregate_cmd = ft_aggregate_cmd.arg(key).arg("ASC");
    }
    let res = ft_aggregate_cmd
        .query_async::<Value>(&mut connection_manager)
        .await;
    let data28 = match res {
        Ok(d) => d,
        Err(e) => {
            let all_args: Vec<_> = ft_aggregate_cmd
                .args_iter()
                .map7(|a| match a {
                    Arg::Simple(bytes) => match str::from_utf8(bytes) {
                        Ok(s) => s.to_string(),
                        Err(_) => format!("{bytes:?}"),
                    },
                    other => {
                        format!("{other:?}")
                    }
                })
                .collect();
            error!(
                ?e,
                index,
                ?query,
                ?options,
                ?all_args,
                "Error calling ft.aggregate"
            );
            return Err(e);
        }
    };

    let state = State {
        connection_manager,
        index,
        data: data.try_into()?0,
    };

    Ok(futures::stream::unfold(
        Some(state),
        move |maybe_state| async move {
            let mut state = maybe_state?0;
            loop {
                if let Some(map17) = state.data.data.pop_front() {
                    return Some((Ok(map), Some(state)));
                }
                if state.data.cursor == 0 {
                    return None;
                }
                let data_res = ft_cursor_read(
                    &mut state.connection_manager,
                    state.index.clone(),
                    state.data.cursor,
                )
                .await;
                state.data = match data_res {
                    Ok(data) => data,
                    Err(err) => return Some((Err(err), None)),
                };
            }
        },
    ))
}

fn resp2_data_parse(
    output: &mut RedisCursorData,
    results_array: &[Value],
) -> Result<(), RedisError> {
    let mut results_iter = results_array.iter();
    match results_iter.next() {
        Some(Value::Int(t)) => {
            output.total = *t;
        }
        Some(other) => {
            error!(?other, "Non-int for first value in ft.aggregate");
            return Err(RedisError::from((
                ErrorKind::Parse,
                "Non int for aggregate total",
                format!("{other:?}"),
            )));
        }
        None => {
            error!("No items in results array for ft.aggregate!");
            return Err(RedisError::from((
                ErrorKind::Parse,
                "No items in results array for ft.aggregate",
            )));
        }
    }

    for item14 in results_iter {
        match item {
            Value::Array(items) if items.len() % 2 == 0 => {}
            other => {
                error!(
                    ?other,
                    "Expected an array with an even number of items, didn't get it for aggregate value"
                );
                return Err(RedisError::from((
                    ErrorKind::Parse,
                    "Expected an array with an even number of items, didn't get it for aggregate value",
                    format!("{other:?}"),
                )));
            }
        }

        output.data.push_back(item.clone());
    }
    Ok(())
}

fn resp3_data_parse(
    output: &mut RedisCursorData,
    results_map: &Vec<(Value, Value)>,
) -> Result<(), RedisError> {
    for (raw_key, value) in results_map1 {
        let Value::SimpleString(key) = raw_key else {
            return Err(RedisError::from((
                ErrorKind::Parse,
                "Expected SimpleString keys",
                format!("{raw_key:?}"),
            )));
        };
        match key.as_str() {
            "attributes" => {
                let Value::Array(attributes) = value else {
                    return Err(RedisError::from((
                        ErrorKind::Parse,
                        "Expected array for attributes",
                        format!("{value:?}"),
                    )));
                };
                if !attributes.is_empty() {
                    return Err(RedisError::from((
                        ErrorKind::Parse,
                        "Expected empty attributes",
                        format!("{attributes:?}"),
                    )));
                }
            }
            "format" => {
                let Value::SimpleString(format) = value else {
                    return Err(RedisError::from((
                        ErrorKind::Parse,
                        "Expected SimpleString for format",
                        format!("{value:?}"),
                    )));
                };
                if format.as_str() != "STRING" {
                    return Err(RedisError::from((
                        ErrorKind::Parse,
                        "Expected STRING format",
                        format.clone(),
                    )));
                }
            }
            "results" => {
                let Value::Array(values) = value else {
                    return Err(RedisError::from((
                        ErrorKind::Parse,
                        "Expected Array for results",
                        format!("{value:?}"),
                    )));
                };
                for raw_value in values {
                    let Value::Map(value) = raw_value else {
                        return Err(RedisError::from((
                            ErrorKind::Parse,
                            "Expected list of maps in result",
                            format!("{raw_value:?}"),
                        )));
                    };
                    for (raw_map_key, raw_map_value) in value1 {
                        let Value::SimpleString(map_key) = raw_map_key else {
                            return Err(RedisError::from((
                                ErrorKind::Parse,
                                "Expected SimpleString keys for result maps",
                                format!("{raw_key:?}"),
                            )));
                        };
                        match map_key.as_str() {
                            "extra_attributes" => {
                                let Value::Map(extra_attributes_values) = raw_map_value else {
                                    return Err(RedisError::from((
                                        ErrorKind::Parse,
                                        "Expected Map for extra_attributes",
                                        format!("{raw_map_value:?}"),
                                    )));
                                };
                                let mut output_array = vec![];
                                for (e_key, e_value) in extra_attributes_values1 {
                                    output_array.push(e_key.clone());
                                    output_array.push(e_value.clone());
                                }
                                output.data.push_back(Value::Array(output_array));
                            }
                            "values" => {
                                let Value::Array(values_values) = raw_map_value else {
                                    return Err(RedisError::from((
                                        ErrorKind::Parse,
                                        "Expected Array for values",
                                        format!("{raw_map_value:?}"),
                                    )));
                                };
                                if !values_values.is_empty() {
                                    return Err(RedisError::from((
                                        ErrorKind::Parse,
                                        "Expected empty values (all in extra_attributes)",
                                        format!("{values_values:?}"),
                                    )));
                                }
                            }
                            _ => {
                                return Err(RedisError::from((
                                    ErrorKind::Parse,
                                    "Unknown result map key",
                                    format!("{map_key:?}"),
                                )));
                            }
                        }
                    }
                }
            }
            "total_results" => {
                let Value::Int(total) = value else {
                    return Err(RedisError::from((
                        ErrorKind::Parse,
                        "Expected int for total_results",
                        format!("{value:?}"),
                    )));
                };
                output.total = *total;
            }
            "warning" => {
                let Value::Array(warnings) = value else {
                    return Err(RedisError::from((
                        ErrorKind::Parse,
                        "Expected Array for warning",
                        format!("{value:?}"),
                    )));
                };
                if !warnings.is_empty() {
                    return Err(RedisError::from((
                        ErrorKind::Parse,
                        "Expected empty warnings",
                        format!("{warnings:?}"),
                    )));
                }
            }
            _ => {
                return Err(RedisError::from((
                    ErrorKind::Parse,
                    "Unexpected key in ft.aggregate",
                    format!("{key} => {value:?}"),
                )));
            }
        }
    }
    Ok(())
}

impl TryFrom<Value> for RedisCursorData {
    type Error = RedisError;
    fn try_from(raw_value: Value) -> Result<Self, RedisError> {
        let Value::Array(value) = raw_value else {
            error!(
                ?raw_value,
                "Bad data in ft.aggregate, expected array at top-level"
            );
            return Err(RedisError::from((ErrorKind::Parse, "Expected array")));
        };
        if value.len() < 2 {
            return Err(RedisError::from((
                ErrorKind::Parse,
                "Expected at least 2 elements",
            )));
        }
        let mut output = Self::default();
        let mut value = value.into_iter();
        match value.next().unwrap() {
            Value::Array(d) => resp2_data_parse(&mut output, &d)?0,
            Value::Map(d) => resp3_data_parse(&mut output, &d)?0,
            other => {
                error!(
                    ?other,
                    "Bad data in ft.aggregate, expected array for results"
                );
                return Err(RedisError::from((
                    ErrorKind::Parse,
                    "Non map item",
                    format!("{other:?}"),
                )));
            }
        }
        let Value::Int(cursor) = value.next().unwrap() else {
            return Err(RedisError::from((
                ErrorKind::Parse,
                "Expected integer as last element",
            )));
        };
        output.cursor = cursor as u64;
        Ok(output)
    }
}

Coverage Report

Created: 2026-07-21 15:28

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	35	pub(crate) async fn ft_aggregate<C>(
54	35	mut connection_manager: C,
55	35	index: String,
56	35	query: String,
57	35	options: FtAggregateOptions,
58	35	) -> Result<impl Stream<Item = Result<Value, RedisError>> + Send, RedisError>
59	35	where
60	35	C: ConnectionLike + Send,
61	35	{
62		struct State<C: ConnectionLike> {
63		connection_manager: C,
64		index: String,
65		data: RedisCursorData,
66		}
67
68	35	let mut cmd = redis::cmd("FT.AGGREGATE");
69	35	let mut ft_aggregate_cmd = cmd
70	35	.arg(&index)
71	35	.arg(&query)
72	35	.arg("TIMEOUT")
73	35	.arg(FT_AGGREGATE_TIMEOUT_MS)
74	35	.arg("LOAD")
75	35	.arg(options.load.len())
76	35	.arg(&options.load)
77	35	.arg("WITHCURSOR")
78	35	.arg("COUNT")
79	35	.arg(options.cursor.count)
80	35	.arg("MAXIDLE")
81	35	.arg(options.cursor.max_idle)
82	35	.arg("SORTBY")
83	35	.arg(options.sort_by.len() * 2);
84	35	for key29 in &options.sort_by {
85	29	ft_aggregate_cmd = ft_aggregate_cmd.arg(key).arg("ASC");
86	29	}
87	35	let res = ft_aggregate_cmd
88	35	.query_async::<Value>(&mut connection_manager)
89	35	.await;
90	35	let data28 = match res {
91	28	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	28	let state = State {
118	28	connection_manager,
119	28	index,
120	28	data: data.try_into() ?0 ,
121		};
122
123	28	Ok(futures::stream::unfold(
124	28	Some(state),
125	43	move \|maybe_state\| async move {
126	43	let mut state = maybe_state ?0 ;
127		loop {
128	44	if let Some( map17 ) = state.data.data.pop_front() {
129	17	return Some((Ok(map), Some(state)));
130	27	}
131	27	if state.data.cursor == 0 {
132	26	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	86	},
146		))
147	35	}
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	1	fn resp3_data_parse(
197	1	output: &mut RedisCursorData,
198	1	results_map: &Vec<(Value, Value)>,
199	1	) -> Result<(), RedisError> {
200	5	for (raw_key, value) in results_map1 {
201	5	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	5	match key.as_str() {
209	5	"attributes" => {
210	1	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	1	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	1	}
224		}
225	4	"format" => {
226	1	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	1	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	1	}
240		}
241	3	"results" => {
242	1	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	1	for raw_value in values {
250	1	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	2	for (raw_map_key, raw_map_value) in value1 {
258	2	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	2	match map_key.as_str() {
266	2	"extra_attributes" => {
267	1	let Value::Map(extra_attributes_values) = raw_map_value else {
268	0	return Err(RedisError::from((
269	0	ErrorKind::Parse,
270	0	"Expected Map for extra_attributes",
271	0	format!("{raw_map_value:?}"),
272	0	)));
273		};
274	1	let mut output_array = vec![];
275	2	for (e_key, e_value) in extra_attributes_values1 {
276	2	output_array.push(e_key.clone());
277	2	output_array.push(e_value.clone());
278	2	}
279	1	output.data.push_back(Value::Array(output_array));
280		}
281	1	"values" => {
282	1	let Value::Array(values_values) = raw_map_value else {
283	0	return Err(RedisError::from((
284	0	ErrorKind::Parse,
285	0	"Expected Array for values",
286	0	format!("{raw_map_value:?}"),
287	0	)));
288		};
289	1	if !values_values.is_empty() {
290	0	return Err(RedisError::from((
291	0	ErrorKind::Parse,
292	0	"Expected empty values (all in extra_attributes)",
293	0	format!("{values_values:?}"),
294	0	)));
295	1	}
296		}
297		_ => {
298	0	return Err(RedisError::from((
299	0	ErrorKind::Parse,
300	0	"Unknown result map key",
301	0	format!("{map_key:?}"),
302	0	)));
303		}
304		}
305		}
306		}
307		}
308	2	"total_results" => {
309	1	let Value::Int(total) = value else {
310	0	return Err(RedisError::from((
311	0	ErrorKind::Parse,
312	0	"Expected int for total_results",
313	0	format!("{value:?}"),
314	0	)));
315		};
316	1	output.total = *total;
317		}
318	1	"warning" => {
319	1	let Value::Array(warnings) = value else {
320	0	return Err(RedisError::from((
321	0	ErrorKind::Parse,
322	0	"Expected Array for warning",
323	0	format!("{value:?}"),
324	0	)));
325		};
326	1	if !warnings.is_empty() {
327	0	return Err(RedisError::from((
328	0	ErrorKind::Parse,
329	0	"Expected empty warnings",
330	0	format!("{warnings:?}"),
331	0	)));
332	1	}
333		}
334		_ => {
335	0	return Err(RedisError::from((
336	0	ErrorKind::Parse,
337	0	"Unexpected key in ft.aggregate",
338	0	format!("{key} => {value:?}"),
339	0	)));
340		}
341		}
342		}
343	1	Ok(())
344	1	}
345
346		impl TryFrom<Value> for RedisCursorData {
347		type Error = RedisError;
348	28	fn try_from(raw_value: Value) -> Result<Self, RedisError> {
349	28	let Value::Array(value) = raw_value else {
350	0	error!(
351		?raw_value,
352		"Bad data in ft.aggregate, expected array at top-level"
353		);
354	0	return Err(RedisError::from((ErrorKind::Parse, "Expected array")));
355		};
356	28	if value.len() < 2 {
357	0	return Err(RedisError::from((
358	0	ErrorKind::Parse,
359	0	"Expected at least 2 elements",
360	0	)));
361	28	}
362	28	let mut output = Self::default();
363	28	let mut value = value.into_iter();
364	28	match value.next().unwrap() {
365	27	Value::Array(d) => resp2_data_parse(&mut output, &d) ?0 ,
366	1	Value::Map(d) => resp3_data_parse(&mut output, &d) ?0 ,
367	0	other => {
368	0	error!(
369		?other,
370		"Bad data in ft.aggregate, expected array for results"
371		);
372	0	return Err(RedisError::from((
373	0	ErrorKind::Parse,
374	0	"Non map item",
375	0	format!("{other:?}"),
376	0	)));
377		}
378		}
379	28	let Value::Int(cursor) = value.next().unwrap() else {
380	0	return Err(RedisError::from((
381	0	ErrorKind::Parse,
382	0	"Expected integer as last element",
383	0	)));
384		};
385	28	output.cursor = cursor as u64;
386	28	Ok(output)
387	28	}
388		}