From 5a7005591465dd7ea34834918299fac3bfe508f6 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Thu, 4 Jul 2024 18:50:36 +0200
Subject: pkg/vminfo: accept context.Context

The object enables a graceful shutdown of machine checks.
---
 pkg/rpcserver/local.go     | 2 +-
 pkg/rpcserver/rpcserver.go | 7 ++++---
 pkg/vminfo/linux_test.go   | 3 ++-
 pkg/vminfo/vminfo.go       | 3 +--
 pkg/vminfo/vminfo_test.go  | 5 +++--
 5 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/pkg/rpcserver/local.go b/pkg/rpcserver/local.go
index 1831259f2..27c0aac4a 100644
--- a/pkg/rpcserver/local.go
+++ b/pkg/rpcserver/local.go
@@ -50,7 +50,7 @@ func RunLocal(cfg *LocalConfig) error {
 		cfg:       cfg,
 		setupDone: make(chan bool),
 	}
-	serv, err := newImpl(&cfg.Config, ctx)
+	serv, err := newImpl(cfg.Context, &cfg.Config, ctx)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/rpcserver/rpcserver.go b/pkg/rpcserver/rpcserver.go
index d6a94094b..81563a077 100644
--- a/pkg/rpcserver/rpcserver.go
+++ b/pkg/rpcserver/rpcserver.go
@@ -5,6 +5,7 @@ package rpcserver
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"maps"
@@ -105,7 +106,7 @@ func New(cfg *mgrconfig.Config, mgr Manager, debug bool) (*Server, error) {
 	if !cfg.Experimental.RemoteCover {
 		features &= ^flatrpc.FeatureExtraCoverage
 	}
-	return newImpl(&Config{
+	return newImpl(context.Background(), &Config{
 		Config: vminfo.Config{
 			Target:     cfg.Target,
 			Features:   features,
@@ -130,9 +131,9 @@ func New(cfg *mgrconfig.Config, mgr Manager, debug bool) (*Server, error) {
 	}, mgr)
 }
 
-func newImpl(cfg *Config, mgr Manager) (*Server, error) {
+func newImpl(ctx context.Context, cfg *Config, mgr Manager) (*Server, error) {
 	cfg.Procs = min(cfg.Procs, prog.MaxPids)
-	checker := vminfo.New(&cfg.Config)
+	checker := vminfo.New(ctx, &cfg.Config)
 	baseSource := queue.DynamicSource(checker)
 	// Note that we use VMArch, rather than Arch. We need the kernel address ranges and bitness.
 	sysTarget := targets.Get(cfg.Target.OS, cfg.VMArch)
diff --git a/pkg/vminfo/linux_test.go b/pkg/vminfo/linux_test.go
index 8f956e5cf..c026eb478 100644
--- a/pkg/vminfo/linux_test.go
+++ b/pkg/vminfo/linux_test.go
@@ -6,6 +6,7 @@ package vminfo
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"fmt"
 	"os"
 	"runtime"
@@ -20,7 +21,7 @@ import (
 
 func TestLinuxSyscalls(t *testing.T) {
 	cfg := testConfig(t, targets.Linux, targets.AMD64)
-	checker := New(cfg)
+	checker := New(context.Background(), cfg)
 	filesystems := []string{
 		// Without sysfs, the checks would also disable mount().
 		"", "sysfs", "ext4", "binder", "",
diff --git a/pkg/vminfo/vminfo.go b/pkg/vminfo/vminfo.go
index b65baaac8..eb012d1de 100644
--- a/pkg/vminfo/vminfo.go
+++ b/pkg/vminfo/vminfo.go
@@ -47,7 +47,7 @@ type Config struct {
 	SandboxArg int64
 }
 
-func New(cfg *Config) *Checker {
+func New(ctx context.Context, cfg *Config) *Checker {
 	var impl checker
 	switch cfg.Target.OS {
 	case targets.Linux:
@@ -59,7 +59,6 @@ func New(cfg *Config) *Checker {
 	default:
 		impl = new(stub)
 	}
-	ctx := context.Background()
 	executor := queue.Plain()
 	return &Checker{
 		checker:      impl,
diff --git a/pkg/vminfo/vminfo_test.go b/pkg/vminfo/vminfo_test.go
index 2be23ca66..049456e19 100644
--- a/pkg/vminfo/vminfo_test.go
+++ b/pkg/vminfo/vminfo_test.go
@@ -4,6 +4,7 @@
 package vminfo
 
 import (
+	"context"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -54,7 +55,7 @@ func TestSyscalls(t *testing.T) {
 			t.Run(target.OS+"/"+target.Arch, func(t *testing.T) {
 				t.Parallel()
 				cfg := testConfig(t, target.OS, target.Arch)
-				checker := New(cfg)
+				checker := New(context.Background(), cfg)
 				stop := make(chan struct{})
 				go createSuccessfulResults(checker, stop)
 				enabled, disabled, _, err := checker.Run(nil, allFeatures())
@@ -119,7 +120,7 @@ func createSuccessfulResults(source queue.Source, stop chan struct{}) {
 
 func hostChecker(t *testing.T) (*Checker, []*flatrpc.FileInfo) {
 	cfg := testConfig(t, runtime.GOOS, runtime.GOARCH)
-	checker := New(cfg)
+	checker := New(context.Background(), cfg)
 	files := readFiles(checker.RequiredFiles())
 	return checker, files
 }
-- 
cgit mrf-deployment