stackrox · md2119 · Nov 2, 2022 · Nov 2, 2022 · Nov 2, 2022
diff --git a/central/alert/datastore/datastore_impl.go b/central/alert/datastore/datastore_impl.go
@@ -23,6 +23,7 @@ import (
 	"github.com/stackrox/rox/pkg/env"
 	"github.com/stackrox/rox/pkg/errorhelpers"
 	"github.com/stackrox/rox/pkg/logging"
+	"github.com/stackrox/rox/pkg/postgres/pgutils"
 	"github.com/stackrox/rox/pkg/sac"
 	searchCommon "github.com/stackrox/rox/pkg/search"
 	"github.com/stackrox/rox/pkg/search/paginated"
@@ -448,8 +449,11 @@ func (ds *datastoreImpl) WalkAll(ctx context.Context, fn func(*storage.ListAlert
 		return sac.ErrResourceAccessDenied
 	}
 
-	return ds.storage.Walk(ctx, func(alert *storage.Alert) error {
-		listAlert := convert.AlertToListAlert(alert)
-		return fn(listAlert)
-	})
+	walkFn := func() error {
+		return ds.storage.Walk(ctx, func(alert *storage.Alert) error {
+			listAlert := convert.AlertToListAlert(alert)
+			return fn(listAlert)
+		})
+	}
+	return pgutils.RetryIfPostgres(walkFn)
 }
diff --git a/central/apitoken/datastore/datastore_impl.go b/central/apitoken/datastore/datastore_impl.go
@@ -7,6 +7,7 @@ import (
 	"github.com/stackrox/rox/central/role/resources"
 	v1 "github.com/stackrox/rox/generated/api/v1"
 	"github.com/stackrox/rox/generated/storage"
+	"github.com/stackrox/rox/pkg/postgres/pgutils"
 	"github.com/stackrox/rox/pkg/sac"
 	"github.com/stackrox/rox/pkg/sync"
 )
@@ -65,14 +66,17 @@ func (b *datastoreImpl) GetTokens(ctx context.Context, req *v1.GetAPITokensReque
 	defer b.Unlock()
 
 	var tokens []*storage.TokenMetadata
-	err := b.storage.Walk(ctx, func(token *storage.TokenMetadata) error {
-		if req.GetRevokedOneof() != nil && req.GetRevoked() != token.GetRevoked() {
+	walkFn := func() error {
+		tokens = tokens[:0]
+		return b.storage.Walk(ctx, func(token *storage.TokenMetadata) error {
+			if req.GetRevokedOneof() != nil && req.GetRevoked() != token.GetRevoked() {
+				return nil
+			}
+			tokens = append(tokens, token)
 			return nil
-		}
-		tokens = append(tokens, token)
-		return nil
-	})
-	if err != nil {
+		})
+	}
+	if err := pgutils.RetryIfPostgres(walkFn); err != nil {
 		return nil, err
 	}
 	return tokens, nil

diff --git a/central/cluster/datastore/datastore_impl.go b/central/cluster/datastore/datastore_impl.go
@@ -39,6 +39,7 @@ import (
 	"github.com/stackrox/rox/pkg/env"
 	"github.com/stackrox/rox/pkg/errox"
 	"github.com/stackrox/rox/pkg/images/defaults"
+	"github.com/stackrox/rox/pkg/postgres/pgutils"
 	"github.com/stackrox/rox/pkg/protoconv"
 	"github.com/stackrox/rox/pkg/sac"
 	pkgSearch "github.com/stackrox/rox/pkg/search"
@@ -156,21 +157,20 @@ func (ds *datastoreImpl) UpdateClusterStatus(ctx context.Context, id string, sta
 }
 
 func (ds *datastoreImpl) buildIndex(ctx context.Context) error {
-	var clusters []*storage.Cluster
-	err := ds.clusterStorage.Walk(ctx, func(cluster *storage.Cluster) error {
-		clusters = append(clusters, cluster)
-		return nil
-	})
+	clusters, err := ds.collectClusters(ctx)
 	if err != nil {
 		return err
 	}
 
 	clusterHealthStatuses := make(map[string]*storage.ClusterHealthStatus)
-	err = ds.clusterHealthStorage.Walk(ctx, func(healthInfo *storage.ClusterHealthStatus) error {
-		clusterHealthStatuses[healthInfo.Id] = healthInfo
-		return nil
-	})
-	if err != nil {
+	walkFn := func() error {
+		clusterHealthStatuses = make(map[string]*storage.ClusterHealthStatus)
+		return ds.clusterHealthStorage.Walk(ctx, func(healthInfo *storage.ClusterHealthStatus) error {
+			clusterHealthStatuses[healthInfo.Id] = healthInfo
+			return nil
+		})
+	}
+	if err := pgutils.RetryIfPostgres(walkFn); err != nil {
 		return err
 	}
 
@@ -188,14 +188,10 @@ func (ds *datastoreImpl) registerClusterForNetworkGraphExtSrcs() error {
 			sac.AccessModeScopeKeys(storage.Access_READ_ACCESS, storage.Access_READ_WRITE_ACCESS),
 			sac.ResourceScopeKeys(resources.Node, resources.NetworkGraph)))
 
-	var clusters []*storage.Cluster
-	if err := ds.clusterStorage.Walk(ctx, func(cluster *storage.Cluster) error {
-		clusters = append(clusters, cluster)
-		return nil
-	}); err != nil {
+	clusters, err := ds.collectClusters(ctx)
+	if err != nil {
 		return err
 	}
-
 	for _, cluster := range clusters {
 		ds.netEntityDataStore.RegisterCluster(ctx, cluster.GetId())
 	}
@@ -245,11 +241,7 @@ func (ds *datastoreImpl) GetClusters(ctx context.Context) ([]*storage.Cluster, e
 	if err != nil {
 		return nil, err
 	} else if ok {
-		var clusters []*storage.Cluster
-		err := ds.clusterStorage.Walk(ctx, func(cluster *storage.Cluster) error {
-			clusters = append(clusters, cluster)
-			return nil
-		})
+		clusters, err := ds.collectClusters(ctx)
 		if err != nil {
 			return nil, err
 		}
@@ -1047,3 +1039,18 @@ func configureFromHelmConfig(cluster *storage.Cluster, helmConfig *storage.Compl
 	cluster.TolerationsConfig = staticConfig.GetTolerationsConfig().Clone()
 	cluster.SlimCollector = staticConfig.GetSlimCollector()
 }
+
+func (ds *datastoreImpl) collectClusters(ctx context.Context) ([]*storage.Cluster, error) {
+	var clusters []*storage.Cluster
+	walkFn := func() error {
+		clusters = clusters[:0]
+		return ds.clusterStorage.Walk(ctx, func(cluster *storage.Cluster) error {
+			clusters = append(clusters, cluster)
+			return nil
+		})
+	}
+	if err := pgutils.RetryIfPostgres(walkFn); err != nil {
+		return nil, err
+	}
+	return clusters, nil
+}
diff --git a/central/compliance/data/repository.go b/central/compliance/data/repository.go
@@ -16,6 +16,7 @@ import (
 	"github.com/stackrox/rox/pkg/compliance/data"
 	"github.com/stackrox/rox/pkg/complianceoperator/api/v1alpha1"
 	"github.com/stackrox/rox/pkg/logging"
+	"github.com/stackrox/rox/pkg/postgres/pgutils"
 	"github.com/stackrox/rox/pkg/search"
 	"github.com/stackrox/rox/pkg/set"
 	"github.com/stackrox/rox/pkg/utils"
@@ -297,19 +298,22 @@ func (r *repository) init(ctx context.Context, domain framework.ComplianceDomain
 	}
 
 	r.complianceOperatorResults = make(map[string][]*storage.ComplianceOperatorCheckResult)
-	err = f.complianceOperatorResultStore.Walk(ctx, func(c *storage.ComplianceOperatorCheckResult) error {
-		if c.GetClusterId() != clusterID {
-			return nil
-		}
-		rule := c.Annotations[v1alpha1.RuleIDAnnotationKey]
-		if rule == "" {
-			log.Errorf("Expected rule annotation for %+v", c)
+	walkFn := func() error {
+		r.complianceOperatorResults = make(map[string][]*storage.ComplianceOperatorCheckResult)
+		return f.complianceOperatorResultStore.Walk(ctx, func(c *storage.ComplianceOperatorCheckResult) error {
+			if c.GetClusterId() != clusterID {
+				return nil
+			}
+			rule := c.Annotations[v1alpha1.RuleIDAnnotationKey]
+			if rule == "" {
+				log.Errorf("Expected rule annotation for %+v", c)
+				return nil
+			}
+			r.complianceOperatorResults[rule] = append(r.complianceOperatorResults[rule], c)
 			return nil
-		}
-		r.complianceOperatorResults[rule] = append(r.complianceOperatorResults[rule], c)
-		return nil
-	})
-	if err != nil {
+		})
+	}
+	if err := pgutils.RetryIfPostgres(walkFn); err != nil {
 		return err
 	}
 

diff --git a/central/complianceoperator/checkresults/datastore/datastore.go b/central/complianceoperator/checkresults/datastore/datastore.go
@@ -36,6 +36,7 @@ func (d *datastoreImpl) Walk(ctx context.Context, fn func(result *storage.Compli
 	} else if !ok {
 		return errors.Wrap(sac.ErrResourceAccessDenied, "compliance operator check results read")
 	}
+	// Retry in the caller
 	return d.store.Walk(ctx, fn)
 }
 

diff --git a/central/detection/lifecycle/manager_impl.go b/central/detection/lifecycle/manager_impl.go
@@ -26,6 +26,7 @@ import (
 	"github.com/stackrox/rox/pkg/env"
 	"github.com/stackrox/rox/pkg/expiringcache"
 	"github.com/stackrox/rox/pkg/policies"
+	"github.com/stackrox/rox/pkg/postgres/pgutils"
 	"github.com/stackrox/rox/pkg/process/filter"
 	processBaselinePkg "github.com/stackrox/rox/pkg/processbaseline"
 	"github.com/stackrox/rox/pkg/protoutils"
@@ -93,28 +94,29 @@ func (m *managerImpl) copyAndResetIndicatorQueue() map[string]*storage.ProcessIn
 
 func (m *managerImpl) buildIndicatorFilter() {
 	ctx := sac.WithAllAccess(context.Background())
-	var processesToRemove []string
-
 	deploymentIDs, err := m.deploymentDataStore.GetDeploymentIDs(ctx)
 	if err != nil {
 		utils.Should(errors.Wrap(err, "error getting deployment IDs"))
 		return
 	}
 
-	deploymentIDSet := set.NewStringSet(deploymentIDs...)
-
-	err = m.processesDataStore.WalkAll(ctx, func(pi *storage.ProcessIndicator) error {
-		if !deploymentIDSet.Contains(pi.GetDeploymentId()) {
-			// Don't remove as these processes will be removed by GC
-			// but don't add to the filter
+	var processesToRemove []string
+	walkFn := func() error {
+		deploymentIDSet := set.NewStringSet(deploymentIDs...)
+		processesToRemove = processesToRemove[:0]
+		return m.processesDataStore.WalkAll(ctx, func(pi *storage.ProcessIndicator) error {
+			if !deploymentIDSet.Contains(pi.GetDeploymentId()) {
+				// Don't remove as these processes will be removed by GC
+				// but don't add to the filter
+				return nil
+			}
+			if !m.processFilter.Add(pi) {
+				processesToRemove = append(processesToRemove, pi.GetId())
+			}
 			return nil
-		}
-		if !m.processFilter.Add(pi) {
-			processesToRemove = append(processesToRemove, pi.GetId())
-		}
-		return nil
-	})
-	if err != nil {
+		})
+	}
+	if err := pgutils.RetryIfPostgres(walkFn); err != nil {
 		utils.Should(errors.Wrap(err, "error building indicator filter"))
 	}
 

diff --git a/central/group/datastore/datastore_impl.go b/central/group/datastore/datastore_impl.go
@@ -8,6 +8,7 @@ import (
 	"github.com/stackrox/rox/central/role/resources"
 	"github.com/stackrox/rox/generated/storage"
 	"github.com/stackrox/rox/pkg/errox"
+	"github.com/stackrox/rox/pkg/postgres/pgutils"
 	"github.com/stackrox/rox/pkg/sac"
 	"github.com/stackrox/rox/pkg/sync"
 	"github.com/stackrox/rox/pkg/utils"
@@ -56,14 +57,19 @@ func (ds *dataStoreImpl) GetFiltered(ctx context.Context, filter func(*storage.G
 	}
 
 	var groups []*storage.Group
-	err := ds.storage.Walk(ctx, func(g *storage.Group) error {
-		if filter == nil || filter(g.GetProps()) {
-			groups = append(groups, g)
-		}
-		return nil
-	})
-
-	return groups, err
+	walkFn := func() error {
+		groups = groups[:0]
+		return ds.storage.Walk(ctx, func(g *storage.Group) error {
+			if filter == nil || filter(g.GetProps()) {
+				groups = append(groups, g)
+			}
+			return nil
+		})
+	}
+	if err := pgutils.RetryIfPostgres(walkFn); err != nil {
+		return nil, err
+	}
+	return groups, nil
 }
 
 // Walk is an optimization that allows to search through the datastore and find
@@ -78,16 +84,21 @@ func (ds *dataStoreImpl) Walk(ctx context.Context, authProviderID string, attrib
 	// Search through the datastore and find all groups that apply to a user within a single transaction.
 	toSearch := getPossibleGroupProperties(authProviderID, attributes)
 	var groups []*storage.Group
-	err := ds.storage.Walk(ctx, func(group *storage.Group) error {
-		for _, check := range toSearch {
-			if propertiesMatch(group.GetProps(), check) {
-				groups = append(groups, group)
+	walkFn := func() error {
+		groups = groups[:0]
+		return ds.storage.Walk(ctx, func(group *storage.Group) error {
+			for _, check := range toSearch {
+				if propertiesMatch(group.GetProps(), check) {
+					groups = append(groups, group)
+				}
 			}
-		}
-		return nil
-	})
-
-	return groups, err
+			return nil
+		})
+	}
+	if err := pgutils.RetryIfPostgres(walkFn); err != nil {
+		return nil, err
+	}
+	return groups, nil
 }
 
 func (ds *dataStoreImpl) Add(ctx context.Context, group *storage.Group) error {

diff --git a/central/integrationhealth/datastore/datastore_impl.go b/central/integrationhealth/datastore/datastore_impl.go
@@ -7,6 +7,7 @@ import (
 	"github.com/stackrox/rox/central/integrationhealth/store"
 	"github.com/stackrox/rox/central/role/resources"
 	"github.com/stackrox/rox/generated/storage"
+	"github.com/stackrox/rox/pkg/postgres/pgutils"
 	"github.com/stackrox/rox/pkg/sac"
 	"github.com/stackrox/rox/pkg/utils"
 )
@@ -116,14 +117,16 @@ func readAllowed(ctx context.Context, typ storage.IntegrationHealth_Type) (bool,
 
 func (ds *datastoreImpl) getIntegrationsOfType(ctx context.Context, integrationType storage.IntegrationHealth_Type) ([]*storage.IntegrationHealth, error) {
 	var integrationHealth []*storage.IntegrationHealth
-	err := ds.store.Walk(ctx, func(obj *storage.IntegrationHealth) error {
-		if obj.GetType() == integrationType {
-			integrationHealth = append(integrationHealth, obj)
-		}
-		return nil
-	})
-
-	if err != nil {
+	walkFn := func() error {
+		integrationHealth = integrationHealth[:0]
+		return ds.store.Walk(ctx, func(obj *storage.IntegrationHealth) error {
+			if obj.GetType() == integrationType {
+				integrationHealth = append(integrationHealth, obj)
+			}
+			return nil
+		})
+	}
+	if err := pgutils.RetryIfPostgres(walkFn); err != nil {
 		return nil, err
 	}
 	return integrationHealth, nil