Implement context cancellation for command based gatherers

internal/factsengine/gatherers/ascsers_cluster.go

@@ -75,18 +75,16 @@ func (g *AscsErsClusterGatherer) SetCache(cache *factscache.FactsCache) {
 }
 
 func (g *AscsErsClusterGatherer) Gather(
-	_ context.Context,
+	ctx context.Context,
 	factsRequests []entities.FactRequest,
 ) ([]entities.Fact, error) {
 	log.Infof("Starting %s facts gathering process", AscsErsClusterGathererName)
 	var cibdata cib.Root
 
-	ctx := context.Background()
-
 	content, err := factscache.GetOrUpdate(
 		g.cache,
 		CibAdminGathererCache,
-		memoizeCibAdmin,
+		makeMemoizeCibAdmin(ctx),
 		g.executor,
 	)
 

internal/factsengine/gatherers/ascsers_cluster_test.go

@@ -8,12 +8,14 @@ import (
 	"os"
 	"testing"
 
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 	sapcontrol "github.com/trento-project/agent/internal/core/sapsystem/sapcontrolapi"
 	sapControlMocks "github.com/trento-project/agent/internal/core/sapsystem/sapcontrolapi/mocks"
 	"github.com/trento-project/agent/internal/factsengine/factscache"
 	"github.com/trento-project/agent/internal/factsengine/gatherers"
 	"github.com/trento-project/agent/pkg/factsengine/entities"
+	"github.com/trento-project/agent/pkg/utils"
 	utilsMocks "github.com/trento-project/agent/pkg/utils/mocks"
 	"github.com/trento-project/agent/test/helpers"
 )
@@ -36,7 +38,7 @@ func (suite *AscsErsClusterTestSuite) SetupTest() {
 }
 
 func (suite *AscsErsClusterTestSuite) TestAscsErsClusterGatherCmdNotFound() {
-	suite.mockExecutor.On("Exec", "cibadmin", "--query", "--local").Return(
+	suite.mockExecutor.On("ExecContext", mock.Anything, "cibadmin", "--query", "--local").Return(
 		[]byte{}, errors.New("cibadmin not found"))
 
 	p := gatherers.NewAscsErsClusterGatherer(suite.mockExecutor, suite.webService, nil)
@@ -84,7 +86,7 @@ func (suite *AscsErsClusterTestSuite) TestAscsErsClusterGatherInvalidInstanceNam
 	lFile, _ := os.Open(helpers.GetFixturePath("gatherers/cibadmin_multisid_invalid.xml"))
 	content, _ := io.ReadAll(lFile)
 
-	suite.mockExecutor.On("Exec", "cibadmin", "--query", "--local").Return(
+	suite.mockExecutor.On("ExecContext", mock.Anything, "cibadmin", "--query", "--local").Return(
 		content, nil)
 
 	p := gatherers.NewAscsErsClusterGatherer(suite.mockExecutor, suite.webService, nil)
@@ -109,7 +111,7 @@ func (suite *AscsErsClusterTestSuite) TestAscsErsClusterGatherInvalidInstanceNum
 		helpers.GetFixturePath("gatherers/cibadmin_multisid_invalid_instance_number.xml"))
 	content, _ := io.ReadAll(lFile)
 
-	suite.mockExecutor.On("Exec", "cibadmin", "--query", "--local").Return(
+	suite.mockExecutor.On("ExecContext", mock.Anything, "cibadmin", "--query", "--local").Return(
 		content, nil)
 
 	p := gatherers.NewAscsErsClusterGatherer(suite.mockExecutor, suite.webService, nil)
@@ -135,7 +137,7 @@ func (suite *AscsErsClusterTestSuite) TestAscsErsClusterGather() {
 	lFile, _ := os.Open(helpers.GetFixturePath("gatherers/cibadmin_multisid.xml"))
 	content, _ := io.ReadAll(lFile)
 
-	suite.mockExecutor.On("Exec", "cibadmin", "--query", "--local").Return(
+	suite.mockExecutor.On("ExecContext", mock.Anything, "cibadmin", "--query", "--local").Return(
 		content, nil)
 
 	mockWebServicePRDASCS00 := new(sapControlMocks.WebService)
@@ -286,3 +288,22 @@ func (suite *AscsErsClusterTestSuite) TestAscsErsClusterGather() {
 	}
 	suite.ElementsMatch(expectedEntries, entries)
 }
+
+func (suite *AscsErsClusterTestSuite) TestAscsErsGathererContextCancelled() {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	c := gatherers.NewAscsErsClusterGatherer(utils.Executor{}, suite.webService, nil)
+	factRequests := []entities.FactRequest{
+		{
+			Name:     "ascsers",
+			Gatherer: "ascsers_cluster",
+			Argument: "",
+			CheckID:  "check1",
+		},
+	}
+	factResults, err := c.Gather(ctx, factRequests)
+
+	suite.Error(err)
+	suite.Empty(factResults)
+}

internal/factsengine/gatherers/cibadmin.go

@@ -47,18 +47,21 @@ func (g *CibAdminGatherer) SetCache(cache *factscache.FactsCache) {
 	g.cache = cache
 }
 
-func memoizeCibAdmin(args ...interface{}) (interface{}, error) {
-	executor, ok := args[0].(utils.CommandExecutor)
-	if !ok {
-		return nil, ImplementationError.Wrap("error using memoizeCibAdmin. executor must be 1st argument")
+func makeMemoizeCibAdmin(ctx context.Context) func(...interface{}) (interface{}, error) {
+	return func(args ...interface{}) (interface{}, error) {
+		executor, ok := args[0].(utils.CommandExecutor)
+		if !ok {
+			return nil, ImplementationError.Wrap("error using memoizeCibAdmin. executor must be 1st argument")
+		}
+		return executor.ExecContext(ctx, "cibadmin", "--query", "--local")
 	}
-	return executor.Exec("cibadmin", "--query", "--local")
+
 }
 
-func (g *CibAdminGatherer) Gather(_ context.Context, factsRequests []entities.FactRequest) ([]entities.Fact, error) {
+func (g *CibAdminGatherer) Gather(ctx context.Context, factsRequests []entities.FactRequest) ([]entities.Fact, error) {
 	log.Infof("Starting %s facts gathering process", CibAdminGathererName)
 
-	content, err := factscache.GetOrUpdate(g.cache, CibAdminGathererCache, memoizeCibAdmin, g.executor)
+	content, err := factscache.GetOrUpdate(g.cache, CibAdminGathererCache, makeMemoizeCibAdmin(ctx), g.executor)
 
 	if err != nil {
 		return nil, CibAdminCommandError.Wrap(err.Error())

internal/factsengine/gatherers/cibadmin_test.go

@@ -7,10 +7,12 @@ import (
 	"os"
 	"testing"
 
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 	"github.com/trento-project/agent/internal/factsengine/factscache"
 	"github.com/trento-project/agent/internal/factsengine/gatherers"
 	"github.com/trento-project/agent/pkg/factsengine/entities"
+	"github.com/trento-project/agent/pkg/utils"
 	utilsMocks "github.com/trento-project/agent/pkg/utils/mocks"
 	"github.com/trento-project/agent/test/helpers"
 )
@@ -37,7 +39,7 @@ func (suite *CibAdminTestSuite) SetupTest() {
 }
 
 func (suite *CibAdminTestSuite) TestCibAdminGatherCmdNotFound() {
-	suite.mockExecutor.On("Exec", "cibadmin", "--query", "--local").Return(
+	suite.mockExecutor.On("ExecContext", mock.Anything, "cibadmin", "--query", "--local").Return(
 		suite.cibAdminOutput, errors.New("cibadmin not found"))
 
 	p := gatherers.NewCibAdminGatherer(suite.mockExecutor, nil)
@@ -58,7 +60,7 @@ func (suite *CibAdminTestSuite) TestCibAdminGatherCmdNotFound() {
 }
 
 func (suite *CibAdminTestSuite) TestCibAdminInvalidXML() {
-	suite.mockExecutor.On("Exec", "cibadmin", "--query", "--local").Return(
+	suite.mockExecutor.On("ExecContext", mock.Anything, "cibadmin", "--query", "--local").Return(
 		[]byte("invalid"), nil)
 
 	p := gatherers.NewCibAdminGatherer(suite.mockExecutor, nil)
@@ -79,7 +81,7 @@ func (suite *CibAdminTestSuite) TestCibAdminInvalidXML() {
 }
 
 func (suite *CibAdminTestSuite) TestCibAdminGather() {
-	suite.mockExecutor.On("Exec", "cibadmin", "--query", "--local").Return(
+	suite.mockExecutor.On("ExecContext", mock.Anything, "cibadmin", "--query", "--local").Return(
 		suite.cibAdminOutput, nil)
 
 	p := gatherers.NewCibAdminGatherer(suite.mockExecutor, nil)
@@ -210,7 +212,7 @@ func (suite *CibAdminTestSuite) TestCibAdminGather() {
 }
 
 func (suite *CibAdminTestSuite) TestCibAdminGatherWithCache() {
-	suite.mockExecutor.On("Exec", "cibadmin", "--query", "--local").
+	suite.mockExecutor.On("ExecContext", mock.Anything, "cibadmin", "--query", "--local").
 		Return(suite.cibAdminOutput, nil).
 		Once()
 
@@ -269,3 +271,25 @@ func (suite *CibAdminTestSuite) TestCibAdminGatherCacheCastingError() {
 	suite.EqualError(err, "fact gathering error: cibadmin-decoding-error - "+
 		"error decoding cibadmin output: error casting the command output")
 }
+
+func (suite *CibAdminTestSuite) TestCibAdminGatherWithContextCancelled() {
+
+	// Create a cancelled context
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	p := gatherers.NewCibAdminGatherer(utils.Executor{}, nil)
+	factRequests := []entities.FactRequest{
+		{
+			Name:     "cib",
+			Gatherer: "cibadmin",
+			Argument: "cib",
+			CheckID:  "check1",
+		},
+	}
+
+	factResults, err := p.Gather(ctx, factRequests)
+
+	suite.Error(err)
+	suite.Empty(factResults)
+}

internal/factsengine/gatherers/corosynccmapctl_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/stretchr/testify/suite"
 	"github.com/trento-project/agent/internal/factsengine/gatherers"
 	"github.com/trento-project/agent/pkg/factsengine/entities"
+	"github.com/trento-project/agent/pkg/utils"
 	utilsMocks "github.com/trento-project/agent/pkg/utils/mocks"
 	"github.com/trento-project/agent/test/helpers"
 )
@@ -209,3 +210,21 @@ func (suite *CorosyncCmapctlTestSuite) TestCorosyncCmapctlGatherer() {
 	suite.NoError(err)
 	suite.ElementsMatch(expectedResults, factResults)
 }
+
+func (suite *CorosyncCmapctlTestSuite) TestCorosyncCmapctlGathererContextCancelled() {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	c := gatherers.NewCorosyncCmapctlGatherer(utils.Executor{})
+	factRequests := []entities.FactRequest{
+		{
+			Name:     "madeup_fact",
+			Gatherer: "corosync-cmapctl",
+			Argument: "madeup.fact",
+		},
+	}
+	factResults, err := c.Gather(ctx, factRequests)
+
+	suite.Error(err)
+	suite.Empty(factResults)
+}

internal/factsengine/gatherers/dispwork.go

@@ -67,7 +67,7 @@ func NewDispWorkGatherer(fs afero.Fs, executor utils.CommandExecutor) *DispWorkG
 	}
 }
 
-func (g *DispWorkGatherer) Gather(_ context.Context, factsRequests []entities.FactRequest) ([]entities.Fact, error) {
+func (g *DispWorkGatherer) Gather(ctx context.Context, factsRequests []entities.FactRequest) ([]entities.Fact, error) {
 	facts := []entities.Fact{}
 	log.Infof("Starting %s facts gathering process", DispWorkGathererName)
 
@@ -82,8 +82,11 @@ func (g *DispWorkGatherer) Gather(_ context.Context, factsRequests []entities.Fa
 		sid := filepath.Base(systemPath)
 		sapUser := fmt.Sprintf("%sadm", strings.ToLower(sid))
 
-		dispWorkOutput, err := g.executor.Exec("su", "-", sapUser, "-c", "\"disp+work\"")
-		if err != nil {
+		dispWorkOutput, err := g.executor.ExecContext(ctx, "su", "-", sapUser, "-c", "\"disp+work\"")
+		switch {
+		case ctx.Err() != nil:
+			return nil, ctx.Err()
+		case err != nil:
 			gatheringError := DispWorkCommandError.Wrap(err.Error())
 			log.Error(gatheringError)
 			dispWorkMap[sid] = dispWorkData{} // fill with empty data

internal/factsengine/gatherers/dispwork_test.go

@@ -8,9 +8,11 @@ import (
 	"testing"
 
 	"github.com/spf13/afero"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 	"github.com/trento-project/agent/internal/factsengine/gatherers"
 	"github.com/trento-project/agent/pkg/factsengine/entities"
+	"github.com/trento-project/agent/pkg/utils"
 	utilsMocks "github.com/trento-project/agent/pkg/utils/mocks"
 	"github.com/trento-project/agent/test/helpers"
 )
@@ -48,13 +50,13 @@ func (suite *DispWorkGathererTestSuite) TestDispWorkGatheringSuccess() {
 	unsortedOutputFile, _ := os.Open(helpers.GetFixturePath("gatherers/dispwork-unsorted.output"))
 	unsortedOutput, _ := io.ReadAll(unsortedOutputFile)
 	suite.mockExecutor.
-		On("Exec", "su", "-", "prdadm", "-c", "\"disp+work\"").
+		On("ExecContext", mock.Anything, "su", "-", "prdadm", "-c", "\"disp+work\"").
 		Return(validOutput, nil).
-		On("Exec", "su", "-", "qasadm", "-c", "\"disp+work\"").
+		On("ExecContext", mock.Anything, "su", "-", "qasadm", "-c", "\"disp+work\"").
 		Return(partialOutput, nil).
-		On("Exec", "su", "-", "qa2adm", "-c", "\"disp+work\"").
+		On("ExecContext", mock.Anything, "su", "-", "qa2adm", "-c", "\"disp+work\"").
 		Return(unsortedOutput, nil).
-		On("Exec", "su", "-", "devadm", "-c", "\"disp+work\"").
+		On("ExecContext", mock.Anything, "su", "-", "devadm", "-c", "\"disp+work\"").
 		Return(nil, errors.New("some error"))
 
 	g := gatherers.NewDispWorkGatherer(suite.fs, suite.mockExecutor)
@@ -132,3 +134,21 @@ func (suite *DispWorkGathererTestSuite) TestDispWorkGatheringEmptyFileSystem() {
 	suite.NoError(err)
 	suite.EqualValues(expectedResults, result)
 }
+
+func (suite *DispWorkGathererTestSuite) TestDispWorkGathererContextCancelled() {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	c := gatherers.NewDispWorkGatherer(suite.fs, utils.Executor{})
+	factRequests := []entities.FactRequest{
+		{
+			Name:     "dispwork",
+			CheckID:  "check1",
+			Gatherer: "disp+work",
+		},
+	}
+	factResults, err := c.Gather(ctx, factRequests)
+
+	suite.Error(err)
+	suite.Empty(factResults)
+}

internal/factsengine/gatherers/mountinfo.go

@@ -66,7 +66,7 @@ func NewMountInfoGatherer(mInfo MountParserInterface, executor utils.CommandExec
 	return &MountInfoGatherer{mInfo: mInfo, executor: executor}
 }
 
-func (g *MountInfoGatherer) Gather(_ context.Context, factsRequests []entities.FactRequest) ([]entities.Fact, error) {
+func (g *MountInfoGatherer) Gather(ctx context.Context, factsRequests []entities.FactRequest) ([]entities.Fact, error) {
 	facts := []entities.Fact{}
 	log.Infof("Starting %s facts gathering process", MountInfoGathererName)
 	mounts, err := g.mInfo.GetMounts(nil)
@@ -92,7 +92,8 @@ func (g *MountInfoGatherer) Gather(_ context.Context, factsRequests []entities.F
 					Options:    mount.Options,
 				}
 
-				if blkidOuptut, err := g.executor.Exec("blkid", foundMountInfoResult.Source, "-o", "export"); err != nil {
+				blkidOuptut, err := g.executor.ExecContext(ctx, "blkid", foundMountInfoResult.Source, "-o", "export")
+				if err != nil {
 					log.Warnf("blkid command failed for source %s: %s", foundMountInfoResult.Source, err)
 				} else if fields, err := envparse.Parse(strings.NewReader(string(blkidOuptut))); err != nil {
 					log.Warnf("error parsing the blkid output: %s", err)

internal/factsengine/gatherers/mountinfo_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/trento-project/agent/internal/factsengine/gatherers"
 	"github.com/trento-project/agent/internal/factsengine/gatherers/mocks"
 	"github.com/trento-project/agent/pkg/factsengine/entities"
+	"github.com/trento-project/agent/pkg/utils"
 	utilsMocks "github.com/trento-project/agent/pkg/utils/mocks"
 )
 
@@ -66,11 +67,11 @@ TYPE=xfs
 `)
 
 	suite.mockExecutor.
-		On("Exec", "blkid", "10.1.1.10:/sapmnt", "-o", "export").
+		On("ExecContext", mock.Anything, "blkid", "10.1.1.10:/sapmnt", "-o", "export").
 		Return(nil, fmt.Errorf("blkid error")).
-		On("Exec", "blkid", "/dev/mapper/vg_hana-lv_data", "-o", "export").
+		On("ExecContext", mock.Anything, "blkid", "/dev/mapper/vg_hana-lv_data", "-o", "export").
 		Return(blkidOutput, nil).
-		On("Exec", "blkid", "/dev/mapper/vg_hana-lv_log", "-o", "export").
+		On("ExecContext", mock.Anything, "blkid", "/dev/mapper/vg_hana-lv_log", "-o", "export").
 		Return(blkidOutputNoUUID, nil)
 
 	requestedFacts := []entities.FactRequest{
@@ -211,3 +212,21 @@ func (suite *MountInfoTestSuite) TestMountInfoParsingError() {
 	suite.EqualError(err, "fact gathering error: mount-info-parsing-error - "+
 		"error parsing mount information: some error")
 }
+
+func (suite *MountInfoTestSuite) TestMountInfoParsingGathererContextCancelled() {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	c := gatherers.NewSapHostCtrlGatherer(utils.Executor{})
+	factRequests := []entities.FactRequest{
+		{Name: "shared",
+			Gatherer: "mount_info",
+			CheckID:  "check1",
+			Argument: "/sapmnt",
+		},
+	}
+	factResults, err := c.Gather(ctx, factRequests)
+
+	suite.Error(err)
+	suite.Empty(factResults)
+}