diff --git a/crates/warp_completer/src/signatures/legacy/miss_cache.rs b/crates/warp_completer/src/signatures/legacy/miss_cache.rs new file mode 100644 index 00000000000..f6cf419bf75 --- /dev/null +++ b/crates/warp_completer/src/signatures/legacy/miss_cache.rs @@ -0,0 +1,79 @@ +use std::collections::{HashSet, VecDeque}; +use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; + +const MAX_CACHED_MISSES: usize = 256; + +/// A bounded, `RwLock`-guarded FIFO set of (lowercased) command names that recently failed to +/// resolve to a signature. FIFO rather than LRU because a lookup here is a pure read that never +/// needs to reorder anything, so the write lock is only needed for a genuinely new miss. +pub(super) struct MissCache { + capacity: usize, + entries: RwLock, +} + +#[derive(Default)] +struct MissCacheEntries { + /// Insertion order, oldest first, used to find the next entry to evict once at capacity. + order: VecDeque, + /// The actual set of currently-remembered misses, for O(1) membership checks in `contains`. + set: HashSet, +} + +impl Default for MissCache { + fn default() -> Self { + Self::new(MAX_CACHED_MISSES) + } +} + +impl MissCache { + fn new(capacity: usize) -> Self { + Self { + capacity, + entries: RwLock::default(), + } + } + + /// Returns `true` if `command` was recently recorded as a miss. A pure read: does not + /// affect eviction order. + pub(super) fn contains(&self, command: &str) -> bool { + self.read().set.contains(command) + } + + /// Returns the number of misses currently recorded, for tests. + #[cfg(test)] + pub(super) fn len(&self) -> usize { + self.read().set.len() + } + + /// Records `command` as a miss, evicting the oldest-recorded miss first if already at + /// capacity. + pub(super) fn insert(&self, command: String) { + let mut entries = self.write(); + if entries.set.contains(&command) { + return; + } + if entries.order.len() >= self.capacity + && let Some(oldest) = entries.order.pop_front() + { + entries.set.remove(&oldest); + } + entries.order.push_back(command.clone()); + entries.set.insert(command); + } + + fn read(&self) -> RwLockReadGuard<'_, MissCacheEntries> { + self.entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn write(&self) -> RwLockWriteGuard<'_, MissCacheEntries> { + self.entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +#[cfg(test)] +#[path = "miss_cache_tests.rs"] +mod tests; diff --git a/crates/warp_completer/src/signatures/legacy/miss_cache_tests.rs b/crates/warp_completer/src/signatures/legacy/miss_cache_tests.rs new file mode 100644 index 00000000000..d04b0e8a580 --- /dev/null +++ b/crates/warp_completer/src/signatures/legacy/miss_cache_tests.rs @@ -0,0 +1,55 @@ +use super::MissCache; + +#[test] +fn test_contains_reflects_recorded_misses() { + let cache = MissCache::new(3); + + assert!(!cache.contains("a")); + cache.insert("a".to_string()); + assert!(cache.contains("a")); +} + +#[test] +fn test_inserting_an_existing_entry_is_a_no_op() { + let cache = MissCache::new(3); + + cache.insert("a".to_string()); + cache.insert("a".to_string()); + assert_eq!(cache.len(), 1); +} + +#[test] +fn test_stops_growing_at_capacity() { + let capacity = 4; + let cache = MissCache::new(capacity); + + for i in 0..capacity * 3 { + cache.insert(format!("miss-{i}")); + assert!(cache.len() <= capacity); + } + assert_eq!(cache.len(), capacity); +} + +#[test] +fn test_evicts_in_fifo_order_regardless_of_lookups() { + let cache = MissCache::new(3); + cache.insert("a".to_string()); + cache.insert("b".to_string()); + cache.insert("c".to_string()); + + for _ in 0..5 { + assert!(cache.contains("a")); + } + + cache.insert("d".to_string()); + assert!( + !cache.contains("a"), + "the oldest-inserted entry should have been evicted regardless of being looked up again" + ); + assert!( + cache.contains("b"), + "a newer entry should not have been evicted" + ); + assert!(cache.contains("c")); + assert!(cache.contains("d")); +} diff --git a/crates/warp_completer/src/signatures/legacy/mod.rs b/crates/warp_completer/src/signatures/legacy/mod.rs index 4096ee9654e..ae8dd8639e7 100644 --- a/crates/warp_completer/src/signatures/legacy/mod.rs +++ b/crates/warp_completer/src/signatures/legacy/mod.rs @@ -2,6 +2,7 @@ use std::sync::{Arc, OnceLock}; use warp_core::channel::Channel; +mod miss_cache; pub mod registry; pub use registry::CommandRegistry; diff --git a/crates/warp_completer/src/signatures/legacy/registry.rs b/crates/warp_completer/src/signatures/legacy/registry.rs index 2ef0036001d..59279f14d78 100644 --- a/crates/warp_completer/src/signatures/legacy/registry.rs +++ b/crates/warp_completer/src/signatures/legacy/registry.rs @@ -4,6 +4,7 @@ use itertools::Itertools; use memo_map::MemoMap; use warp_command_signatures::{Argument, DynamicCompletionData, IsArgumentOptional, Signature}; +use super::miss_cache::MissCache; use crate::completer::{CommandExitStatus, CompletionContext, TopLevelCommandCaseSensitivity}; use crate::parsers::SignatureAtTokenIndex; @@ -18,6 +19,8 @@ pub enum SignatureResult<'a> { type SignatureLookupFn = dyn 'static + Send + Sync + Fn(&str) -> Option; +const MAX_CACHEABLE_COMMAND_LEN: usize = 255; + /// A simple structure to cache parsed command signatures. These are stored as /// JSON, so this makes it easy for us to lazily load and parse the JSON when /// a command signature is needed, and only need to do that parsing work once @@ -27,12 +30,13 @@ struct SignatureCache { /// for it. Should return None if there is no signature available for the /// given command. lookup_fn: Box, - /// A map from command name to the signature for the command, if any. The - /// use of [`MemoMap`] here allows us to safely return references to the - /// contained signatures (as the map internally is an append-only - /// structure). This stores an `Option` in order to also store - /// our knowledge of commands for which we do _not_ have a signature. - signatures: MemoMap>, + /// A map from (lowercased) command name to its signature. The use of [`MemoMap`] here allows + /// us to safely return references to the contained signatures (as the map internally is an + /// append-only structure). + signatures: MemoMap, + /// A bounded set of (lowercased) command names that recently failed to resolve to a + /// signature. + misses: MissCache, } impl SignatureCache { @@ -40,6 +44,7 @@ impl SignatureCache { Self { lookup_fn, signatures: Default::default(), + misses: MissCache::default(), } } @@ -49,19 +54,44 @@ impl SignatureCache { } else { command }; + + if command.len() > MAX_CACHEABLE_COMMAND_LEN { + // No known command/subcommand name comes anywhere close to this length, so a token + // this long can never resolve to anything. Return before the lowercase allocation + // and before touching either cache -- in particular, an oversized token must never + // be admitted into `misses`, or the leak this cap exists to prevent would just move + // to the negative cache. + return None; + } + let command = command.to_lowercase(); - self.signatures - .get_or_insert(&command, || (self.lookup_fn)(&command)) - .as_ref() + + if let Some(signature) = self.signatures.get(command.as_str()) { + return Some(signature); + } + + if self.misses.contains(command.as_str()) { + return None; + } + + match (self.lookup_fn)(&command) { + Some(signature) => Some( + self.signatures + .get_or_insert(command.as_str(), || signature), + ), + None => { + self.misses.insert(command); + None + } + } } /// Inserts the given `Signature` into the underlying map, keyed by `Signature::name`. /// - /// If there is already a cached value for the given `Signature::name`, this is a no-op (even - /// if the cached value is `None`). + /// If there is already a cached value for the given `Signature::name`, this is a no-op. fn insert(&self, signature: Signature) { self.signatures - .insert(signature.name.to_lowercase(), Some(signature)); + .insert(signature.name.to_lowercase(), signature); } } @@ -126,12 +156,13 @@ impl CommandRegistry { pub fn registered_commands(&self) -> impl Iterator { // Note we need to collect the keys because MemoMap uses a mutex under the hood to control // access to the underlying signature data. This means the mutex is locked as long as the - // iterator returned from `keys()` lives, which means we need to collect keys into a vec - // and return an owned iterator. + // iterator returned from `iter()` lives, which means we need to collect keys into a vec + // and return an owned iterator. Every entry in `signatures` corresponds to a real + // signature (see its doc comment), so no filtering is needed here. self.signatures .signatures .iter() - .filter_map(|(key, signature)| signature.as_ref().map(|_| key.as_str())) + .map(|(key, _)| key.as_str()) .collect::>() .into_iter() } @@ -411,11 +442,13 @@ impl CommandRegistry { self.signatures.get(name) } - /// Registers the given `Signature`. + /// Registers the given `Signature`, making it resolvable via `signature()` and friends. /// - /// Note the underlying map caches the lookup result for a given signature (regardless of - /// whether or not it is `Some` or `None`), which means that if there is already a cached - /// `None` value for the command corresponding to this signature, this is a no-op. + /// `get` always checks the positive cache before the negative one (see + /// `SignatureCache::misses`), so this takes effect immediately even if a lookup for this + /// name previously missed and is currently sitting in the negative cache -- there's no + /// stale `None` result to invalidate. If a signature is already registered for this name, + /// this is a no-op. pub fn register_signature(&self, signature: Signature) { self.signatures.insert(signature); } diff --git a/crates/warp_completer/src/signatures/legacy/registry_tests.rs b/crates/warp_completer/src/signatures/legacy/registry_tests.rs index 2f5668a4f7a..a494d95be82 100644 --- a/crates/warp_completer/src/signatures/legacy/registry_tests.rs +++ b/crates/warp_completer/src/signatures/legacy/registry_tests.rs @@ -1,8 +1,65 @@ +use warp_command_signatures::{Priority, Signature}; +use warp_core::channel::Channel; + use crate::completer::testing::FakeCompletionContext; use crate::completer::{CompletionContext, TopLevelCommandCaseSensitivity}; -use crate::signatures::registry::SignatureResult; +use crate::signatures::registry::{MAX_CACHEABLE_COMMAND_LEN, SignatureResult}; use crate::signatures::testing::{create_test_command_registry, test_signature}; +/// A minimal signature with the given `name`, for exercising `SignatureCache` boundary +/// conditions that don't care about arguments, subcommands, or options. +fn signature_with_name(name: &str) -> Signature { + Signature { + name: name.to_string(), + alias_generator: None, + description: None, + arguments: None, + subcommands: None, + options: None, + priority: Priority::default(), + parser_directives: Default::default(), + } +} + +/// Recursively finds the longest `Signature::name` in `signature` or any of its (possibly +/// nested) subcommands, updating `longest` in place if a longer one is found. +fn track_longest_name(signature: &Signature, longest: &mut (usize, String)) { + if signature.name.len() > longest.0 { + *longest = (signature.name.len(), signature.name.clone()); + } + for subcommand in signature.subcommands() { + track_longest_name(subcommand, longest); + } +} + +#[test] +fn test_all_known_signature_names_are_within_the_length_cap() { + let mut longest = (0, String::new()); + + for signature in warp_command_signatures::commands() { + track_longest_name(&signature, &mut longest); + } + + for channel in [Channel::Stable, Channel::Preview, Channel::Dev] { + let mut clap_cmd = ::command(); + let signature = crate::signatures::clap::signature_from_clap_command( + &mut clap_cmd, + channel.cli_command_name(), + ); + track_longest_name(&signature, &mut longest); + } + + let (max_len, longest_name) = longest; + assert!( + max_len <= MAX_CACHEABLE_COMMAND_LEN, + "found a command/subcommand name of length {max_len} ({longest_name:?}), longer than \ + MAX_CACHEABLE_COMMAND_LEN ({MAX_CACHEABLE_COMMAND_LEN}); SignatureCache::get's \ + oversized-token fast path assumes no such name exists and would make this one \ + unresolvable -- raise MAX_CACHEABLE_COMMAND_LEN or add back a fallback lookup path for \ + oversized tokens" + ); +} + #[test] fn test_find_command_from_a_top_level_signature() { let bundle = warp_command_signatures::signature_by_name("bundle") @@ -204,6 +261,130 @@ fn test_alias_expansion_path_skips_flag_with_value_before_subcommand() { assert_eq!(found_signature.token_index, 3); } +#[test] +fn test_oversized_command_is_not_cached_and_resolves_to_none() { + let registry = create_test_command_registry([test_signature()]); + let positive_len_before = registry.signatures.signatures.len(); + let negative_len_before = registry.signatures.misses.len(); + + let oversized_command = "a".repeat(MAX_CACHEABLE_COMMAND_LEN + 1); + assert_eq!(registry.signature(&oversized_command), None); + assert_eq!( + registry.signatures.signatures.len(), + positive_len_before, + "looking up an oversized command should not add an entry to the positive cache" + ); + assert_eq!( + registry.signatures.misses.len(), + negative_len_before, + "looking up an oversized command should not add an entry to the negative cache either -- \ + that would just move the leak this cap exists to fix into the negative cache" + ); +} + +#[test] +fn test_misses_are_never_cached_in_the_positive_cache() { + let registry = create_test_command_registry([test_signature()]); + let len_before = registry.signatures.signatures.len(); + + assert_eq!(registry.signature("not-a-real-command"), None); + assert_eq!(registry.signature("not-a-real-command"), None); + assert_eq!(registry.signatures.signatures.len(), len_before); +} + +#[test] +fn test_oversized_later_token_does_not_bypass_the_length_guard() { + let sudo = warp_command_signatures::signature_by_name("sudo") + .expect("global command signatures should include 'sudo'"); + let registry = create_test_command_registry([sudo]); + let positive_len_before = registry.signatures.signatures.len(); + let negative_len_before = registry.signatures.misses.len(); + + let oversized_token = "a".repeat(MAX_CACHEABLE_COMMAND_LEN + 1); + let tokens = ["sudo", oversized_token.as_str()]; + + let found_signature = registry + .signature_from_tokens( + &tokens, + false, + TopLevelCommandCaseSensitivity::CaseSensitive, + ) + .expect("sudo signature from tokens should exist"); + assert_eq!( + found_signature.signature.name(), + "sudo", + "an oversized later token shouldn't resolve to a replacement signature" + ); + assert_eq!( + registry.signatures.signatures.len(), + positive_len_before, + "looking up an oversized later token should not add an entry to the positive cache" + ); + assert_eq!( + registry.signatures.misses.len(), + negative_len_before, + "looking up an oversized later token should not add an entry to the negative cache" + ); + + let sudo = warp_command_signatures::signature_by_name("sudo") + .expect("global command signatures should include 'sudo'"); + let ctx = + FakeCompletionContext::new(create_test_command_registry([sudo])).with_case_sensitivity(); + let result = warpui_core::r#async::block_on( + ctx.command_registry() + .signature_with_alias_expansion(&tokens, false, &ctx), + ); + let SignatureResult::Success(found_signature) = result else { + panic!("expected SignatureResult::Success"); + }; + assert_eq!(found_signature.signature.name(), "sudo"); +} + +#[test] +fn test_ordinary_commands_still_resolve_and_are_cached() { + let registry = create_test_command_registry([test_signature()]); + + let found = registry.signature("TEST"); + assert_eq!(found.map(|s| s.name.as_str()), Some("test")); + + let len_after_first_lookup = registry.signatures.signatures.len(); + assert_eq!(registry.signature("test"), found); + assert_eq!(registry.signatures.signatures.len(), len_after_first_lookup); +} + +#[test] +fn test_registered_signature_longer_than_the_cap_is_unresolvable() { + let long_name = "a".repeat(MAX_CACHEABLE_COMMAND_LEN + 1); + let registry = create_test_command_registry([signature_with_name(&long_name)]); + + assert_eq!(registry.signature(&long_name), None); +} + +#[cfg(windows)] +#[test] +fn test_exe_suffix_is_trimmed_before_the_length_check() { + let max_length_name = "a".repeat(MAX_CACHEABLE_COMMAND_LEN); + let registry = create_test_command_registry([signature_with_name(&max_length_name)]); + + let token = format!("{max_length_name}.exe"); + assert_eq!( + registry.signature(&token).map(|s| s.name.as_str()), + Some(max_length_name.as_str()) + ); +} + +#[test] +fn test_registered_commands_unaffected_by_oversized_lookups() { + let registry = create_test_command_registry([test_signature()]); + + let oversized_command = "a".repeat(MAX_CACHEABLE_COMMAND_LEN + 1); + assert_eq!(registry.signature(&oversized_command), None); + assert_eq!(registry.signature("not-a-real-command"), None); + + let registered = registry.registered_commands().collect::>(); + assert_eq!(registered, vec!["test"]); +} + #[test] fn test_alias_expansion_path_skips_multiple_flags_before_subcommand() { // Exercises signature_with_alias_expansion with multiple flags (valued and