forked from bazelbuild/rules_kotlin
-
Notifications
You must be signed in to change notification settings - Fork 0
/
BazelWorker.kt
164 lines (155 loc) · 6.29 KB
/
BazelWorker.kt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
/*
* Copyright 2018 The Bazel Authors. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* 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.
*/
package io.bazel.kotlin.builder.tasks
import com.google.devtools.build.lib.worker.WorkerProtocol
import io.bazel.kotlin.builder.utils.rootCause
import java.io.*
import java.nio.charset.StandardCharsets.UTF_8
import java.nio.file.Files
import java.nio.file.Paths
/**
* Interface for command line programs.
*
* This is the same thing as a main function, except not static.
*/
interface CommandLineProgram {
/**
* Runs blocking program start to finish.
*
* This function might be called multiple times throughout the life of this object. Output
* must be sent to [System.out] and [System.err].
*
* @param args command line arguments
* @return program exit code, i.e. 0 for success, non-zero for failure
*/
fun apply(args: List<String>): Int
}
/**
* Bazel worker runner.
*
* This class adapts a traditional command line program so it can be spawned by Bazel as a
* persistent worker process that handles multiple invocations per JVM. It will also be backwards
* compatible with being run as a normal single-invocation command.
*
* @param <T> delegate program type
</T> */
class BazelWorker(
private val delegate: CommandLineProgram,
private val output: PrintStream,
private val mnemonic: String
) : CommandLineProgram {
companion object {
private const val INTERUPTED_STATUS = 0
private const val ERROR_STATUS = 1
}
override fun apply(args: List<String>): Int {
return if (args.contains("--persistent_worker"))
runAsPersistentWorker(args)
else delegate.apply(loadArguments(args, false))
}
@Suppress("UNUSED_PARAMETER")
private fun runAsPersistentWorker(ignored: List<String>): Int {
val realStdIn = System.`in`
val realStdOut = System.out
val realStdErr = System.err
try {
ByteArrayInputStream(ByteArray(0)).use { emptyIn ->
ByteArrayOutputStream().use { buffer ->
PrintStream(buffer).use { ps ->
System.setIn(emptyIn)
System.setOut(ps)
System.setErr(ps)
while (true) {
val request = WorkerProtocol.WorkRequest.parseDelimitedFrom(realStdIn) ?: return 0
val exitCode = try {
delegate.apply(loadArguments(request.argumentsList, true))
} catch (e: RuntimeException) {
val innerExitCode = if (wasInterrupted(e)) INTERUPTED_STATUS
else ERROR_STATUS.also {
System.err.println(
"ERROR: Worker threw uncaught exception with args: ${request.argumentsList.joinToString(" ")}"
)
e.printStackTrace(System.err)
}
WorkerProtocol.WorkResponse.newBuilder()
.setOutput(buffer.toString())
.setExitCode(innerExitCode)
.build()
.writeDelimitedTo(realStdOut)
realStdOut.flush()
return innerExitCode
}
WorkerProtocol.WorkResponse.newBuilder()
.setOutput(buffer.toString())
.setExitCode(exitCode)
.build()
.writeDelimitedTo(realStdOut)
realStdOut.flush()
buffer.reset()
System.gc() // be a good little worker process and consume less memory when idle
}
}
}
}
} catch (e: IOException) {
if (wasInterrupted(e)) {
return INTERUPTED_STATUS
}
throw e
} catch (e: RuntimeException) {
if (wasInterrupted(e)) {
return INTERUPTED_STATUS
}
throw e
} finally {
System.setIn(realStdIn)
System.setOut(realStdOut)
System.setErr(realStdErr)
}
throw RuntimeException("drop through")
}
private fun loadArguments(args: List<String>, isWorker: Boolean): List<String> {
if (args.isNotEmpty()) {
val lastArg = args[args.size - 1]
if (lastArg.startsWith("@")) {
val pathElement = lastArg.substring(1)
val flagFile = Paths.get(pathElement)
if (isWorker && lastArg.startsWith("@@") || Files.exists(flagFile)) {
if (!isWorker && !mnemonic.isEmpty()) {
output.printf(
"HINT: %s will compile faster if you run: " + "echo \"build --strategy=%s=worker\" >>~/.bazelrc\n",
mnemonic, mnemonic
)
}
try {
return Files.readAllLines(flagFile, UTF_8)
} catch (e: IOException) {
throw RuntimeException(e)
}
}
}
}
return args
}
private fun wasInterrupted(e: Throwable): Boolean {
val cause = e.rootCause
if (cause is InterruptedException || cause is InterruptedIOException) {
output.println("Terminating worker due to interrupt signal")
return true
}
return false
}
}