pkg/vminfo: gracefully handle context abortion

On context abortion, return a special error.
On the pkg/rpcserver side, recognize and process it.

2 files changed, 11 insertions, 2 deletions

diff --git a/pkg/vminfo/syscalls.go b/pkg/vminfo/syscalls.go
index 9793cfee8..f4a6919a5 100644
--- a/pkg/vminfo/syscalls.go
+++ b/pkg/vminfo/syscalls.go
@@ -115,7 +115,10 @@ func (ctx *checkContext) do(fileInfos []*flatrpc.FileInfo, featureInfos []*flatr
 	globs := make(map[string][]string)
 	for _, req := range globReqs {
 		res := req.Wait(ctx.ctx)
-		if res.Status != queue.Success {
+		if res.Err == queue.ErrRequestAborted {
+			// Don't return an error on context cancellation.
+			return nil, nil, nil, nil
+		} else if res.Status != queue.Success {
 			return nil, nil, nil, fmt.Errorf("failed to execute glob: %w (%v)\n%s\n%s",
 				res.Err, res.Status, req.GlobPattern, res.Output)
 		}
diff --git a/pkg/vminfo/vminfo.go b/pkg/vminfo/vminfo.go
index acacdc9d5..41fadf59a 100644
--- a/pkg/vminfo/vminfo.go
+++ b/pkg/vminfo/vminfo.go
@@ -101,10 +101,16 @@ func (checker *Checker) MachineInfo(fileInfos []*flatrpc.FileInfo) ([]*KernelMod
 	return modules, info.Bytes(), nil
 }
 
+var ErrAborted = errors.New("aborted through the context")
+
 func (checker *Checker) Run(ctx context.Context, files []*flatrpc.FileInfo, featureInfos []*flatrpc.FeatureInfo) (
 	map[*prog.Syscall]bool, map[*prog.Syscall]string, Features, error) {
 	cc := newCheckContext(ctx, checker.cfg, checker.checker, checker.executor)
-	return cc.do(files, featureInfos)
+	enabled, disabled, features, err := cc.do(files, featureInfos)
+	if ctx.Err() != nil {
+		return nil, nil, nil, ErrAborted
+	}
+	return enabled, disabled, features, err
 }
 
 // Implementation of the queue.Source interface.