Skip to content

Commit e9e6978

Browse files
larsrc-googlecopybara-github
larsrc-google
authored and
copybara-github
committed
Server-side implementation of worker cancellation.
RELNOTES: None. PiperOrigin-RevId: 368598866
1 parent 7e790a0 commit e9e6978

9 files changed

Lines changed: 158 additions & 30 deletions

File tree

‎src/main/java/com/google/devtools/build/lib/worker/WorkerKey.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,8 @@ final class WorkerKey {
5353
private final boolean isSpeculative;
5454
/** A WorkerProxy will be instantiated if true, instantiate a regular Worker if false. */
5555
private final boolean proxied;
56+
/** If true, the workers for this key are able to cancel work requests. */
57+
private final boolean cancellable;
5658
/**
5759
* Cached value for the hash of this key, because the value is expensive to calculate
5860
* (ImmutableMap and ImmutableList do not cache their hashcodes.
@@ -70,6 +72,7 @@ final class WorkerKey {
7072
SortedMap<PathFragment, HashCode> workerFilesWithHashes,
7173
boolean isSpeculative,
7274
boolean proxied,
75+
boolean cancellable,
7376
WorkerProtocolFormat protocolFormat) {
7477
this.args = Preconditions.checkNotNull(args);
7578
this.env = Preconditions.checkNotNull(env);
@@ -79,8 +82,8 @@ final class WorkerKey {
7982
this.workerFilesWithHashes = Preconditions.checkNotNull(workerFilesWithHashes);
8083
this.isSpeculative = isSpeculative;
8184
this.proxied = proxied;
85+
this.cancellable = cancellable;
8286
this.protocolFormat = protocolFormat;
83-
8487
hash = calculateHashCode();
8588
}
8689

@@ -128,6 +131,10 @@ public boolean isMultiplex() {
128131
return getProxied() && !isSpeculative;
129132
}
130133

134+
public boolean isCancellable() {
135+
return cancellable;
136+
}
137+
131138
/** Returns the format of the worker protocol. */
132139
public WorkerProtocolFormat getProtocolFormat() {
133140
return protocolFormat;

‎src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnRunner.java‎

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,6 @@ final class WorkerSpawnRunner implements SpawnRunner {
7777
public static final String REASON_NO_FLAGFILE =
7878
"because the command-line arguments do not contain at least one @flagfile or --flagfile=";
7979
public static final String REASON_NO_TOOLS = "because the action has no tools";
80-
public static final String REASON_NO_EXECUTION_INFO =
81-
"because the action's execution info does not contain 'supports-workers=1'";
8280

8381
/** Pattern for @flagfile.txt and --flagfile=flagfile.txt */
8482
private static final Pattern FLAG_FILE_PATTERN = Pattern.compile("(?:@|--?flagfile=)(.+)");
@@ -205,6 +203,7 @@ public SpawnResult exec(Spawn spawn, SpawnExecutionContext context)
205203
workerFiles,
206204
context.speculating(),
207205
multiplex && Spawns.supportsMultiplexWorkers(spawn),
206+
Spawns.supportsWorkerCancellation(spawn),
208207
protocolFormat);
209208

210209
SpawnMetrics.Builder spawnMetrics =
@@ -458,7 +457,11 @@ WorkResponse execInWorker(
458457
try {
459458
response = worker.getResponse(request.getRequestId());
460459
} catch (InterruptedException e) {
461-
finishWorkAsync(key, worker, request);
460+
finishWorkAsync(
461+
key,
462+
worker,
463+
request,
464+
workerOptions.workerCancellation && Spawns.supportsWorkerCancellation(spawn));
462465
worker = null;
463466
throw e;
464467
} catch (IOException e) {
@@ -480,6 +483,12 @@ WorkResponse execInWorker(
480483
throw createEmptyResponseException(worker.getLogFile());
481484
}
482485

486+
if (response.getWasCancelled()) {
487+
throw createUserExecException(
488+
"Received cancel response for " + response.getRequestId() + " without having cancelled",
489+
Code.FINISH_FAILURE);
490+
}
491+
483492
try {
484493
Stopwatch processOutputsStopwatch = Stopwatch.createStarted();
485494
context.lockOutputFiles();
@@ -525,12 +534,21 @@ WorkResponse execInWorker(
525534
* interrupted. This takes ownership of the worker for purposes of returning it to the worker
526535
* pool.
527536
*/
528-
private void finishWorkAsync(WorkerKey key, Worker worker, WorkRequest request) {
537+
private void finishWorkAsync(
538+
WorkerKey key, Worker worker, WorkRequest request, boolean canCancel) {
529539
Thread reaper =
530540
new Thread(
531541
() -> {
532542
Worker w = worker;
533543
try {
544+
if (canCancel) {
545+
WorkRequest cancelRequest =
546+
WorkRequest.newBuilder()
547+
.setRequestId(request.getRequestId())
548+
.setCancel(true)
549+
.build();
550+
w.putRequest(cancelRequest);
551+
}
534552
w.getResponse(request.getRequestId());
535553
} catch (IOException | InterruptedException e1) {
536554
// If this happens, we either can't trust the output of the worker, or we got
@@ -549,7 +567,8 @@ private void finishWorkAsync(WorkerKey key, Worker worker, WorkRequest request)
549567
workers.returnObject(key, w);
550568
}
551569
}
552-
});
570+
},
571+
"AsyncFinish-Worker-" + worker.workerId);
553572
reaper.start();
554573
}
555574

‎src/test/java/com/google/devtools/build/lib/worker/ExampleWorker.java‎

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.google.common.collect.ImmutableSet;
2323
import com.google.devtools.build.lib.actions.ExecutionRequirements.WorkerProtocolFormat;
2424
import com.google.devtools.build.lib.worker.ExampleWorkerOptions.ExampleWorkOptions;
25+
import com.google.devtools.build.lib.worker.WorkRequestHandler.WorkerMessageProcessor;
2526
import com.google.devtools.build.lib.worker.WorkerProtocol.Input;
2627
import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
2728
import com.google.devtools.common.options.OptionsParser;
@@ -42,12 +43,9 @@
4243
import java.util.Map;
4344
import java.util.Random;
4445
import java.util.UUID;
45-
import java.util.concurrent.Semaphore;
4646
import java.util.function.BiFunction;
4747
import java.util.regex.Matcher;
4848
import java.util.regex.Pattern;
49-
import sun.misc.Signal;
50-
import sun.misc.SignalHandler;
5149

5250
/** An example implementation of a worker process that is used for integration tests. */
5351
public final class ExampleWorker {
@@ -70,6 +68,7 @@ public final class ExampleWorker {
7068

7169
// The options passed to this worker on a per-worker-lifetime basis.
7270
static ExampleWorkerOptions workerOptions;
71+
private static WorkerMessageProcessor messageProcessor;
7372

7473
private static class InterruptableWorkRequestHandler extends WorkRequestHandler {
7574

@@ -118,7 +117,7 @@ public static void main(String[] args) throws Exception {
118117
parser.parse(args);
119118
workerOptions = parser.getOptions(ExampleWorkerOptions.class);
120119
WorkerProtocolFormat protocolFormat = workerOptions.workerProtocol;
121-
WorkRequestHandler.WorkerMessageProcessor messageProcessor = null;
120+
messageProcessor = null;
122121
switch (protocolFormat) {
123122
case JSON:
124123
messageProcessor =
@@ -147,21 +146,23 @@ private static int doWork(List<String> args, PrintWriter err) {
147146
PrintStream originalStdOut = System.out;
148147
PrintStream originalStdErr = System.err;
149148

150-
if (workerOptions.waitForSignal) {
151-
Semaphore signalSem = new Semaphore(0);
152-
Signal.handle(
153-
new Signal("HUP"),
154-
new SignalHandler() {
155-
@Override
156-
public void handle(Signal sig) {
157-
signalSem.release();
158-
}
159-
});
149+
if (workerOptions.waitForCancel) {
160150
try {
161-
signalSem.acquire();
162-
} catch (InterruptedException e) {
163-
System.out.println("Interrupted while waiting for signal");
164-
e.printStackTrace();
151+
WorkRequest workRequest = messageProcessor.readWorkRequest();
152+
if (workRequest.getRequestId() != currentRequest.getRequestId()) {
153+
System.err.format(
154+
"Got cancel request for %d while expecting cancel request for %d%n",
155+
workRequest.getRequestId(), currentRequest.getRequestId());
156+
return 1;
157+
}
158+
if (!workRequest.getCancel()) {
159+
System.err.format(
160+
"Got non-cancel request for %d while expecting cancel request%n",
161+
workRequest.getRequestId());
162+
return 1;
163+
}
164+
} catch (IOException e) {
165+
throw new RuntimeException("Exception while waiting for cancel request", e);
165166
}
166167
}
167168
try (PrintStream ps = new PrintStream(baos)) {

‎src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerOptions.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -136,12 +136,12 @@ public static class ExampleWorkOptions extends OptionsBase {
136136
public boolean hardPoison;
137137

138138
@Option(
139-
name = "wait_for_signal",
139+
name = "wait_for_cancel",
140140
documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
141141
effectTags = {OptionEffectTag.NO_OP},
142142
defaultValue = "false",
143-
help = "Don't send a response until receiving a SIGXXXX.")
144-
public boolean waitForSignal;
143+
help = "Don't send a response until receiving a cancel request.")
144+
public boolean waitForCancel;
145145

146146
/** Enum converter for --worker_protocol. */
147147
public static class WorkerProtocolEnumConverter

‎src/test/java/com/google/devtools/build/lib/worker/TestUtils.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ static WorkerKey createWorkerKey(
4545
/* workerFilesWithHashes= */ ImmutableSortedMap.of(),
4646
/* mustBeSandboxed= */ false,
4747
/* proxied= */ proxied,
48+
/* cancellable= */ false,
4849
WorkerProtocolFormat.PROTO);
4950
}
5051

@@ -58,6 +59,7 @@ static WorkerKey createWorkerKey(WorkerProtocolFormat protocolFormat, FileSystem
5859
/* workerFilesWithHashes= */ ImmutableSortedMap.of(),
5960
/* mustBeSandboxed= */ true,
6061
/* proxied= */ true,
62+
/* cancellable= */ false,
6163
protocolFormat);
6264
}
6365

‎src/test/java/com/google/devtools/build/lib/worker/WorkerFactoryTest.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ protected WorkerKey createWorkerKey(boolean mustBeSandboxed, boolean proxied, St
5858
/* workerFilesWithHashes= */ ImmutableSortedMap.of(),
5959
/* mustBeSandboxed= */ mustBeSandboxed,
6060
/* proxied= */ proxied,
61+
/* cancellable= */ false,
6162
WorkerProtocolFormat.PROTO);
6263
}
6364

‎src/test/java/com/google/devtools/build/lib/worker/WorkerKeyTest.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ private WorkerKey makeWorkerKey(boolean multiplex, boolean dynamic) {
4343
/* workerFilesWithHashes= */ ImmutableSortedMap.of(),
4444
/* isSpeculative= */ dynamic,
4545
/* proxied= */ multiplex,
46+
/* cancellable=*/ false,
4647
WorkerProtocolFormat.PROTO);
4748
}
4849

@@ -90,6 +91,7 @@ public void testWorkerKeyEquality() {
9091
workerKey.getWorkerFilesWithHashes(),
9192
workerKey.isSpeculative(),
9293
workerKey.getProxied(),
94+
workerKey.isCancellable(),
9395
workerKey.getProtocolFormat());
9496
assertThat(workerKey).isEqualTo(workerKeyWithSameFields);
9597
}
@@ -107,6 +109,7 @@ public void testWorkerKeyInequality_protocol() {
107109
workerKey.getWorkerFilesWithHashes(),
108110
workerKey.isSpeculative(),
109111
workerKey.getProxied(),
112+
workerKey.isCancellable(),
110113
WorkerProtocolFormat.JSON);
111114
assertThat(workerKey).isNotEqualTo(workerKeyWithDifferentProtocol);
112115
}

‎src/test/java/com/google/devtools/build/lib/worker/WorkerMultiplexerManagerTest.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ public void instanceCreationRemovalTest() throws Exception {
5959
ImmutableSortedMap.of(),
6060
false,
6161
false,
62+
/* cancellable= */ false,
6263
WorkerProtocolFormat.PROTO);
6364
WorkerMultiplexer wm1 = WorkerMultiplexerManager.getInstance(workerKey1, logFile);
6465

@@ -77,6 +78,7 @@ public void instanceCreationRemovalTest() throws Exception {
7778
ImmutableSortedMap.of(),
7879
false,
7980
false,
81+
/* cancellable= */ false,
8082
WorkerProtocolFormat.PROTO);
8183
WorkerMultiplexer wm2 = WorkerMultiplexerManager.getInstance(workerKey2, logFile);
8284

0 commit comments

Comments
 (0)