Skip to content

Commit 37a95cd

Browse files
loxlox
authored
fix: allow console to force local exit on second interrupt (#82)
* fix: allow console to force local exit on second interrupt Amp-Thread-ID: https://ampcode.com/threads/T-019cc2a9-921d-703a-b59d-22a1cb4b4bc0 * fix: tighten console interrupt escape handling Amp-Thread-ID: https://ampcode.com/threads/T-019cc2a9-921d-703a-b59d-22a1cb4b4bc0 * fix: dispatch first console interrupt synchronously Amp-Thread-ID: https://ampcode.com/threads/T-019cc2a9-921d-703a-b59d-22a1cb4b4bc0
1 parent a95df92 commit 37a95cd

File tree

2 files changed

+179
-7
lines changed

2 files changed

+179
-7
lines changed

‎internal/cli/cli.go‎

Lines changed: 73 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ import (
1717
"runtime"
1818
"strconv"
1919
"strings"
20+
"sync"
21+
"sync/atomic"
2022
"syscall"
2123
"text/tabwriter"
2224
"time"
@@ -51,10 +53,11 @@ import (
5153
const defaultBumpRefSource = "ghcr.io/buildkite/cleanroom-base/alpine:latest"
5254

5355
const (
54-
systemdServiceName = "cleanroom.service"
55-
launchdServiceName = "com.buildkite.cleanroom"
56-
defaultDaemonListen = "unix://" + endpoint.DefaultSystemSocketPath
57-
sandboxTerminateTimeout = 2 * time.Second
56+
systemdServiceName = "cleanroom.service"
57+
launchdServiceName = "com.buildkite.cleanroom"
58+
defaultDaemonListen = "unix://" + endpoint.DefaultSystemSocketPath
59+
sandboxTerminateTimeout = 2 * time.Second
60+
interruptForceExitWindow = 1200 * time.Millisecond
5861
)
5962

6063
type policyLoader interface {
@@ -1106,6 +1109,38 @@ func (c *ConsoleCommand) Run(ctx *runtimeContext) error {
11061109
notifySignals(signalCh, os.Interrupt, syscall.SIGTERM)
11071110
defer stopSignals(signalCh)
11081111

1112+
var interruptCount atomic.Int32
1113+
var lastInterruptAt atomic.Int64
1114+
forceLocalExit := make(chan struct{})
1115+
var forceLocalExitOnce sync.Once
1116+
requestInterrupt := func(signal int32) {
1117+
now := time.Now()
1118+
last := time.Unix(0, lastInterruptAt.Load())
1119+
if !last.IsZero() && now.Sub(last) > interruptForceExitWindow {
1120+
interruptCount.Store(0)
1121+
}
1122+
count := interruptCount.Add(1)
1123+
lastInterruptAt.Store(now.UnixNano())
1124+
if count == 1 {
1125+
_ = interactiveSession.SendSignal(signal)
1126+
return
1127+
}
1128+
forceLocalExitOnce.Do(func() {
1129+
close(forceLocalExit)
1130+
go func() {
1131+
_ = interactiveSession.Close()
1132+
}()
1133+
})
1134+
}
1135+
isForceLocalExit := func() bool {
1136+
select {
1137+
case <-forceLocalExit:
1138+
return true
1139+
default:
1140+
return false
1141+
}
1142+
}
1143+
11091144
if rawMode {
11101145
resizeSignalCh := make(chan os.Signal, 4)
11111146
signal.Notify(resizeSignalCh, syscall.SIGWINCH)
@@ -1127,7 +1162,7 @@ func (c *ConsoleCommand) Run(ctx *runtimeContext) error {
11271162
if sig == syscall.SIGTERM {
11281163
num = 15
11291164
}
1130-
_ = interactiveSession.SendSignal(num)
1165+
requestInterrupt(num)
11311166
}
11321167
}()
11331168

@@ -1137,8 +1172,24 @@ func (c *ConsoleCommand) Run(ctx *runtimeContext) error {
11371172
n, readErr := os.Stdin.Read(buf)
11381173
if n > 0 {
11391174
payload := append([]byte(nil), buf[:n]...)
1140-
if sendErr := interactiveSession.WriteStdin(payload); sendErr != nil {
1141-
return
1175+
if rawMode {
1176+
filtered := payload[:0]
1177+
for _, b := range payload {
1178+
if b == 0x03 {
1179+
requestInterrupt(2)
1180+
continue
1181+
}
1182+
filtered = append(filtered, b)
1183+
}
1184+
payload = filtered
1185+
if isForceLocalExit() {
1186+
return
1187+
}
1188+
}
1189+
if len(payload) > 0 {
1190+
if sendErr := interactiveSession.WriteStdin(payload); sendErr != nil {
1191+
return
1192+
}
11421193
}
11431194
}
11441195
if readErr != nil {
@@ -1179,6 +1230,9 @@ func (c *ConsoleCommand) Run(ctx *runtimeContext) error {
11791230
buf := make([]byte, 4096)
11801231
for {
11811232
n, readErr := interactiveSession.ReadPTY(buf)
1233+
if isForceLocalExit() {
1234+
return exitCodeError{code: 130}
1235+
}
11821236
if n > 0 {
11831237
chunk := append([]byte(nil), buf[:n]...)
11841238
if rawMode {
@@ -1189,6 +1243,9 @@ func (c *ConsoleCommand) Run(ctx *runtimeContext) error {
11891243
}
11901244
}
11911245
if polledExitCode, gotExitCode, pollErr := pollInteractiveExitOrControlErr(exitCodeCh, &controlErrCh); pollErr != nil {
1246+
if isForceLocalExit() {
1247+
return exitCodeError{code: 130}
1248+
}
11921249
return pollErr
11931250
} else if gotExitCode {
11941251
exitCode = polledExitCode
@@ -1201,17 +1258,26 @@ func (c *ConsoleCommand) Run(ctx *runtimeContext) error {
12011258
break
12021259
}
12031260
}
1261+
if isForceLocalExit() {
1262+
return exitCodeError{code: 130}
1263+
}
12041264

12051265
if !haveExitCode {
12061266
waitedExitCode, gotExitCode, waitErr := waitForInteractiveExitOrControlErr(exitCodeCh, &controlErrCh, 2*time.Second)
12071267
if waitErr != nil {
1268+
if isForceLocalExit() {
1269+
return exitCodeError{code: 130}
1270+
}
12081271
return waitErr
12091272
}
12101273
if gotExitCode {
12111274
exitCode = waitedExitCode
12121275
haveExitCode = true
12131276
}
12141277
}
1278+
if isForceLocalExit() {
1279+
return exitCodeError{code: 130}
1280+
}
12151281

12161282
if !haveExitCode {
12171283
if fetchedExitCode, ok := getFinalExecutionExitCode(client, sandboxID, executionID); ok {

‎internal/cli/console_integration_test.go‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,112 @@ func TestConsoleIntegrationInterruptCancelsExecution(t *testing.T) {
141141
}
142142
}
143143

144+
func TestConsoleIntegrationSecondInterruptForcesLocalExitWhenExecutionIgnoresCancel(t *testing.T) {
145+
started := make(chan struct{}, 1)
146+
releaseRun := make(chan struct{})
147+
t.Cleanup(func() {
148+
close(releaseRun)
149+
})
150+
adapter := &integrationAdapter{
151+
runStreamFn: func(_ context.Context, req backend.RunRequest, stream backend.OutputStream) (*backend.RunResult, error) {
152+
if stream.OnAttach != nil {
153+
stream.OnAttach(backend.AttachIO{
154+
WriteStdin: func(_ []byte) error { return nil },
155+
})
156+
}
157+
select {
158+
case started <- struct{}{}:
159+
default:
160+
}
161+
<-releaseRun
162+
return &backend.RunResult{RunID: req.RunID, ExitCode: 0, Message: "released"}, nil
163+
},
164+
}
165+
166+
host, _ := startIntegrationServer(t, adapter)
167+
signalCh := withTestSignalChannel(t)
168+
cwd := t.TempDir()
169+
170+
done := make(chan execOutcome, 1)
171+
go func() {
172+
done <- runConsoleWithCapture(ConsoleCommand{
173+
clientFlags: clientFlags{Host: host},
174+
Chdir: cwd,
175+
}, "", runtimeContext{
176+
CWD: cwd,
177+
Loader: integrationLoader{},
178+
})
179+
}()
180+
181+
_ = mustReceiveWithin(t, started, 2*time.Second, "timed out waiting for console execution to start")
182+
signalCh <- os.Interrupt
183+
signalCh <- os.Interrupt
184+
185+
outcome := mustReceiveWithin(t, done, 2*time.Second, "timed out waiting for second interrupt to force local console exit")
186+
if outcome.cause != nil {
187+
t.Fatalf("capture failure: %v", outcome.cause)
188+
}
189+
if outcome.err == nil {
190+
t.Fatal("expected non-zero exit from forced local interrupt")
191+
}
192+
if got, want := ExitCode(outcome.err), 130; got != want {
193+
t.Fatalf("unexpected console exit code: got %d want %d (err=%v)", got, want, outcome.err)
194+
}
195+
}
196+
197+
func TestConsoleIntegrationSecondInterruptAfterWindowDoesNotForceLocalExit(t *testing.T) {
198+
started := make(chan struct{}, 1)
199+
releaseRun := make(chan struct{})
200+
adapter := &integrationAdapter{
201+
runStreamFn: func(_ context.Context, req backend.RunRequest, stream backend.OutputStream) (*backend.RunResult, error) {
202+
if stream.OnAttach != nil {
203+
stream.OnAttach(backend.AttachIO{WriteStdin: func(_ []byte) error { return nil }})
204+
}
205+
select {
206+
case started <- struct{}{}:
207+
default:
208+
}
209+
<-releaseRun
210+
return &backend.RunResult{RunID: req.RunID, ExitCode: 0, Message: "released"}, nil
211+
},
212+
}
213+
214+
host, _ := startIntegrationServer(t, adapter)
215+
signalCh := withTestSignalChannel(t)
216+
cwd := t.TempDir()
217+
218+
done := make(chan execOutcome, 1)
219+
go func() {
220+
done <- runConsoleWithCapture(ConsoleCommand{
221+
clientFlags: clientFlags{Host: host},
222+
Chdir: cwd,
223+
}, "", runtimeContext{CWD: cwd, Loader: integrationLoader{}})
224+
}()
225+
226+
_ = mustReceiveWithin(t, started, 2*time.Second, "timed out waiting for console execution to start")
227+
signalCh <- os.Interrupt
228+
time.Sleep(2*interruptForceExitWindow + 100*time.Millisecond)
229+
signalCh <- os.Interrupt
230+
231+
select {
232+
case outcome := <-done:
233+
t.Fatalf("unexpected forced local exit after interrupt window elapsed: err=%v cause=%v", outcome.err, outcome.cause)
234+
case <-time.After(300 * time.Millisecond):
235+
}
236+
237+
close(releaseRun)
238+
outcome := mustReceiveWithin(t, done, 2*time.Second, "timed out waiting for console execution to finish after release")
239+
if outcome.cause != nil {
240+
t.Fatalf("capture failure: %v", outcome.cause)
241+
}
242+
if outcome.err == nil {
243+
t.Fatal("expected canceled exit after interrupt")
244+
}
245+
if got, want := ExitCode(outcome.err), 130; got != want {
246+
t.Fatalf("unexpected console exit code: got %d want %d (err=%v)", got, want, outcome.err)
247+
}
248+
}
249+
144250
func TestConsoleRejectsUnsupportedHostScheme(t *testing.T) {
145251
outcome := runConsoleWithCapture(ConsoleCommand{
146252
clientFlags: clientFlags{Host: "tssvc://cleanroom"},

0 commit comments

Comments
 (0)