stackrox · vikin91 · Sep 11, 2025 · Aug 26, 2025 · Aug 26, 2025 · Aug 29, 2025
diff --git a/pkg/env/sensor.go b/pkg/env/sensor.go
@@ -48,6 +48,9 @@ var (
 	// NetworkFlowClosedConnRememberDuration controls how long the categorized update computer will track
 	// timestamps for closed connections to handle late-arriving updates.
 	NetworkFlowClosedConnRememberDuration = registerDurationSetting("ROX_NETFLOW_CLOSED_CONN_REMEMBER_DURATION", 6*time.Minute)
+	// NetworkFlowUseLegacyUpdateComputer enables the Legacy update computer for the network flow enrichment pipeline
+	// updates sent to Central. Setting this to `true` enables the behavior as in 4.8 and earlier.
+	NetworkFlowUseLegacyUpdateComputer = RegisterBooleanSetting("ROX_NETFLOW_USE_LEGACY_UPDATE_COMPUTER", false)
 
 	// NetworkFlowDeduperHashingAlgorithm selects the hashing algorithm used for the deduper in the process of
 	// computing the updates for Central.

@@ -194,6 +194,16 @@ func (c *connStatus) pastContainerResolutionDeadline(now timestamp.MicroTS) bool
 	return c.timeElapsedSinceFirstSeen(now) > env.ContainerIDResolutionGracePeriod.DurationSetting()
 }
 
+// checkRemoveCondition returns true when the given entity can be removed from the enrichment queue.
+// It returns false if the enrichment should be retried in the next tick.
+func (c *connStatus) checkRemoveCondition(useLegacy, isConsumed bool) bool {
+	// Legacy UpdateComputer requires keeping all open entities in the enrichment queue until they are closed.
+	if useLegacy {
+		return c.rotten || (c.isClosed() && isConsumed)
+	}
+	return c.rotten || isConsumed
+}
+
 // containerResolutionResult holds the result of container ID resolution
 type containerResolutionResult struct {
 	Container          clusterentities.ContainerMetadata

@@ -11,7 +11,6 @@ import (
 	"github.com/stackrox/rox/pkg/networkgraph"
 	"github.com/stackrox/rox/pkg/timestamp"
 	"github.com/stackrox/rox/sensor/common/networkflow/manager/indicator"
-	"github.com/stackrox/rox/sensor/common/networkflow/updatecomputer"
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/mock/gomock"
 )
@@ -241,7 +240,7 @@ func TestEnrichConnection_BusinessLogicPaths(t *testing.T) {
 			enrichTickerC := make(chan time.Time)
 			defer close(enrichTickerC)
 
-			m, mockEntityStore, mockExternalSrc, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+			m, mockEntityStore, mockExternalSrc, _ := createManager(mockCtrl, enrichTickerC)
 
 			// Setup mocks
 			mocks := newMockExpectations(mockEntityStore, mockExternalSrc)
@@ -364,7 +363,7 @@ func TestEnrichContainerEndpoint_EdgeCases(t *testing.T) {
 			enrichTickerC := make(chan time.Time)
 			defer close(enrichTickerC)
 
-			m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+			m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 
 			// Setup mocks
 			mocks := newMockExpectations(mockEntityStore, nil)
@@ -394,3 +393,102 @@ func TestEnrichContainerEndpoint_EdgeCases(t *testing.T) {
 		})
 	}
 }
+
+func Test_connStatus_checkRemoveCondition(t *testing.T) {
+	tests := map[string]struct {
+		rotten     bool
+		closed     bool
+		useLegacy  bool
+		isConsumed bool
+		want       bool
+	}{
+		// Legacy
+		"Legacy shall remove closed, consumed EEs": {
+			rotten:     false,
+			closed:     true,
+			useLegacy:  true,
+			isConsumed: true,
+			want:       true,
+		},
+		"Legacy shall keep closed, unconsumed EEs": {
+			rotten:     false,
+			closed:     true,
+			useLegacy:  true,
+			isConsumed: false,
+			want:       false,
+		},
+		"Legacy shall keep open, consumed EEs": {
+			rotten:     false,
+			closed:     false,
+			useLegacy:  true,
+			isConsumed: true,
+			want:       false,
+		},
+		"Legacy shall keep open, unconsumed EEs": {
+			rotten:     false,
+			closed:     false,
+			useLegacy:  true,
+			isConsumed: false,
+			want:       false,
+		},
+		// TransitionBased (current impl),
+		"Current impl shall remove closed, consumed EEs": {
+			rotten:     false,
+			closed:     true,
+			useLegacy:  false,
+			isConsumed: true,
+			want:       true,
+		},
+		"Current impl shall keep closed, unconsumed EEs": {
+			rotten:     false,
+			closed:     true,
+			useLegacy:  false,
+			isConsumed: false,
+			want:       false,
+		},
+		"Current impl shall remove open, consumed EEs": { // difference to legacy
+			rotten:     false,
+			closed:     false,
+			useLegacy:  false,
+			isConsumed: true,
+			want:       true,
+		},
+		"Current impl shall keep open, unconsumed EEs": {
+			rotten:     false,
+			closed:     false,
+			useLegacy:  false,
+			isConsumed: false,
+			want:       false,
+		},
+		// Rotten
+		"Legacy shall remove rotten": {
+			rotten:     true,
+			closed:     false,
+			useLegacy:  true,
+			isConsumed: false,
+			want:       true,
+		},
+		"Current impl shall remove rotten": {
+			rotten:     true,
+			closed:     false,
+			useLegacy:  false,
+			isConsumed: false,
+			want:       true,
+		},
+	}
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			ts := timestamp.InfiniteFuture
+			if tt.closed {
+				ts = timestamp.Now()
+			}
+			c := &connStatus{
+				rotten:   tt.rotten,
+				lastSeen: ts,
+			}
+			assert.Equalf(t, tt.want,
+				c.checkRemoveCondition(tt.useLegacy, tt.isConsumed),
+				"checkRemoveCondition(%v, %v)", tt.useLegacy, tt.isConsumed)
+		})
+	}
+}
@@ -6,6 +6,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stackrox/rox/pkg/centralsensor"
 	"github.com/stackrox/rox/pkg/concurrency"
+	"github.com/stackrox/rox/pkg/env"
 	"github.com/stackrox/rox/pkg/features"
 	"github.com/stackrox/rox/pkg/net"
 	"github.com/stackrox/rox/pkg/networkgraph"
@@ -39,7 +40,7 @@ func (m *networkFlowManager) executeConnectionAction(
 	case PostEnrichmentActionRetry:
 		// noop, retry happens through not removing from `hostConns.connections`
 	case PostEnrichmentActionCheckRemove:
-		if status.rotten || (status.isClosed() && status.enrichmentConsumption.consumedNetworkGraph) {
+		if status.checkRemoveCondition(env.NetworkFlowUseLegacyUpdateComputer.BooleanSetting(), status.enrichmentConsumption.consumedNetworkGraph) {
 			delete(hostConns.connections, *conn)
 			flowMetrics.HostConnectionsOperations.WithLabelValues("remove", "connections").Inc()
 		}

@@ -37,7 +37,7 @@ func (m *networkFlowManager) executeEndpointAction(
 	case PostEnrichmentActionRetry:
 		// noop, retry happens through not removing from `hostConns.endpoints`
 	case PostEnrichmentActionCheckRemove:
-		if status.rotten || (status.isClosed() && status.enrichmentConsumption.IsConsumed()) {
+		if status.checkRemoveCondition(env.NetworkFlowUseLegacyUpdateComputer.BooleanSetting(), status.enrichmentConsumption.IsConsumed()) {
 			delete(hostConns.endpoints, *ep)
 			flowMetrics.HostConnectionsOperations.WithLabelValues("remove", "endpoints").Inc()
 			flowMetrics.HostProcessesEvents.WithLabelValues("remove").Inc()

@@ -12,7 +12,6 @@ import (
 	"github.com/stackrox/rox/pkg/networkgraph"
 	"github.com/stackrox/rox/pkg/timestamp"
 	"github.com/stackrox/rox/sensor/common/networkflow/manager/indicator"
-	"github.com/stackrox/rox/sensor/common/networkflow/updatecomputer"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/mock/gomock"
 )
@@ -31,7 +30,7 @@ func (s *TestNetworkFlowManagerEnrichmentTestSuite) TestEnrichConnection() {
 	enrichTickerC := make(chan time.Time)
 	defer close(enrichTickerC)
 	defer mockCtrl.Finish()
-	m, mockEntityStore, mockExternalSrc, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+	m, mockEntityStore, mockExternalSrc, _ := createManager(mockCtrl, enrichTickerC)
 	srcID := "src-id"
 	dstID := "dst-id"
 
@@ -494,7 +493,7 @@ func (s *TestNetworkFlowManagerEnrichmentTestSuite) TestEnrichContainerEndpoint(
 
 	for name, tc := range cases {
 		s.Run(name, func() {
-			m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+			m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 
 			// Setup environment variables
 			s.T().Setenv(env.ProcessesListeningOnPort.EnvVar(), strconv.FormatBool(tc.plopFeatEnabled))

@@ -16,7 +16,6 @@ import (
 	"github.com/stackrox/rox/pkg/timestamp"
 	"github.com/stackrox/rox/sensor/common"
 	"github.com/stackrox/rox/sensor/common/clusterentities"
-	"github.com/stackrox/rox/sensor/common/networkflow/updatecomputer"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
@@ -212,7 +211,7 @@ func (s *NetworkFlowManagerTestSuite) TestManagerOfflineMode() {
 	enrichTickerC := make(chan time.Time)
 	defer close(enrichTickerC)
 	defer mockCtrl.Finish()
-	m, mockEntity, _, mockDetector := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+	m, mockEntity, _, mockDetector := createManager(mockCtrl, enrichTickerC)
 	states := []struct {
 		testName                    string
 		notify                      common.SensorComponentEvent
@@ -375,7 +374,7 @@ func (s *NetworkFlowManagerTestSuite) TestExpireMessage() {
 	enrichTickerC := make(chan time.Time)
 	defer close(enrichTickerC)
 	defer mockCtrl.Finish()
-	m, mockEntity, _, mockDetector := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+	m, mockEntity, _, mockDetector := createManager(mockCtrl, enrichTickerC)
 	go m.enrichConnections(enrichTickerC)
 	mockEntity.EXPECT().LookupByContainerID(gomock.Any()).Times(1).DoAndReturn(func(_ any) (clusterentities.ContainerMetadata, bool, bool) {
 		return clusterentities.ContainerMetadata{

@@ -13,7 +13,6 @@ import (
 	"github.com/stackrox/rox/sensor/common/clusterentities"
 	mocksDetector "github.com/stackrox/rox/sensor/common/detector/mocks"
 	mocksManager "github.com/stackrox/rox/sensor/common/networkflow/manager/mocks"
-	"github.com/stackrox/rox/sensor/common/networkflow/updatecomputer"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/mock/gomock"
@@ -48,7 +47,7 @@ func (b *sendNetflowsSuite) SetupTest() {
 	b.mockCtrl = gomock.NewController(b.T())
 	enrichTickerC := make(chan time.Time)
 	defer close(enrichTickerC)
-	b.m, b.mockEntity, _, b.mockDetector = createManager(b.mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+	b.m, b.mockEntity, _, b.mockDetector = createManager(b.mockCtrl, enrichTickerC)
 
 	b.fakeTicker = make(chan time.Time)
 	go b.m.enrichConnections(b.fakeTicker)
@@ -163,7 +162,7 @@ func (b *sendNetflowsSuite) TestCloseOldEndpointFailedLookup() {
 }
 
 func (b *sendNetflowsSuite) TestUnchangedConnection() {
-	b.expectLookups(2)
+	b.expectLookups(1)
 	b.expectDetections(1)
 
 	b.updateConn(createConnectionPair().lastSeen(timestamp.InfiniteFuture))

@@ -10,7 +10,6 @@ import (
 	"github.com/stackrox/rox/pkg/timestamp"
 	"github.com/stackrox/rox/sensor/common"
 	"github.com/stackrox/rox/sensor/common/networkflow/manager/indicator"
-	"github.com/stackrox/rox/sensor/common/networkflow/updatecomputer"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/mock/gomock"
 )
@@ -31,7 +30,7 @@ func (s *NetworkFlowPurgerTestSuite) TestPurgerStartWithTicker() {
 	s.T().Setenv(env.EnrichmentPurgerTickerCycle.EnvVar(), "1s")
 	s.Equal(time.Second, nonZeroPurgerCycle())
 
-	m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+	m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 	purger := NewNetworkFlowPurger(mockEntityStore, time.Hour, WithManager(m))
 	s.NoError(purger.Start())
 	// Enable the ticker after going online - send the same signal that activates the manager
@@ -50,7 +49,7 @@ func (s *NetworkFlowPurgerTestSuite) TestDisabledPurger() {
 	// Disabling the purger
 	s.T().Setenv(env.EnrichmentPurgerTickerCycle.EnvVar(), "0s")
 
-	m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+	m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 	purger := NewNetworkFlowPurger(mockEntityStore, time.Hour, WithManager(m), WithPurgerTicker(s.T(), purgerTickerC))
 
 	s.ErrorContains(purger.Start(), "purger is disabled")
@@ -72,7 +71,7 @@ func (s *NetworkFlowPurgerTestSuite) TestPurgerWithoutManager() {
 	enrichTickerC := make(chan time.Time)
 	defer close(enrichTickerC)
 	defer mockCtrl.Finish()
-	_, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+	_, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 	// Don't set manager to explicitly simulate disconnected purger
 	purger := NewNetworkFlowPurger(mockEntityStore, time.Hour, WithPurgerTicker(s.T(), purgerTickerC))
 
@@ -124,7 +123,7 @@ func (s *NetworkFlowPurgerTestSuite) TestPurgerWithManager() {
 			now := time.Now()
 			lastUpdateTS := timestamp.FromGoTime(now.Add(-tc.lastUpdateTime))
 
-			m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+			m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 			purger := NewNetworkFlowPurger(mockEntityStore, tc.purgerMaxAge, WithManager(m), WithPurgerTicker(s.T(), purgerTickerC))
 
 			expectationsEndpointPurger(mockEntityStore, tc.isKnownEndpoint, true, false)
@@ -214,7 +213,7 @@ func (s *NetworkFlowPurgerTestSuite) TestPurgerHostConnsEndpoints() {
 			now := time.Now()
 			lastUpdateTS := timestamp.FromGoTime(now.Add(-tc.lastUpdateTime))
 
-			m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+			m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 			expectationsEndpointPurger(mockEntityStore, tc.isKnownEndpoint, tc.foundContainerID, false)
 			ep := createEndpointPair(timestamp.FromGoTime(now.Add(-tc.firstSeen)), lastUpdateTS)
 			concurrency.WithLock(&m.connectionsByHostMutex, func() {
@@ -283,7 +282,7 @@ func (s *NetworkFlowPurgerTestSuite) TestPurgerHostConnsConnections() {
 			now := time.Now()
 			lastUpdateTS := timestamp.FromGoTime(now.Add(-tc.lastUpdateTime))
 
-			m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+			m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 			expectationsEndpointPurger(mockEntityStore, true, tc.foundContainerID, tc.containerIDHistorical)
 
 			pair := createConnectionPair().
@@ -348,7 +347,7 @@ func (s *NetworkFlowPurgerTestSuite) TestPurgerActiveConnections() {
 			now := time.Now()
 			lastUpdateTS := timestamp.FromGoTime(now.Add(-tc.lastUpdateTime))
 
-			m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+			m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 			expectationsEndpointPurger(mockEntityStore, true, tc.foundContainerID, tc.containerIDHistorical)
 
 			pair := createConnectionPair().
@@ -438,7 +437,7 @@ func (s *NetworkFlowPurgerTestSuite) TestPurgerActiveEndpoints() {
 			now := time.Now()
 			lastUpdateTS := timestamp.FromGoTime(now.Add(-tc.lastUpdateTime))
 
-			m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewTransitionBased(), enrichTickerC)
+			m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
 			expectationsEndpointPurger(mockEntityStore, tc.isKnownEndpoint, tc.foundContainerID, tc.containerIDHistorical)
 
 			ep := createEndpointPair(timestamp.FromGoTime(now.Add(-tc.firstSeen)), lastUpdateTS)

@@ -21,15 +21,15 @@ import (
 	"go.uber.org/mock/gomock"
 )
 
-func createManager(mockCtrl *gomock.Controller, updateComputer updatecomputer.UpdateComputer, enrichTicker <-chan time.Time) (*networkFlowManager, *mocksManager.MockEntityStore, *mocksExternalSrc.MockStore, *mocksDetector.MockDetector) {
+func createManager(mockCtrl *gomock.Controller, enrichTicker <-chan time.Time) (*networkFlowManager, *mocksManager.MockEntityStore, *mocksExternalSrc.MockStore, *mocksDetector.MockDetector) {
 	mockEntityStore := mocksManager.NewMockEntityStore(mockCtrl)
 	mockExternalStore := mocksExternalSrc.NewMockStore(mockCtrl)
 	mockDetector := mocksDetector.NewMockDetector(mockCtrl)
 	mgr := &networkFlowManager{
 		clusterEntities:   mockEntityStore,
 		externalSrcs:      mockExternalStore,
 		policyDetector:    mockDetector,
-		updateComputer:    updateComputer,
+		updateComputer:    updatecomputer.New(),
 		connectionsByHost: make(map[string]*hostConnections),
 		sensorUpdates:     make(chan *message.ExpiringMessage, 5),
 		publicIPs:         newPublicIPsManager(),

@@ -5,6 +5,7 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stackrox/rox/generated/storage"
+	"github.com/stackrox/rox/pkg/env"
 	"github.com/stackrox/rox/pkg/timestamp"
 	"github.com/stackrox/rox/sensor/common/networkflow/manager/indicator"
 )
@@ -30,3 +31,10 @@ type UpdateComputer interface {
 	// RecordSizeMetrics records metrics for length and byte-size of the collections used in updateComputer.
 	RecordSizeMetrics(gv1, gv2 *prometheus.GaugeVec)
 }
+
+func New() UpdateComputer {
+	if env.NetworkFlowUseLegacyUpdateComputer.BooleanSetting() {
+		return NewLegacy()
+	}
+	return NewTransitionBased()
+}
@@ -128,8 +128,9 @@ func CreateSensor(cfg *CreateOptions) (*sensor.Sensor, error) {
 	} else {
 		processSignals = signalService.New(processPipeline, indicators, signalService.WithTraceWriter(cfg.processIndicatorWriter))
 	}
+
 	networkFlowManager :=
-		manager.NewManager(storeProvider.Entities(), externalsrcs.StoreInstance(), policyDetector, pubSub, updatecomputer.NewTransitionBased(), manager.WithEnrichTicker(cfg.networkFlowTicker))
+		manager.NewManager(storeProvider.Entities(), externalsrcs.StoreInstance(), policyDetector, pubSub, updatecomputer.New(), manager.WithEnrichTicker(cfg.networkFlowTicker))
 	enhancer := deploymentenhancer.CreateEnhancer(storeProvider)
 	components := []common.SensorComponent{
 		admCtrlMsgForwarder,