Skip to content

Commit 72e5358

Browse files
tarasmadantarasmadan
committed
vm: func Run accepts context
It allows to use context as a single termination signal source.
1 parent cfde826 commit 72e5358

File tree

20 files changed

+76
-94
lines changed

20 files changed

+76
-94
lines changed

‎pkg/build/netbsd.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package build
55

66
import (
7+
"context"
78
"encoding/json"
89
"fmt"
910
"os"
@@ -155,7 +156,9 @@ func (ctx netbsd) copyKernelToDisk(targetArch, vmType, outputDir, kernel string)
155156
}
156157
commands = append(commands, "mknod /dev/vhci c 355 0")
157158
commands = append(commands, "sync") // Run sync so that the copied image is stored properly.
158-
_, rep, err := inst.Run(time.Minute, reporter, strings.Join(commands, ";"))
159+
ctxTimeout, cancel := context.WithTimeout(context.Background(), time.Minute)
160+
defer cancel()
161+
_, rep, err := inst.Run(ctxTimeout, reporter, strings.Join(commands, ";"))
159162
if err != nil {
160163
return fmt.Errorf("error syncing the instance %w", err)
161164
}

‎pkg/instance/execprog.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package instance
55

66
import (
7+
"context"
78
"fmt"
89
"os"
910
"time"
@@ -123,7 +124,9 @@ func (inst *ExecProgInstance) runCommand(command string, duration time.Duration,
123124
if inst.BeforeContextLen != 0 {
124125
opts = append(opts, vm.OutputSize(inst.BeforeContextLen))
125126
}
126-
output, rep, err := inst.VMInstance.Run(duration, inst.reporter, command, opts...)
127+
ctxTimeout, cancel := context.WithTimeout(context.Background(), duration)
128+
defer cancel()
129+
output, rep, err := inst.VMInstance.Run(ctxTimeout, inst.reporter, command, opts...)
127130
if err != nil {
128131
return nil, fmt.Errorf("failed to run command in VM: %w", err)
129132
}

‎pkg/manager/diff.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -535,8 +535,10 @@ func (kc *kernelContext) runInstance(ctx context.Context, inst *vm.Instance,
535535
return nil, fmt.Errorf("failed to parse manager's address")
536536
}
537537
cmd := fmt.Sprintf("%v runner %v %v %v", executorBin, inst.Index(), host, port)
538-
_, rep, err := inst.Run(kc.cfg.Timeouts.VMRunningTime, kc.reporter, cmd,
539-
vm.ExitTimeout, vm.StopContext(ctx), vm.InjectExecuting(injectExec),
538+
ctxTimeout, cancel := context.WithTimeout(ctx, kc.cfg.Timeouts.VMRunningTime)
539+
defer cancel()
540+
_, rep, err := inst.Run(ctxTimeout, kc.reporter, cmd, vm.ExitTimeout,
541+
vm.InjectExecuting(injectExec),
540542
vm.EarlyFinishCb(func() {
541543
// Depending on the crash type and kernel config, fuzzing may continue
542544
// running for several seconds even after kernel has printed a crash report.

‎syz-manager/manager.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -652,8 +652,10 @@ func (mgr *Manager) runInstanceInner(ctx context.Context, inst *vm.Instance, inj
652652
return nil, nil, fmt.Errorf("failed to parse manager's address")
653653
}
654654
cmd := fmt.Sprintf("%v runner %v %v %v", executorBin, inst.Index(), host, port)
655-
_, rep, err := inst.Run(mgr.cfg.Timeouts.VMRunningTime, mgr.reporter, cmd,
656-
vm.ExitTimeout, vm.StopContext(ctx), vm.InjectExecuting(injectExec),
655+
ctxTimeout, cancel := context.WithTimeout(ctx, mgr.cfg.Timeouts.VMRunningTime)
656+
defer cancel()
657+
_, rep, err := inst.Run(ctxTimeout, mgr.reporter, cmd,
658+
vm.ExitTimeout, vm.InjectExecuting(injectExec),
657659
finishCb,
658660
)
659661
if err != nil {

‎syz-manager/snapshot.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,9 @@ func (mgr *Manager) snapshotLoop(ctx context.Context, inst *vm.Instance) error {
4141
// All network connections (including ssh) will break once we start restoring snapshots.
4242
// So we start a background process and log to /dev/kmsg.
4343
cmd := fmt.Sprintf("nohup %v exec snapshot 1>/dev/null 2>/dev/kmsg </dev/null &", executor)
44-
if _, _, err := inst.Run(time.Hour, mgr.reporter, cmd); err != nil {
44+
ctxTimeout, cancel := context.WithTimeout(ctx, time.Hour)
45+
defer cancel()
46+
if _, _, err := inst.Run(ctxTimeout, mgr.reporter, cmd); err != nil {
4547
return err
4648
}
4749

‎vm/adb/adb.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ package adb
77

88
import (
99
"bytes"
10+
"context"
1011
"encoding/json"
1112
"fmt"
1213
"io"
@@ -521,7 +522,7 @@ func isRemoteCuttlefish(dev string) (bool, string) {
521522
return true, ip
522523
}
523524

524-
func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
525+
func (inst *instance) Run(ctx context.Context, command string) (
525526
<-chan []byte, <-chan error, error) {
526527
var tty io.ReadCloser
527528
var err error
@@ -566,9 +567,8 @@ func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command strin
566567
merger.Add("console", tty)
567568
merger.Add("adb", adbRpipe)
568569

569-
return vmimpl.Multiplex(adb, merger, timeout, vmimpl.MultiplexConfig{
570+
return vmimpl.Multiplex(ctx, adb, merger, vmimpl.MultiplexConfig{
570571
Console: tty,
571-
Stop: stop,
572572
Close: inst.closed,
573573
Debug: inst.debug,
574574
Scale: inst.timeouts.Scale,

‎vm/bhyve/bhyve.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package bhyve
55

66
import (
7+
"context"
78
"fmt"
89
"io"
910
"os"
@@ -324,7 +325,7 @@ func (inst *instance) Copy(hostSrc string) (string, error) {
324325
return vmDst, nil
325326
}
326327

327-
func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
328+
func (inst *instance) Run(ctx context.Context, command string) (
328329
<-chan []byte, <-chan error, error) {
329330
rpipe, wpipe, err := osutil.LongPipe()
330331
if err != nil {
@@ -360,9 +361,7 @@ func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command strin
360361

361362
go func() {
362363
select {
363-
case <-time.After(timeout):
364-
signal(vmimpl.ErrTimeout)
365-
case <-stop:
364+
case <-ctx.Done():
366365
signal(vmimpl.ErrTimeout)
367366
case err := <-inst.merger.Err:
368367
cmd.Process.Kill()

‎vm/cuttlefish/cuttlefish.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
package cuttlefish
1212

1313
import (
14+
"context"
1415
"fmt"
1516
"os/exec"
1617
"path/filepath"
@@ -167,9 +168,9 @@ func (inst *instance) Close() error {
167168
return inst.gceInst.Close()
168169
}
169170

170-
func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
171+
func (inst *instance) Run(ctx context.Context, command string) (
171172
<-chan []byte, <-chan error, error) {
172-
return inst.gceInst.Run(timeout, stop, fmt.Sprintf("adb shell 'cd %s; %s'", deviceRoot, command))
173+
return inst.gceInst.Run(ctx, fmt.Sprintf("adb shell 'cd %s; %s'", deviceRoot, command))
173174
}
174175

175176
func (inst *instance) Diagnose(rep *report.Report) ([]byte, bool) {

‎vm/gce/gce.go‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ func (inst *instance) Copy(hostSrc string) (string, error) {
271271
return vmDst, nil
272272
}
273273

274-
func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
274+
func (inst *instance) Run(ctx context.Context, command string) (
275275
<-chan []byte, <-chan error, error) {
276276
conRpipe, conWpipe, err := osutil.LongPipe()
277277
if err != nil {
@@ -340,9 +340,8 @@ func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command strin
340340
sshWpipe.Close()
341341
merger.Add("ssh", sshRpipe)
342342

343-
return vmimpl.Multiplex(ssh, merger, timeout, vmimpl.MultiplexConfig{
343+
return vmimpl.Multiplex(ctx, ssh, merger, vmimpl.MultiplexConfig{
344344
Console: vmimpl.CmdCloser{Cmd: con},
345-
Stop: stop,
346345
Close: inst.closed,
347346
Debug: inst.debug,
348347
Scale: inst.timeouts.Scale,

‎vm/gvisor/gvisor.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ package gvisor
77

88
import (
99
"bytes"
10+
"context"
1011
"fmt"
1112
"io"
1213
"net"
@@ -286,7 +287,7 @@ func (inst *instance) Copy(hostSrc string) (string, error) {
286287
return filepath.Join("/", fname), nil
287288
}
288289

289-
func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
290+
func (inst *instance) Run(ctx context.Context, command string) (
290291
<-chan []byte, <-chan error, error) {
291292
args := []string{"exec", "-user=0:0"}
292293
for _, c := range sandboxCaps {
@@ -327,9 +328,7 @@ func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command strin
327328

328329
go func() {
329330
select {
330-
case <-time.After(timeout):
331-
signal(vmimpl.ErrTimeout)
332-
case <-stop:
331+
case <-ctx.Done():
333332
signal(vmimpl.ErrTimeout)
334333
case err := <-inst.merger.Err:
335334
cmd.Process.Kill()

0 commit comments

Comments
 (0)