vm: add context to Pool.Create()

Enable external abortion of the instance creation process. This is
especially useful for the qemu case where we retry the creation/boot up
to 1000 times, which can take significant time (e.g. it timeouts
syz-cluster pods on unstable kernels).

The context can be further propagated to WaitForSSH, but that requires
another quite significant vm/ refactoring.

17 files changed, 34 insertions, 31 deletions

diff --git a/vm/adb/adb.go b/vm/adb/adb.go
index 202f352c0..6b546678e 100644
--- a/vm/adb/adb.go
+++ b/vm/adb/adb.go
@@ -129,7 +129,7 @@ func (pool *Pool) Count() int {
 	return len(pool.cfg.Devices)
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	device, err := loadDevice(pool.cfg.Devices[index])
 	if err != nil {
 		return nil, err
diff --git a/vm/bhyve/bhyve.go b/vm/bhyve/bhyve.go
index 6cccf7a9f..7be5fdfb4 100644
--- a/vm/bhyve/bhyve.go
+++ b/vm/bhyve/bhyve.go
@@ -83,7 +83,7 @@ func (pool *Pool) Count() int {
 	return pool.cfg.Count
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	inst := &instance{
 		cfg:   pool.cfg,
 		debug: pool.env.Debug,
diff --git a/vm/cuttlefish/cuttlefish.go b/vm/cuttlefish/cuttlefish.go
index 55e56d2e5..dbb95d51d 100644
--- a/vm/cuttlefish/cuttlefish.go
+++ b/vm/cuttlefish/cuttlefish.go
@@ -66,8 +66,8 @@ func (pool *Pool) Count() int {
 	return pool.gcePool.Count()
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
-	gceInst, err := pool.gcePool.Create(workdir, index)
+func (pool *Pool) Create(ctx context.Context, workdir string, index int) (vmimpl.Instance, error) {
+	gceInst, err := pool.gcePool.Create(ctx, workdir, index)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create underlying gce instance: %w", err)
 	}
diff --git a/vm/dispatcher/pool.go b/vm/dispatcher/pool.go
index 08bc1ebfa..424724c25 100644
--- a/vm/dispatcher/pool.go
+++ b/vm/dispatcher/pool.go
@@ -19,7 +19,7 @@ type Instance interface {
 
 type UpdateInfo func(cb func(info *Info))
 type Runner[T Instance] func(ctx context.Context, inst T, updInfo UpdateInfo)
-type CreateInstance[T Instance] func(int) (T, error)
+type CreateInstance[T Instance] func(context.Context, int) (T, error)
 
 // Pool[T] provides the functionality of a generic pool of instances.
 // The instance is assumed to boot, be controlled by one Runner and then be re-created.
@@ -125,7 +125,7 @@ func (p *Pool[T]) runInstance(ctx context.Context, inst *poolInstance[T]) {
 	inst.status(StateBooting)
 	defer inst.status(StateOffline)
 
-	obj, err := p.creator(inst.idx)
+	obj, err := p.creator(ctx, inst.idx)
 	if err != nil {
 		p.reportBootError(ctx, err)
 		return
diff --git a/vm/dispatcher/pool_test.go b/vm/dispatcher/pool_test.go
index 738ddfe27..9ac11d202 100644
--- a/vm/dispatcher/pool_test.go
+++ b/vm/dispatcher/pool_test.go
@@ -21,7 +21,7 @@ func TestPoolDefault(t *testing.T) {
 
 	mgr := NewPool[*testInstance](
 		count,
-		func(idx int) (*testInstance, error) {
+		func(_ context.Context, idx int) (*testInstance, error) {
 			pool[idx].reset()
 			return &pool[idx], nil
 		},
@@ -62,7 +62,7 @@ func TestPoolSplit(t *testing.T) {
 
 	mgr := NewPool[*testInstance](
 		count,
-		func(idx int) (*testInstance, error) {
+		func(_ context.Context, idx int) (*testInstance, error) {
 			pool[idx].reset()
 			return &pool[idx], nil
 		},
@@ -131,7 +131,7 @@ func TestPoolStress(t *testing.T) {
 	// The test to aid the race detector.
 	mgr := NewPool[*nilInstance](
 		10,
-		func(idx int) (*nilInstance, error) {
+		func(_ context.Context, idx int) (*nilInstance, error) {
 			return &nilInstance{}, nil
 		},
 		func(ctx context.Context, _ *nilInstance, _ UpdateInfo) {
@@ -166,7 +166,7 @@ func TestPoolNewDefault(t *testing.T) {
 	// The test to aid the race detector.
 	mgr := NewPool[*nilInstance](
 		10,
-		func(idx int) (*nilInstance, error) {
+		func(_ context.Context, idx int) (*nilInstance, error) {
 			return &nilInstance{}, nil
 		},
 		func(ctx context.Context, _ *nilInstance, _ UpdateInfo) {
@@ -205,7 +205,7 @@ func TestPoolNewDefault(t *testing.T) {
 func TestPoolPause(t *testing.T) {
 	mgr := NewPool[*nilInstance](
 		10,
-		func(idx int) (*nilInstance, error) {
+		func(_ context.Context, idx int) (*nilInstance, error) {
 			return &nilInstance{}, nil
 		},
 		func(ctx context.Context, _ *nilInstance, _ UpdateInfo) {
@@ -242,7 +242,7 @@ func TestPoolCancelRun(t *testing.T) {
 	// The test to aid the race detector.
 	mgr := NewPool[*nilInstance](
 		10,
-		func(idx int) (*nilInstance, error) {
+		func(_ context.Context, idx int) (*nilInstance, error) {
 			return &nilInstance{}, nil
 		},
 		func(ctx context.Context, _ *nilInstance, _ UpdateInfo) {
@@ -293,7 +293,7 @@ func TestPoolBootErrors(t *testing.T) {
 
 	mgr := NewPool[*testInstance](
 		3,
-		func(idx int) (*testInstance, error) {
+		func(_ context.Context, idx int) (*testInstance, error) {
 			failCount.Add(1)
 			return nil, fmt.Errorf("boot error")
 		},
diff --git a/vm/gce/gce.go b/vm/gce/gce.go
index 1a887ba9f..a830a9889 100644
--- a/vm/gce/gce.go
+++ b/vm/gce/gce.go
@@ -176,7 +176,7 @@ func (pool *Pool) Count() int {
 	return pool.cfg.Count
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	name := fmt.Sprintf("%v-%v", pool.env.Name, index)
 	// Create SSH key for the instance.
 	gceKey := filepath.Join(workdir, "key")
diff --git a/vm/gvisor/gvisor.go b/vm/gvisor/gvisor.go
index 3f3fc6fa5..368b9a9fc 100644
--- a/vm/gvisor/gvisor.go
+++ b/vm/gvisor/gvisor.go
@@ -88,7 +88,7 @@ func (pool *Pool) Count() int {
 	return pool.cfg.Count
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	rootDir := filepath.Clean(filepath.Join(workdir, "..", "gvisor_root"))
 	imageDir := filepath.Join(workdir, "image")
 	bundleDir := filepath.Join(workdir, "bundle")
diff --git a/vm/isolated/isolated.go b/vm/isolated/isolated.go
index 6e57fef9c..769665ec3 100755
--- a/vm/isolated/isolated.go
+++ b/vm/isolated/isolated.go
@@ -92,7 +92,7 @@ func (pool *Pool) Count() int {
 	return len(pool.cfg.Targets)
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	targetAddr, targetPort, _ := splitTargetPort(pool.cfg.Targets[index])
 	inst := &instance{
 		cfg: pool.cfg,
diff --git a/vm/proxyapp/proxyappclient.go b/vm/proxyapp/proxyappclient.go
index dce7b623e..4ffa12b57 100644
--- a/vm/proxyapp/proxyappclient.go
+++ b/vm/proxyapp/proxyappclient.go
@@ -154,7 +154,7 @@ func (p *pool) Count() int {
 	return p.count
 }
 
-func (p *pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (p *pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	p.mu.Lock()
 	proxy := p.proxy
 	p.mu.Unlock()
diff --git a/vm/proxyapp/proxyappclient_test.go b/vm/proxyapp/proxyappclient_test.go
index 7053f2411..1bf7d0571 100644
--- a/vm/proxyapp/proxyappclient_test.go
+++ b/vm/proxyapp/proxyappclient_test.go
@@ -239,7 +239,7 @@ func TestPool_Create_Ok(t *testing.T) {
 		On("CreateInstance", mock.Anything, mock.Anything).
 		Return(nil)
 
-	inst, err := p.Create("workdir", 0)
+	inst, err := p.Create(t.Context(), "workdir", 0)
 	assert.NotNil(t, inst)
 	assert.Nil(t, err)
 }
@@ -257,7 +257,7 @@ func TestPool_Create_ProxyNilError(t *testing.T) {
 
 	p.(io.Closer).Close()
 
-	inst, err := p.Create("workdir", 0)
+	inst, err := p.Create(t.Context(), "workdir", 0)
 	assert.Nil(t, inst)
 	assert.NotNil(t, err)
 }
@@ -272,7 +272,7 @@ func TestPool_Create_OutOfPoolError(t *testing.T) {
 		}).
 		Return(fmt.Errorf("out of pool size"))
 
-	inst, err := p.Create("workdir", p.Count())
+	inst, err := p.Create(t.Context(), "workdir", p.Count())
 	assert.Nil(t, inst)
 	assert.NotNil(t, err)
 }
@@ -283,7 +283,7 @@ func TestPool_Create_ProxyFailure(t *testing.T) {
 		On("CreateInstance", mock.Anything, mock.Anything).
 		Return(fmt.Errorf("create instance failure"))
 
-	inst, err := p.Create("workdir", 0)
+	inst, err := p.Create(t.Context(), "workdir", 0)
 	assert.Nil(t, inst)
 	assert.NotNil(t, err)
 }
@@ -300,7 +300,7 @@ func createInstanceFixture(t *testing.T) (*mock.Mock, vmimpl.Instance) {
 		}).
 		Return(nil)
 
-	inst, err := p.Create("workdir", 0)
+	inst, err := p.Create(t.Context(), "workdir", 0)
 	assert.Nil(t, err)
 	assert.NotNil(t, inst)
 
diff --git a/vm/qemu/qemu.go b/vm/qemu/qemu.go
index 3bf546aa2..63598d4bb 100644
--- a/vm/qemu/qemu.go
+++ b/vm/qemu/qemu.go
@@ -324,7 +324,7 @@ func (pool *Pool) Count() int {
 	return pool.cfg.Count
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(ctx context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	sshkey := pool.env.SSHKey
 	sshuser := pool.env.SSHUser
 	if pool.env.Image == "9p" {
@@ -341,6 +341,9 @@ func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
 	}
 
 	for i := 0; ; i++ {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
 		inst, err := pool.ctor(workdir, sshkey, sshuser, index)
 		if err == nil {
 			return inst, nil
diff --git a/vm/starnix/starnix.go b/vm/starnix/starnix.go
index 665b13321..e6da96dc0 100644
--- a/vm/starnix/starnix.go
+++ b/vm/starnix/starnix.go
@@ -100,7 +100,7 @@ func (pool *Pool) Count() int {
 	return pool.count
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	inst := &instance{
 		fuchsiaDir: pool.env.KernelSrc,
 		ffxDir:     pool.ffxDir,
diff --git a/vm/vm.go b/vm/vm.go
index e155fd8e5..42d758517 100644
--- a/vm/vm.go
+++ b/vm/vm.go
@@ -157,7 +157,7 @@ func (pool *Pool) Count() int {
 	return pool.count
 }
 
-func (pool *Pool) Create(index int) (*Instance, error) {
+func (pool *Pool) Create(ctx context.Context, index int) (*Instance, error) {
 	if index < 0 || index >= pool.count {
 		return nil, fmt.Errorf("invalid VM index %v (count %v)", index, pool.count)
 	}
@@ -170,7 +170,7 @@ func (pool *Pool) Create(index int) (*Instance, error) {
 			return nil, err
 		}
 	}
-	impl, err := pool.impl.Create(workdir, index)
+	impl, err := pool.impl.Create(ctx, workdir, index)
 	if err != nil {
 		os.RemoveAll(workdir)
 		return nil, err
diff --git a/vm/vm_test.go b/vm/vm_test.go
index 446898bd6..23e3f2c7d 100644
--- a/vm/vm_test.go
+++ b/vm/vm_test.go
@@ -26,7 +26,7 @@ func (pool *testPool) Count() int {
 	return 1
 }
 
-func (pool *testPool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *testPool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	return &testInstance{
 		outc: make(chan []byte, 10),
 		errc: make(chan error, 1),
@@ -376,7 +376,7 @@ func makeLinuxAMD64Futex(t *testing.T) (*Instance, *report.Reporter) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	inst, err := pool.Create(0)
+	inst, err := pool.Create(t.Context(), 0)
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/vm/vmimpl/vmimpl.go b/vm/vmimpl/vmimpl.go
index 68e7030c3..480fb85b5 100644
--- a/vm/vmimpl/vmimpl.go
+++ b/vm/vmimpl/vmimpl.go
@@ -33,7 +33,7 @@ type Pool interface {
 	Count() int
 
 	// Create creates and boots a new VM instance.
-	Create(workdir string, index int) (Instance, error)
+	Create(ctx context.Context, workdir string, index int) (Instance, error)
 }
 
 // Instance represents a single VM.
diff --git a/vm/vmm/vmm.go b/vm/vmm/vmm.go
index 5190e32b3..884622fc8 100644
--- a/vm/vmm/vmm.go
+++ b/vm/vmm/vmm.go
@@ -93,7 +93,7 @@ func (pool *Pool) Count() int {
 	return pool.cfg.Count
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	var tee io.Writer
 	if pool.env.Debug {
 		tee = os.Stdout
diff --git a/vm/vmware/vmware.go b/vm/vmware/vmware.go
index e4379717b..d7c2a5eff 100644
--- a/vm/vmware/vmware.go
+++ b/vm/vmware/vmware.go
@@ -77,7 +77,7 @@ func (pool *Pool) Count() int {
 	return pool.cfg.Count
 }
 
-func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+func (pool *Pool) Create(_ context.Context, workdir string, index int) (vmimpl.Instance, error) {
 	createTime := strconv.FormatInt(time.Now().UnixNano(), 10)
 	vmx := filepath.Join(workdir, createTime, "syzkaller.vmx")
 	sshkey := pool.env.SSHKey