Skip to content

Commit f395157

Browse files
alexjskicopybara-github
alexjski
authored and
copybara-github
committed
Allow running an extra spawn for local branch of dynamic execution.
Allow `DynamicExecutionModule` to specify an extra spawn to be ran in the local branch. Add support in `DynamicSpawnStrategy` for running the extra spawn when it is provided. PiperOrigin-RevId: 344826682
1 parent d155376 commit f395157

File tree

4 files changed

+327
-3
lines changed

4 files changed

+327
-3
lines changed

‎src/main/java/com/google/devtools/build/lib/dynamic/DynamicExecutionModule.java

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import java.util.HashMap;
3939
import java.util.List;
4040
import java.util.Map;
41+
import java.util.Optional;
4142
import java.util.concurrent.ExecutorService;
4243
import java.util.concurrent.Executors;
4344

@@ -140,7 +141,12 @@ final void registerSpawnStrategies(
140141
if (options.legacySpawnScheduler) {
141142
strategy = new LegacyDynamicSpawnStrategy(executorService, options, this::getExecutionPolicy);
142143
} else {
143-
strategy = new DynamicSpawnStrategy(executorService, options, this::getExecutionPolicy);
144+
strategy =
145+
new DynamicSpawnStrategy(
146+
executorService,
147+
options,
148+
this::getExecutionPolicy,
149+
this::getPostProcessingSpawnForLocalExecution);
144150
}
145151
registryBuilder.registerStrategy(strategy, "dynamic", "dynamic_worker");
146152

@@ -183,6 +189,18 @@ protected ExecutionPolicy getExecutionPolicy(Spawn spawn) {
183189
return ExecutionPolicy.ANYWHERE;
184190
}
185191

192+
/**
193+
* Returns a post processing {@link Spawn} if one needs to be executed after given {@link Spawn}
194+
* when running locally.
195+
*
196+
* <p>The intention of this is to allow post-processing of the original {@linkplain Spawn spawn}
197+
* when executing it locally. In particular, such spawn should never create outputs which are not
198+
* included in the generating action of the original one.
199+
*/
200+
protected Optional<Spawn> getPostProcessingSpawnForLocalExecution(Spawn spawn) {
201+
return Optional.empty();
202+
}
203+
186204
@Override
187205
public void afterCommand() {
188206
ExecutorUtil.interruptibleShutdown(executorService);

‎src/main/java/com/google/devtools/build/lib/dynamic/DynamicSpawnStrategy.java

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import com.google.devtools.build.lib.actions.SandboxedSpawnStrategy;
3434
import com.google.devtools.build.lib.actions.Spawn;
3535
import com.google.devtools.build.lib.actions.SpawnResult;
36+
import com.google.devtools.build.lib.actions.SpawnResult.Status;
3637
import com.google.devtools.build.lib.actions.SpawnStrategy;
3738
import com.google.devtools.build.lib.events.Event;
3839
import com.google.devtools.build.lib.exec.ExecutionPolicy;
@@ -42,6 +43,7 @@
4243
import com.google.devtools.build.lib.util.io.FileOutErr;
4344
import com.google.devtools.build.lib.vfs.Path;
4445
import java.io.IOException;
46+
import java.util.Optional;
4547
import java.util.concurrent.Callable;
4648
import java.util.concurrent.CancellationException;
4749
import java.util.concurrent.ExecutionException;
@@ -88,6 +90,8 @@ public class DynamicSpawnStrategy implements SpawnStrategy {
8890
*/
8991
private final AtomicBoolean delayLocalExecution = new AtomicBoolean(false);
9092

93+
private final Function<Spawn, Optional<Spawn>> getExtraSpawnForLocalExecution;
94+
9195
/**
9296
* Constructs a {@code DynamicSpawnStrategy}.
9397
*
@@ -96,10 +100,12 @@ public class DynamicSpawnStrategy implements SpawnStrategy {
96100
public DynamicSpawnStrategy(
97101
ExecutorService executorService,
98102
DynamicExecutionOptions options,
99-
Function<Spawn, ExecutionPolicy> getExecutionPolicy) {
103+
Function<Spawn, ExecutionPolicy> getExecutionPolicy,
104+
Function<Spawn, Optional<Spawn>> getPostProcessingSpawnForLocalExecution) {
100105
this.executorService = MoreExecutors.listeningDecorator(executorService);
101106
this.options = options;
102107
this.getExecutionPolicy = getExecutionPolicy;
108+
this.getExtraSpawnForLocalExecution = getPostProcessingSpawnForLocalExecution;
103109
}
104110

105111
/**
@@ -463,7 +469,34 @@ private static FileOutErr getSuffixedFileOutErr(FileOutErr fileOutErr, String su
463469
outDir.getChild(outBaseName + suffix), errDir.getChild(errBaseName + suffix));
464470
}
465471

466-
private static ImmutableList<SpawnResult> runLocally(
472+
private ImmutableList<SpawnResult> runLocally(
473+
Spawn spawn,
474+
ActionExecutionContext actionExecutionContext,
475+
@Nullable SandboxedSpawnStrategy.StopConcurrentSpawns stopConcurrentSpawns)
476+
throws ExecException, InterruptedException {
477+
ImmutableList<SpawnResult> spawnResult =
478+
runSpawnLocally(spawn, actionExecutionContext, stopConcurrentSpawns);
479+
if (spawnResult.stream().anyMatch(result -> result.status() != Status.SUCCESS)) {
480+
return spawnResult;
481+
}
482+
483+
Optional<Spawn> extraSpawn = getExtraSpawnForLocalExecution.apply(spawn);
484+
if (!extraSpawn.isPresent()) {
485+
return spawnResult;
486+
}
487+
488+
// The remote branch was already cancelled -- we are holding the output lock during the
489+
// execution of the extra spawn.
490+
ImmutableList<SpawnResult> extraSpawnResult =
491+
runSpawnLocally(extraSpawn.get(), actionExecutionContext, null);
492+
return ImmutableList.<SpawnResult>builderWithExpectedSize(
493+
spawnResult.size() + extraSpawnResult.size())
494+
.addAll(spawnResult)
495+
.addAll(extraSpawnResult)
496+
.build();
497+
}
498+
499+
private static ImmutableList<SpawnResult> runSpawnLocally(
467500
Spawn spawn,
468501
ActionExecutionContext actionExecutionContext,
469502
@Nullable SandboxedSpawnStrategy.StopConcurrentSpawns stopConcurrentSpawns)

‎src/test/java/com/google/devtools/build/lib/dynamic/BUILD

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,25 @@ filegroup(
1212
visibility = ["//src:__subpackages__"],
1313
)
1414

15+
java_test(
16+
name = "DynamicSpawnStrategyUnitTest",
17+
size = "small",
18+
srcs = ["DynamicSpawnStrategyUnitTest.java"],
19+
deps = [
20+
"//src/main/java/com/google/devtools/build/lib/actions",
21+
"//src/main/java/com/google/devtools/build/lib/dynamic",
22+
"//src/main/java/com/google/devtools/build/lib/exec:execution_policy",
23+
"//src/main/protobuf:failure_details_java_proto",
24+
"//src/test/java/com/google/devtools/build/lib/exec/util",
25+
"//src/test/java/com/google/devtools/build/lib/testutil",
26+
"//src/test/java/com/google/devtools/build/lib/testutil:TestUtils",
27+
"//third_party:guava",
28+
"//third_party:junit4",
29+
"//third_party:mockito",
30+
"//third_party:truth",
31+
],
32+
)
33+
1534
java_test(
1635
name = "DynamicSpawnStrategyTest",
1736
size = "small",

‎src/test/java/com/google/devtools/build/lib/dynamic/DynamicSpawnStrategyUnitTest.java

Lines changed: 254 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,254 @@
1+
// Copyright 2020 The Bazel Authors. All rights reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (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+
// http://www.apache.org/licenses/LICENSE-2.0
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+
package com.google.devtools.build.lib.dynamic;
15+
16+
import static com.google.common.base.Preconditions.checkState;
17+
import static com.google.common.truth.Truth.assertThat;
18+
import static java.util.concurrent.TimeUnit.MILLISECONDS;
19+
import static org.mockito.ArgumentMatchers.any;
20+
import static org.mockito.ArgumentMatchers.eq;
21+
import static org.mockito.ArgumentMatchers.isNotNull;
22+
import static org.mockito.ArgumentMatchers.isNull;
23+
import static org.mockito.Mockito.mock;
24+
import static org.mockito.Mockito.never;
25+
import static org.mockito.Mockito.verify;
26+
import static org.mockito.Mockito.verifyZeroInteractions;
27+
import static org.mockito.Mockito.when;
28+
29+
import com.google.common.collect.ImmutableList;
30+
import com.google.devtools.build.lib.actions.ActionExecutionContext;
31+
import com.google.devtools.build.lib.actions.DynamicStrategyRegistry;
32+
import com.google.devtools.build.lib.actions.SandboxedSpawnStrategy;
33+
import com.google.devtools.build.lib.actions.SandboxedSpawnStrategy.StopConcurrentSpawns;
34+
import com.google.devtools.build.lib.actions.Spawn;
35+
import com.google.devtools.build.lib.actions.SpawnResult;
36+
import com.google.devtools.build.lib.actions.SpawnResult.Status;
37+
import com.google.devtools.build.lib.exec.ExecutionPolicy;
38+
import com.google.devtools.build.lib.exec.util.SpawnBuilder;
39+
import com.google.devtools.build.lib.server.FailureDetails.Execution;
40+
import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
41+
import com.google.devtools.build.lib.testutil.TestFileOutErr;
42+
import com.google.devtools.build.lib.testutil.TestUtils;
43+
import java.util.Optional;
44+
import java.util.concurrent.ExecutorService;
45+
import java.util.concurrent.Executors;
46+
import java.util.concurrent.Semaphore;
47+
import java.util.function.Function;
48+
import org.junit.After;
49+
import org.junit.Before;
50+
import org.junit.Test;
51+
import org.junit.runner.RunWith;
52+
import org.junit.runners.JUnit4;
53+
import org.mockito.ArgumentCaptor;
54+
import org.mockito.Mock;
55+
import org.mockito.MockitoAnnotations;
56+
57+
/** Unit tests for {@link DynamicSpawnStrategy}. */
58+
@RunWith(JUnit4.class)
59+
public class DynamicSpawnStrategyUnitTest {
60+
61+
private static final SpawnResult SUCCESSFUL_SPAWN_RESULT =
62+
new SpawnResult.Builder().setRunnerName("test").setStatus(Status.SUCCESS).build();
63+
private static final FailureDetail FAILURE_DETAIL =
64+
FailureDetail.newBuilder().setExecution(Execution.getDefaultInstance()).build();
65+
66+
private ExecutorService executorServiceForCleanup;
67+
68+
@Mock private Function<Spawn, Optional<Spawn>> mockGetPostProcessingSpawn;
69+
70+
@Before
71+
public void initMocks() {
72+
MockitoAnnotations.initMocks(this);
73+
}
74+
75+
@After
76+
public void stopExecutorService() throws InterruptedException {
77+
executorServiceForCleanup.shutdown();
78+
assertThat(
79+
executorServiceForCleanup.awaitTermination(
80+
TestUtils.WAIT_TIMEOUT_MILLISECONDS, MILLISECONDS))
81+
.isTrue();
82+
}
83+
84+
@Test
85+
public void exec_remoteOnlySpawn_doesNotGetLocalPostProcessingSpawn() throws Exception {
86+
DynamicSpawnStrategy dynamicSpawnStrategy =
87+
createDynamicSpawnStrategy(
88+
ExecutionPolicy.REMOTE_EXECUTION_ONLY, mockGetPostProcessingSpawn);
89+
SandboxedSpawnStrategy local = createMockSpawnStrategy();
90+
SandboxedSpawnStrategy remote = createMockSpawnStrategy();
91+
ArgumentCaptor<Spawn> remoteSpawnCaptor = ArgumentCaptor.forClass(Spawn.class);
92+
when(remote.exec(remoteSpawnCaptor.capture(), any(), any()))
93+
.thenReturn(ImmutableList.of(SUCCESSFUL_SPAWN_RESULT));
94+
ActionExecutionContext actionExecutionContext = createMockActionExecutionContext(local, remote);
95+
Spawn spawn = new SpawnBuilder().build();
96+
97+
ImmutableList<SpawnResult> results = dynamicSpawnStrategy.exec(spawn, actionExecutionContext);
98+
99+
assertThat(results).containsExactly(SUCCESSFUL_SPAWN_RESULT);
100+
verify(mockGetPostProcessingSpawn, never()).apply(any());
101+
verify(local, never()).exec(any(), any(), any());
102+
assertThat(remoteSpawnCaptor.getAllValues()).containsExactly(spawn);
103+
}
104+
105+
@Test
106+
public void exec_localOnlySpawn_runsLocalPostProcessingSpawn() throws Exception {
107+
Spawn spawn = new SpawnBuilder("command").build();
108+
Spawn postProcessingSpawn = new SpawnBuilder("extra_command").build();
109+
DynamicSpawnStrategy dynamicSpawnStrategy =
110+
createDynamicSpawnStrategy(
111+
ExecutionPolicy.LOCAL_EXECUTION_ONLY, ignored -> Optional.of(postProcessingSpawn));
112+
SandboxedSpawnStrategy local = createMockSpawnStrategy();
113+
ArgumentCaptor<Spawn> localSpawnCaptor = ArgumentCaptor.forClass(Spawn.class);
114+
when(local.exec(localSpawnCaptor.capture(), any(), any()))
115+
.thenReturn(ImmutableList.of(SUCCESSFUL_SPAWN_RESULT));
116+
SandboxedSpawnStrategy remote = createMockSpawnStrategy();
117+
ActionExecutionContext actionExecutionContext = createMockActionExecutionContext(local, remote);
118+
119+
ImmutableList<SpawnResult> results = dynamicSpawnStrategy.exec(spawn, actionExecutionContext);
120+
121+
assertThat(results).containsExactly(SUCCESSFUL_SPAWN_RESULT, SUCCESSFUL_SPAWN_RESULT);
122+
verifyZeroInteractions(remote);
123+
assertThat(localSpawnCaptor.getAllValues())
124+
.containsExactly(spawn, postProcessingSpawn)
125+
.inOrder();
126+
}
127+
128+
@Test
129+
public void exec_failedLocalSpawn_doesNotGetLocalPostProcessingSpawn() throws Exception {
130+
testExecFailedLocalSpawnDoesNotGetLocalPostProcessingSpawn(
131+
new SpawnResult.Builder()
132+
.setRunnerName("test")
133+
.setStatus(Status.TIMEOUT)
134+
.setExitCode(SpawnResult.POSIX_TIMEOUT_EXIT_CODE)
135+
.setFailureDetail(FAILURE_DETAIL)
136+
.build());
137+
}
138+
139+
@Test
140+
public void exec_nonZeroExitCodeLocalSpawn_doesNotGetLocalPostProcessingSpawn() throws Exception {
141+
testExecFailedLocalSpawnDoesNotGetLocalPostProcessingSpawn(
142+
new SpawnResult.Builder()
143+
.setRunnerName("test")
144+
.setStatus(Status.EXECUTION_FAILED)
145+
.setExitCode(123)
146+
.setFailureDetail(FAILURE_DETAIL)
147+
.build());
148+
}
149+
150+
private void testExecFailedLocalSpawnDoesNotGetLocalPostProcessingSpawn(SpawnResult failedResult)
151+
throws Exception {
152+
DynamicSpawnStrategy dynamicSpawnStrategy =
153+
createDynamicSpawnStrategy(
154+
ExecutionPolicy.LOCAL_EXECUTION_ONLY, mockGetPostProcessingSpawn);
155+
SandboxedSpawnStrategy local = createMockSpawnStrategy();
156+
ArgumentCaptor<Spawn> localSpawnCaptor = ArgumentCaptor.forClass(Spawn.class);
157+
when(local.exec(localSpawnCaptor.capture(), any(), any()))
158+
.thenReturn(ImmutableList.of(failedResult));
159+
SandboxedSpawnStrategy remote = createMockSpawnStrategy();
160+
ActionExecutionContext actionExecutionContext = createMockActionExecutionContext(local, remote);
161+
Spawn spawn = new SpawnBuilder().build();
162+
163+
ImmutableList<SpawnResult> results = dynamicSpawnStrategy.exec(spawn, actionExecutionContext);
164+
165+
assertThat(results).containsExactly(failedResult);
166+
assertThat(localSpawnCaptor.getAllValues()).containsExactly(spawn);
167+
verify(remote, never()).exec(any(), any(), any());
168+
verify(mockGetPostProcessingSpawn, never()).apply(any());
169+
}
170+
171+
@Test
172+
public void exec_runAnywhereSpawn_runsLocalPostProcessingSpawn() throws Exception {
173+
Spawn spawn = new SpawnBuilder().build();
174+
Spawn postProcessingSpawn = new SpawnBuilder("extra_command").build();
175+
DynamicSpawnStrategy dynamicSpawnStrategy =
176+
createDynamicSpawnStrategy(
177+
ExecutionPolicy.ANYWHERE, ignored -> Optional.of(postProcessingSpawn));
178+
SandboxedSpawnStrategy local = createMockSpawnStrategy();
179+
// Make sure that local execution does not win the race before remote starts.
180+
Semaphore remoteStarted = new Semaphore(0);
181+
// Only the first spawn should be able to stop the concurrent remote execution (get the output
182+
// lock).
183+
when(local.exec(eq(spawn), any(), /*stopConcurrentSpawns=*/ isNotNull()))
184+
.thenAnswer(
185+
invocation -> {
186+
remoteStarted.acquire();
187+
StopConcurrentSpawns stopConcurrentSpawns = invocation.getArgument(2);
188+
stopConcurrentSpawns.stop();
189+
return ImmutableList.of(SUCCESSFUL_SPAWN_RESULT);
190+
});
191+
when(local.exec(eq(postProcessingSpawn), any(), /*stopConcurrentSpawns=*/ isNull()))
192+
.thenReturn(ImmutableList.of(SUCCESSFUL_SPAWN_RESULT));
193+
SandboxedSpawnStrategy remote = createMockSpawnStrategy();
194+
when(remote.exec(eq(spawn), any(), any()))
195+
.thenAnswer(
196+
invocation -> {
197+
remoteStarted.release();
198+
Thread.sleep(TestUtils.WAIT_TIMEOUT_MILLISECONDS);
199+
throw new AssertionError("Timed out waiting for interruption");
200+
});
201+
ActionExecutionContext actionExecutionContext = createMockActionExecutionContext(local, remote);
202+
203+
ImmutableList<SpawnResult> results = dynamicSpawnStrategy.exec(spawn, actionExecutionContext);
204+
205+
assertThat(results).containsExactly(SUCCESSFUL_SPAWN_RESULT, SUCCESSFUL_SPAWN_RESULT);
206+
}
207+
208+
private DynamicSpawnStrategy createDynamicSpawnStrategy(
209+
ExecutionPolicy executionPolicy,
210+
Function<Spawn, Optional<Spawn>> getPostProcessingSpawnForLocalExecution) {
211+
checkState(
212+
executorServiceForCleanup == null,
213+
"Creating the DynamicSpawnStrategy twice in the same test is not supported.");
214+
executorServiceForCleanup = Executors.newCachedThreadPool();
215+
return new DynamicSpawnStrategy(
216+
executorServiceForCleanup,
217+
new DynamicExecutionOptions(),
218+
ignored -> executionPolicy,
219+
getPostProcessingSpawnForLocalExecution);
220+
}
221+
222+
private static ActionExecutionContext createMockActionExecutionContext(
223+
SandboxedSpawnStrategy localStrategy, SandboxedSpawnStrategy remoteStrategy) {
224+
ActionExecutionContext actionExecutionContext = mock(ActionExecutionContext.class);
225+
when(actionExecutionContext.getFileOutErr()).thenReturn(new TestFileOutErr());
226+
when(actionExecutionContext.getContext(DynamicStrategyRegistry.class))
227+
.thenReturn(
228+
new DynamicStrategyRegistry() {
229+
@Override
230+
public ImmutableList<SandboxedSpawnStrategy> getDynamicSpawnActionContexts(
231+
Spawn spawn, DynamicMode dynamicMode) {
232+
switch (dynamicMode) {
233+
case LOCAL:
234+
return ImmutableList.of(localStrategy);
235+
case REMOTE:
236+
return ImmutableList.of(remoteStrategy);
237+
}
238+
throw new AssertionError("Unexpected mode: " + dynamicMode);
239+
}
240+
241+
@Override
242+
public void notifyUsedDynamic(ActionContextRegistry actionContextRegistry) {}
243+
});
244+
when(actionExecutionContext.withFileOutErr(any())).thenReturn(actionExecutionContext);
245+
return actionExecutionContext;
246+
}
247+
248+
private static SandboxedSpawnStrategy createMockSpawnStrategy() throws InterruptedException {
249+
SandboxedSpawnStrategy strategy = mock(SandboxedSpawnStrategy.class);
250+
when(strategy.canExec(any(), any())).thenReturn(true);
251+
when(strategy.beginExecution(any(), any())).thenThrow(UnsupportedOperationException.class);
252+
return strategy;
253+
}
254+
}

0 commit comments

Comments
 (0)