ROX-15834: Add context to nodes.Analyze to allow quitting early (#1118

)

benchmarks/analyzeNode/analyze_node_test.go

@@ -1,6 +1,7 @@
 package detectconent
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"runtime"
@@ -25,7 +26,7 @@ func runBenchmarkAnalyzeNode(b *testing.B, pathName string) {
 	runtime.GC()
 
 	for i := 0; i < b.N; i++ {
-		node.Analyze("testNode", pathName, node.AnalyzeOpts{UncertifiedRHEL: false, IsRHCOSRequired: false})
+		node.Analyze(context.Background(), "testNode", pathName, node.AnalyzeOpts{UncertifiedRHEL: false, IsRHCOSRequired: false})
 	}
 	// Memory measuring command: go test -bench=foo -benchmem
 

go.mod

@@ -44,6 +44,7 @@ require (
 	github.com/stackrox/rox v0.0.0-20210914215712-9ac265932e28
 	github.com/stretchr/testify v1.8.2
 	go.etcd.io/bbolt v1.3.7
+	go.uber.org/goleak v1.2.0
 	go.uber.org/ratelimit v0.2.0
 	golang.org/x/exp v0.0.0-20220823124025-807a23277127
 	golang.org/x/sys v0.6.0

go.sum

@@ -960,6 +960,8 @@ go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
 go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
+go.uber.org/goleak v1.2.0 h1:xqgm/S+aQvhWFTtR0XK3Jvg7z8kGV8P4X14IzwN3Eqk=
+go.uber.org/goleak v1.2.0/go.mod h1:XJYK+MuIchqpmGmUSAzotztawfKvYLUIgg7guXrwVUo=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU=
 go.uber.org/multierr v1.7.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak=

pkg/analyzer/nodes/node.go

@@ -1,6 +1,7 @@
 package nodes
 
 import (
+	"context"
 	"io/fs"
 	"os"
 	"path/filepath"
@@ -65,11 +66,11 @@ type AnalyzeOpts struct {
 }
 
 // Analyze performs analysis of node's hosts filesystem and return the detected components.
-func Analyze(nodeName, rootFSdir string, opts AnalyzeOpts) (*Components, error) {
+func Analyze(ctx context.Context, nodeName, rootFSdir string, opts AnalyzeOpts) (*Components, error) {
 	// Currently, the node analyzer can only identify operating system components
 	// without active vulnerability, so we use the OS matcher.
 	matcher := requiredfilenames.SingletonOSMatcher()
-	files, err := extractFilesFromDirectory(rootFSdir, matcher)
+	files, err := extractFilesFromDirectory(ctx, rootFSdir, matcher)
 	if err != nil {
 		return nil, err
 	}
@@ -91,7 +92,7 @@ func Analyze(nodeName, rootFSdir string, opts AnalyzeOpts) (*Components, error)
 
 // extractFilesFromDirectory extracts files from the specified root directory
 // using a file matcher.
-func extractFilesFromDirectory(root string, matcher matcher.PrefixMatcher) (*filesMap, error) {
+func extractFilesFromDirectory(ctx context.Context, root string, matcher matcher.PrefixMatcher) (*filesMap, error) {
 	absRoot, err := filepath.Abs(root)
 	if err != nil {
 		return nil, errors.Wrapf(err, "invalid root path %q", root)
@@ -103,22 +104,51 @@ func extractFilesFromDirectory(root string, matcher matcher.PrefixMatcher) (*fil
 	m := metrics.FileExtractionMetrics{}
 	// TODO(ROX-13771): Use `range matcher.GetCommonPrefixDirs()` again after fixing.
 	for _, dir := range []string{"etc/", "usr/share/rpm", "var/lib/rpm", "usr/share/buildinfo"} {
-		if err := n.addFiles(filepath.FromSlash(dir), matcher, &m); err != nil {
-			return nil, errors.Wrapf(err, "failed to match filesMap at %q (at %q)", dir, n.root)
+		if err := n.addFiles(ctx, filepath.FromSlash(dir), matcher, &m); err != nil {
+			if errors.Is(err, context.DeadlineExceeded) {
+				err = errors.Wrapf(err, "operation took longer than expected and was interrupted")
+			}
+			return nil, errors.Wrapf(err, "failed to extract files from %q at root %q", dir, n.root)
 		}
 	}
 	m.Emit()
 	return n, nil
 }
 
+func walkDirWithContext(ctx context.Context, dir string, fn fs.WalkDirFunc) error {
+	// Buffering so we can give up on reading the channel and not block the
+	// child goroutine.
+	errC := make(chan error, 1)
+	go func() {
+		defer close(errC)
+		errC <- filepath.WalkDir(dir, func(path string, info fs.DirEntry, err error) error {
+			if err = fn(path, info, err); err != nil {
+				return err
+			}
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+				return nil
+			}
+		})
+	}()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-errC:
+		return err
+	}
+}
+
 // addFiles searches the directory for files using the matcher and adds them to the file map.
 // TODO(ROX-14050): Improve handling of symlinks - if possible, follow instead of ignoring them
-func (n *filesMap) addFiles(dir string, matcher matcher.Matcher, m *metrics.FileExtractionMetrics) error {
+func (n *filesMap) addFiles(ctx context.Context, dir string, matcher matcher.Matcher, m *metrics.FileExtractionMetrics) error {
 	logrus.WithFields(logrus.Fields{
 		"root":      n.root,
 		"directory": dir,
 	}).Info("add files from directory")
-	return filepath.WalkDir(filepath.Join(n.root, dir), func(fullPath string, entry fs.DirEntry, err error) error {
+	return walkDirWithContext(ctx, filepath.Join(n.root, dir), func(fullPath string, entry fs.DirEntry, err error) error {
 		if err != nil {
 			if filesIsNotAccessible(err) {
 				m.InaccessibleCount()

pkg/analyzer/nodes/node_test.go

@@ -1,18 +1,21 @@
 package nodes
 
 import (
+	"context"
 	"io"
 	"io/fs"
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/pkg/errors"
 	"github.com/stackrox/scanner/pkg/analyzer"
 	"github.com/stackrox/scanner/pkg/analyzer/analyzertest"
 	"github.com/stackrox/scanner/singletons/requiredfilenames"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 )
 
 type matcherMock struct {
@@ -415,7 +418,7 @@ func Test_extractFilesFromDirectory(t *testing.T) {
 
 	for _, testcase := range testcases {
 		t.Run(testcase.name, func(t *testing.T) {
-			filesMap, err := extractFilesFromDirectory(testcase.root, requiredfilenames.SingletonOSMatcher())
+			filesMap, err := extractFilesFromDirectory(context.Background(), testcase.root, requiredfilenames.SingletonOSMatcher())
 			assert.NoError(t, err)
 			assert.NoError(t, filesMap.readError)
 
@@ -431,6 +434,47 @@ func Test_extractFilesFromDirectory(t *testing.T) {
 	}
 }
 
+func Test_extractFilesFromDirectory_Context(t *testing.T) {
+	defer goleak.VerifyNone(t)
+	testDataRoot := makeTestData(t)
+	testcases := []struct {
+		name            string
+		root            string
+		ctxDeadline     time.Duration
+		expectedErr     error
+		expectNilResult bool
+	}{
+		{
+			name:            "context cancellation should return early with error",
+			root:            testDataRoot,
+			ctxDeadline:     5 * time.Microsecond,
+			expectedErr:     context.DeadlineExceeded,
+			expectNilResult: true,
+		},
+		{
+			name:            "function should return before context deadline",
+			root:            testDataRoot,
+			ctxDeadline:     5 * time.Second,
+			expectedErr:     nil,
+			expectNilResult: false,
+		},
+	}
+
+	for _, testcase := range testcases {
+		t.Run(testcase.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), testcase.ctxDeadline)
+			filesMap, err := extractFilesFromDirectory(ctx, testcase.root, requiredfilenames.SingletonOSMatcher())
+			cancel()
+			assert.ErrorIs(t, err, testcase.expectedErr)
+			if testcase.expectNilResult {
+				assert.Nil(t, filesMap)
+			} else {
+				assert.NotNil(t, filesMap)
+			}
+		})
+	}
+}
+
 func makeTestData(t *testing.T) string {
 	mkdirs := func(file string) {
 		require.NoError(t, os.MkdirAll(filepath.Dir(file), 0755))