Skip to content

Commit a11cb87

Browse files
ceekay47facebook-github-bot
ceekay47
authored and
facebook-github-bot
committed
[presto] Refactor StaticSelector.match to avoid throwing exceptions
Summary: The `addVariableValues` method in `StaticSelector` previously iterated over all `variableNames` and caught `IllegalArgumentException` for each name not present as a named group in the pattern. This is inefficient and relies on exception handling for control flow. This change pre-computes the named groups for each regex pattern at construction time into an `ImmutableMap<Pattern, ImmutableSet<String>>`. The `addVariableValues` method now looks up the pre-computed groups for the given pattern and only iterates over matching group names, avoiding the need to catch exceptions entirely. Added unit tests covering named group extraction from userRegex, sourceRegex, multiple named groups, combined user+source named groups, patterns without named groups, and unresolved variable validation. Reviewed By: swapsmagic Differential Revision: D93437091
1 parent ee7c62a commit a11cb87

File tree

2 files changed

+111
-11
lines changed

2 files changed

+111
-11
lines changed

‎presto-resource-group-managers/src/main/java/com/facebook/presto/resourceGroups/StaticSelector.java‎

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.facebook.presto.spi.resourceGroups.SelectionCriteria;
1919
import com.google.common.annotations.VisibleForTesting;
2020
import com.google.common.collect.ImmutableList;
21+
import com.google.common.collect.ImmutableMap;
2122
import com.google.common.collect.ImmutableSet;
2223
import com.google.common.collect.Sets;
2324

@@ -55,6 +56,7 @@ public class StaticSelector
5556
private final Optional<Pattern> principalRegex;
5657
private final ResourceGroupIdTemplate group;
5758
private final Set<String> variableNames;
59+
private final ImmutableMap<Pattern, ImmutableSet<String>> patternNamedGroups;
5860

5961
public StaticSelector(
6062
Optional<Pattern> userRegex,
@@ -83,6 +85,11 @@ public StaticSelector(
8385
sourceRegex.ifPresent(s -> addNamedGroups(s, variableNames));
8486
this.variableNames = ImmutableSet.copyOf(variableNames);
8587

88+
ImmutableMap.Builder<Pattern, ImmutableSet<String>> patternGroupsBuilder = ImmutableMap.builder();
89+
userRegex.ifPresent(u -> patternGroupsBuilder.put(u, extractNamedGroups(u)));
90+
sourceRegex.ifPresent(s -> patternGroupsBuilder.put(s, extractNamedGroups(s)));
91+
this.patternNamedGroups = patternGroupsBuilder.build();
92+
8693
Set<String> unresolvedVariables = Sets.difference(group.getVariableNames(), variableNames);
8794
checkArgument(unresolvedVariables.isEmpty(), "unresolved variables %s in resource group ID '%s', available: %s\"", unresolvedVariables, group, variableNames);
8895
}
@@ -167,19 +174,28 @@ private static void addNamedGroups(Pattern pattern, HashSet<String> variables)
167174
}
168175
}
169176

177+
private static ImmutableSet<String> extractNamedGroups(Pattern pattern)
178+
{
179+
ImmutableSet.Builder<String> groups = ImmutableSet.builder();
180+
Matcher matcher = NAMED_GROUPS_PATTERN.matcher(pattern.pattern());
181+
while (matcher.find()) {
182+
groups.add(matcher.group(1));
183+
}
184+
return groups.build();
185+
}
186+
170187
private void addVariableValues(Pattern pattern, String candidate, Map<String, String> mapping)
171188
{
172-
for (String key : variableNames) {
173-
Matcher keyMatcher = pattern.matcher(candidate);
174-
if (keyMatcher.find()) {
175-
try {
176-
String value = keyMatcher.group(key);
177-
if (value != null) {
178-
mapping.put(key, value);
179-
}
180-
}
181-
catch (IllegalArgumentException ignored) {
182-
// there was no capturing group with the specified name
189+
Set<String> groups = patternNamedGroups.getOrDefault(pattern, ImmutableSet.of());
190+
if (groups.isEmpty()) {
191+
return;
192+
}
193+
Matcher keyMatcher = pattern.matcher(candidate);
194+
if (keyMatcher.find()) {
195+
for (String key : groups) {
196+
String value = keyMatcher.group(key);
197+
if (value != null) {
198+
mapping.put(key, value);
183199
}
184200
}
185201
}

‎presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/TestStaticSelector.java‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,90 @@ public void testSourcePrincipalRegex()
277277
assertEquals(selector.match(newSelectionCriteria("A.user", "a source b", "a " + principal + " b", EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
278278
}
279279

280+
@Test
281+
public void testUserRegexNamedGroup()
282+
{
283+
ResourceGroupId resourceGroupId = new ResourceGroupId(new ResourceGroupId("global"), "teamA");
284+
StaticSelector selector = new StaticSelector(
285+
Optional.of(Pattern.compile("(?<team>[a-zA-Z]+)_user")),
286+
Optional.empty(),
287+
Optional.empty(),
288+
Optional.empty(),
289+
Optional.empty(),
290+
Optional.empty(),
291+
Optional.empty(),
292+
Optional.empty(),
293+
new ResourceGroupIdTemplate("global.${team}"));
294+
assertEquals(
295+
selector.match(newSelectionCriteria("teamA_user", null, ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId),
296+
Optional.of(resourceGroupId));
297+
assertEquals(
298+
selector.match(newSelectionCriteria("noMatch", null, ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)),
299+
Optional.empty());
300+
}
301+
302+
@Test
303+
public void testSourceRegexNamedGroup()
304+
{
305+
ResourceGroupId resourceGroupId = new ResourceGroupId(new ResourceGroupId("global"), "pipeline");
306+
StaticSelector selector = new StaticSelector(
307+
Optional.empty(),
308+
Optional.of(Pattern.compile("(?<pipeline>[a-zA-Z]+)\\.query")),
309+
Optional.empty(),
310+
Optional.empty(),
311+
Optional.empty(),
312+
Optional.empty(),
313+
Optional.empty(),
314+
Optional.empty(),
315+
new ResourceGroupIdTemplate("global.${pipeline}"));
316+
assertEquals(
317+
selector.match(newSelectionCriteria("user", "pipeline.query", ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId),
318+
Optional.of(resourceGroupId));
319+
assertEquals(
320+
selector.match(newSelectionCriteria("user", "other", ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)),
321+
Optional.empty());
322+
}
323+
324+
@Test
325+
public void testMultipleNamedGroupsInSinglePattern()
326+
{
327+
ResourceGroupId resourceGroupId = new ResourceGroupId(
328+
new ResourceGroupId(new ResourceGroupId("global"), "teamA"), "role1");
329+
StaticSelector selector = new StaticSelector(
330+
Optional.of(Pattern.compile("(?<team>[a-zA-Z]+)_(?<role>[a-zA-Z0-9]+)")),
331+
Optional.empty(),
332+
Optional.empty(),
333+
Optional.empty(),
334+
Optional.empty(),
335+
Optional.empty(),
336+
Optional.empty(),
337+
Optional.empty(),
338+
new ResourceGroupIdTemplate("global.${team}.${role}"));
339+
assertEquals(
340+
selector.match(newSelectionCriteria("teamA_role1", null, ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId),
341+
Optional.of(resourceGroupId));
342+
}
343+
344+
@Test
345+
public void testNamedGroupsFromBothUserAndSourceRegex()
346+
{
347+
ResourceGroupId resourceGroupId = new ResourceGroupId(
348+
new ResourceGroupId(new ResourceGroupId("global"), "alice"), "etl");
349+
StaticSelector selector = new StaticSelector(
350+
Optional.of(Pattern.compile("(?<name>[a-zA-Z]+)")),
351+
Optional.of(Pattern.compile("(?<pipeline>[a-zA-Z]+)_job")),
352+
Optional.empty(),
353+
Optional.empty(),
354+
Optional.empty(),
355+
Optional.empty(),
356+
Optional.empty(),
357+
Optional.empty(),
358+
new ResourceGroupIdTemplate("global.${name}.${pipeline}"));
359+
assertEquals(
360+
selector.match(newSelectionCriteria("alice", "etl_job", ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId),
361+
Optional.of(resourceGroupId));
362+
}
363+
280364
private SelectionCriteria newSelectionCriteria(String user, String source, Set<String> tags, ResourceEstimates resourceEstimates)
281365
{
282366
return new SelectionCriteria(true, user, Optional.ofNullable(source), tags, resourceEstimates, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());

0 commit comments

Comments
 (0)