Skip to content

ROX-30390: Extract Update Computer #16452

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vikin91 merged 17 commits into from
Sep 1, 2025
Merged

ROX-30390: Extract Update Computer #16452

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
2468aac
Empty commit
vikin91 Aug 19, 2025
4c7091a
Extract update calculation logic into updatecomputer
vikin91 Aug 19, 2025
16f0f9b
Minor: Add comments
vikin91 Aug 21, 2025
a448b13
Update comments and minimize the diff
vikin91 Aug 21, 2025
7c0f995
Fix metric for processes size
vikin91 Aug 21, 2025
ee1c30b
Cleanup tests and comments
vikin91 Aug 21, 2025
57ef6e6
Add more test cases to TestComputeUpdatedEndpoints
vikin91 Aug 21, 2025
689786e
Add more test cases to TestComputeUpdatedProcesses
vikin91 Aug 21, 2025
47e8a87
Extract updateComputer as dependency to NetworkManager
vikin91 Aug 21, 2025
d257454
Update internal state only if sending to Central was successful
vikin91 Aug 27, 2025
e59950d
Use maps.Clone in UpdateState. Skip updates on nil
vikin91 Aug 27, 2025
f92a387
Delete unused const loggingRateLimiter
vikin91 Aug 27, 2025
f289a25
Use open var in test for the ease of reading
vikin91 Aug 27, 2025
4fa2881
Add testcases for handling nils
vikin91 Aug 27, 2025
c3e542c
Improve metric help texts
vikin91 Aug 27, 2025
b51eec1
Remove name arg from RecordSizeMetrics
vikin91 Aug 27, 2025
9f663c4
Ensure that call to UpdateState(nil, nil, nil) changes no state
vikin91 Aug 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions sensor/common/networkflow/manager/enrichment_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ 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"
)
Expand Down Expand Up @@ -240,7 +241,7 @@ func TestEnrichConnection_BusinessLogicPaths(t *testing.T) {
enrichTickerC := make(chan time.Time)
defer close(enrichTickerC)

m, mockEntityStore, mockExternalSrc, _ := createManager(mockCtrl, enrichTickerC)
m, mockEntityStore, mockExternalSrc, _ := createManager(mockCtrl, updatecomputer.NewLegacy(), enrichTickerC)

// Setup mocks
mocks := newMockExpectations(mockEntityStore, mockExternalSrc)
Expand Down Expand Up @@ -363,7 +364,7 @@ func TestEnrichContainerEndpoint_EdgeCases(t *testing.T) {
enrichTickerC := make(chan time.Time)
defer close(enrichTickerC)

m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewLegacy(), enrichTickerC)

// Setup mocks
mocks := newMockExpectations(mockEntityStore, nil)
Expand Down
5 changes: 3 additions & 2 deletions sensor/common/networkflow/manager/manager_enrich_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ 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"
)
Expand All @@ -30,7 +31,7 @@ func (s *TestNetworkFlowManagerEnrichmentTestSuite) TestEnrichConnection() {
enrichTickerC := make(chan time.Time)
defer close(enrichTickerC)
defer mockCtrl.Finish()
m, mockEntityStore, mockExternalSrc, _ := createManager(mockCtrl, enrichTickerC)
m, mockEntityStore, mockExternalSrc, _ := createManager(mockCtrl, updatecomputer.NewLegacy(), enrichTickerC)
srcID := "src-id"
dstID := "dst-id"

Expand Down Expand Up @@ -493,7 +494,7 @@ func (s *TestNetworkFlowManagerEnrichmentTestSuite) TestEnrichContainerEndpoint(

for name, tc := range cases {
s.Run(name, func() {
m, mockEntityStore, _, _ := createManager(mockCtrl, enrichTickerC)
m, mockEntityStore, _, _ := createManager(mockCtrl, updatecomputer.NewLegacy(), enrichTickerC)

// Setup environment variables
s.T().Setenv(env.ProcessesListeningOnPort.EnvVar(), strconv.FormatBool(tc.plopFeatEnabled))
Expand Down
176 changes: 29 additions & 147 deletions sensor/common/networkflow/manager/manager_impl.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,9 @@ import (
"github.com/stackrox/rox/pkg/concurrency"
"github.com/stackrox/rox/pkg/env"
"github.com/stackrox/rox/pkg/features"
"github.com/stackrox/rox/pkg/logging"
"github.com/stackrox/rox/pkg/net"
"github.com/stackrox/rox/pkg/netutil"
"github.com/stackrox/rox/pkg/process/normalize"

"github.com/stackrox/rox/pkg/net"
"github.com/stackrox/rox/pkg/protocompat"
"github.com/stackrox/rox/pkg/sensor/queue"
"github.com/stackrox/rox/pkg/sync"
Expand All @@ -32,13 +30,11 @@ import (
"github.com/stackrox/rox/sensor/common/metrics"
"github.com/stackrox/rox/sensor/common/networkflow/manager/indicator"
flowMetrics "github.com/stackrox/rox/sensor/common/networkflow/metrics"
"github.com/stackrox/rox/sensor/common/networkflow/updatecomputer"
"github.com/stackrox/rox/sensor/common/unimplemented"
)

const (
connectionDeletionGracePeriod = 5 * time.Minute
loggingRateLimiter = "plop-feature-disabled"
)
const connectionDeletionGracePeriod = 5 * time.Minute

var (
emptyProcessInfo = indicator.ProcessInfo{}
Expand Down Expand Up @@ -147,10 +143,10 @@ func NewManager(
externalSrcs externalsrcs.Store,
policyDetector detector.Detector,
pubSub *internalmessage.MessageSubscriber,
updateComputer updatecomputer.UpdateComputer,
opts ...Option,
) Manager {
enricherTicker := time.NewTicker(enricherCycle)

mgr := &networkFlowManager{
connectionsByHost: make(map[string]*hostConnections),
clusterEntities: clusterEntities,
Expand All @@ -159,6 +155,7 @@ func NewManager(
policyDetector: policyDetector,
enricherTicker: enricherTicker,
enricherTickerC: enricherTicker.C,
updateComputer: updateComputer,
initialSync: &atomic.Bool{},
activeConnections: make(map[connection]*networkConnIndicatorWithAge),
activeEndpoints: make(map[containerEndpoint]*containerEndpointIndicatorWithAge),
Expand Down Expand Up @@ -212,10 +209,8 @@ type networkFlowManager struct {
clusterEntities EntityStore
externalSrcs externalsrcs.Store

lastSentStateMutex sync.RWMutex
enrichedConnsLastSentState map[indicator.NetworkConn]timestamp.MicroTS
enrichedEndpointsLastSentState map[indicator.ContainerEndpoint]timestamp.MicroTS
enrichedProcessesLastSentState map[indicator.ProcessListening]timestamp.MicroTS
// UpdateComputer implementation for computing flow updates that are sent to Central on each tick.
updateComputer updatecomputer.UpdateComputer

activeConnectionsMutex sync.RWMutex
// activeConnections tracks all connections reported by Collector that are believed to be active.
Expand Down Expand Up @@ -345,24 +340,10 @@ func (m *networkFlowManager) sendToCentral(msg *central.MsgFromSensor) bool {
}

func (m *networkFlowManager) resetLastSentState() {
m.lastSentStateMutex.Lock()
defer m.lastSentStateMutex.Unlock()
m.enrichedConnsLastSentState = nil
m.enrichedEndpointsLastSentState = nil
m.enrichedProcessesLastSentState = nil
}

func (m *networkFlowManager) updateConnectionStates(newConns map[indicator.NetworkConn]timestamp.MicroTS, newEndpoints map[indicator.ContainerEndpoint]timestamp.MicroTS) {
m.lastSentStateMutex.Lock()
defer m.lastSentStateMutex.Unlock()
m.enrichedConnsLastSentState = newConns
m.enrichedEndpointsLastSentState = newEndpoints
}

func (m *networkFlowManager) updateProcessesState(newProcesses map[indicator.ProcessListening]timestamp.MicroTS) {
m.lastSentStateMutex.Lock()
defer m.lastSentStateMutex.Unlock()
m.enrichedProcessesLastSentState = newProcesses
// Reset state in the UpdateComputer implementation
if m.updateComputer != nil {
m.updateComputer.ResetState()
}
}

func (m *networkFlowManager) enrichConnections(tickerC <-chan time.Time) {
Expand Down Expand Up @@ -390,6 +371,7 @@ func (m *networkFlowManager) getCurrentContext() context.Context {
}

func (m *networkFlowManager) updateEnrichmentCollectionsSize() {
// Number of entities (connections, endpoints) waiting for enrichment
numConnections := 0
numEndpoints := 0
concurrency.WithRLock(&m.connectionsByHostMutex, func() {
Expand All @@ -403,16 +385,16 @@ func (m *networkFlowManager) updateEnrichmentCollectionsSize() {
flowMetrics.EnrichmentCollectionsSize.WithLabelValues("connectionsInEnrichQueue", "connections").Set(float64(numConnections))
flowMetrics.EnrichmentCollectionsSize.WithLabelValues("endpointsInEnrichQueue", "endpoints").Set(float64(numEndpoints))

// Number of entities (connections, endpoints) stored in memory for the purposes of not losing data while offline.
concurrency.WithRLock(&m.activeConnectionsMutex, func() {
flowMetrics.EnrichmentCollectionsSize.WithLabelValues("activeConnections", "connections").Set(float64(len(m.activeConnections)))
flowMetrics.EnrichmentCollectionsSize.WithLabelValues("activeEndpoints", "endpoints").Set(float64(len(m.activeEndpoints)))
})

concurrency.WithRLock(&m.lastSentStateMutex, func() {
flowMetrics.EnrichmentCollectionsSize.WithLabelValues("enrichedConnectionsLastSentState", "connections").Set(float64(len(m.enrichedConnsLastSentState)))
flowMetrics.EnrichmentCollectionsSize.WithLabelValues("enrichedEndpointsLastSentState", "endpoints").Set(float64(len(m.enrichedEndpointsLastSentState)))
flowMetrics.EnrichmentCollectionsSize.WithLabelValues("enrichedProcessesLastSentState", "processes").Set(float64(len(m.enrichedProcessesLastSentState)))
})
// Length and byte sizes of collections used internally by updatecomputer
if m.updateComputer != nil {
m.updateComputer.RecordSizeMetrics(flowMetrics.EnrichmentCollectionsSize, flowMetrics.EnrichmentCollectionsSizeBytes)
}
}

func (m *networkFlowManager) enrichAndSend() {
Expand All @@ -421,27 +403,28 @@ func (m *networkFlowManager) enrichAndSend() {
// Updates m.activeEndpoints and m.activeConnections if lastSeen was reported as null by the Collector.
currentConns, currentEndpoints, currentProcesses := m.currentEnrichedConnsAndEndpoints()

// Compares currently enriched connections & endpoints with those enriched in the previous cycle.
// The new changes are sent to Central.
updatedConns := computeUpdatedConns(currentConns, m.enrichedConnsLastSentState, &m.lastSentStateMutex)
updatedEndpoints := computeUpdatedEndpoints(currentEndpoints, m.enrichedEndpointsLastSentState, &m.lastSentStateMutex)
updatedProcesses := computeUpdatedProcesses(currentProcesses, m.enrichedProcessesLastSentState, &m.lastSentStateMutex)
flowMetrics.NumUpdatesSentToCentral.WithLabelValues("connections").Add(float64(len(updatedConns)))
flowMetrics.NumUpdatesSentToCentral.WithLabelValues("endpoints").Add(float64(len(updatedEndpoints)))
flowMetrics.NumUpdatesSentToCentral.WithLabelValues("processes").Add(float64(len(updatedProcesses)))
// The new changes are sent to Central using the update computer implementation.
updatedConns := m.updateComputer.ComputeUpdatedConns(currentConns)
updatedEndpoints := m.updateComputer.ComputeUpdatedEndpoints(currentEndpoints)
updatedProcesses := m.updateComputer.ComputeUpdatedProcesses(currentProcesses)

if len(updatedConns)+len(updatedEndpoints) > 0 {
if sent := m.sendConnsEps(updatedConns, updatedEndpoints); sent {
m.updateConnectionStates(currentConns, currentEndpoints)
// Update the UpdateComputer's internal state after sending updates to Central.
// This is important for update computers that rely on the state from the previous tick.
m.updateComputer.UpdateState(currentConns, currentEndpoints, nil)
}
metrics.SetNetworkFlowBufferSizeGauge(len(m.sensorUpdates))
}
if env.ProcessesListeningOnPort.BooleanSetting() && len(updatedProcesses) > 0 {
if sent := m.sendProcesses(updatedProcesses); sent {
m.updateProcessesState(currentProcesses)
// Update the UpdateComputer's internal state after sending updates to Central.
// This is important for update computers that rely on the state from the previous tick.
m.updateComputer.UpdateState(nil, nil, currentProcesses)
}
}
metrics.SetNetworkFlowBufferSizeGauge(len(m.sensorUpdates))
}

func (m *networkFlowManager) sendConnsEps(conns []*storage.NetworkFlow, eps []*storage.NetworkEndpoint) bool {
protoToSend := &central.NetworkFlowUpdate{
Updated: conns,
Expand Down Expand Up @@ -498,107 +481,6 @@ func (m *networkFlowManager) currentEnrichedConnsAndEndpoints() (
return enrichedConnections, enrichedEndpoints, enrichedProcesses
}

// isUpdated determines if a connection/endpoint should be sent to Central based on timestamp comparison.
//
// Timestamp Convention:
// - timestamp.InfiniteFuture = connection/endpoint is OPEN (still active)
// - Any other timestamp value = connection/endpoint is CLOSED (lastSeen/closeTime)
//
// This function detects updates when:
// 1. New connection/endpoint (not seen before)
// 2. More recent activity (newer timestamp)
// 3. State transition from OPEN -> CLOSED (InfiniteFuture -> actual timestamp)
func isUpdated(prevTS, currTS timestamp.MicroTS, seenPreviously bool) bool {
// Connection has not been seen in the last tick.
if !seenPreviously {
return true
}
// Collector saw this connection more recently.
if currTS > prevTS {
return true
}
// Connection was active (unclosed) in the last tick, now it is closed.
if prevTS == timestamp.InfiniteFuture && currTS != timestamp.InfiniteFuture {
return true
}
return false
}

func computeUpdatedConns(current map[indicator.NetworkConn]timestamp.MicroTS, previous map[indicator.NetworkConn]timestamp.MicroTS, previousMutex *sync.RWMutex) []*storage.NetworkFlow {
previousMutex.RLock()
defer previousMutex.RUnlock()
var updates []*storage.NetworkFlow

for conn, currTS := range current {
prevTS, seenPreviously := previous[conn]
if isUpdated(prevTS, currTS, seenPreviously) {
updates = append(updates, conn.ToProto(currTS))
}
}

for conn, prevTS := range previous {
if _, ok := current[conn]; !ok {
updates = append(updates, conn.ToProto(prevTS))
}
}

return updates
}

func computeUpdatedEndpoints(current map[indicator.ContainerEndpoint]timestamp.MicroTS, previous map[indicator.ContainerEndpoint]timestamp.MicroTS, previousMutex *sync.RWMutex) []*storage.NetworkEndpoint {
previousMutex.RLock()
defer previousMutex.RUnlock()
var updates []*storage.NetworkEndpoint

for ep, currTS := range current {
prevTS, seenPreviously := previous[ep]
if isUpdated(prevTS, currTS, seenPreviously) {
updates = append(updates, ep.ToProto(currTS))
}
}

for ep, prevTS := range previous {
if _, ok := current[ep]; !ok {
updates = append(updates, ep.ToProto(prevTS))
}
}

return updates
}

func computeUpdatedProcesses(current map[indicator.ProcessListening]timestamp.MicroTS, previous map[indicator.ProcessListening]timestamp.MicroTS, previousMutex *sync.RWMutex) []*storage.ProcessListeningOnPortFromSensor {
if !env.ProcessesListeningOnPort.BooleanSetting() {
if len(current) > 0 {
logging.GetRateLimitedLogger().WarnL(loggingRateLimiter,
"Received process while ProcessesListeningOnPort feature is disabled. This may indicate a misconfiguration.")
}
return []*storage.ProcessListeningOnPortFromSensor{}
}
previousMutex.RLock()
defer previousMutex.RUnlock()
var updates []*storage.ProcessListeningOnPortFromSensor

for pl, currTS := range current {
prevTS, seenPreviously := previous[pl]
if isUpdated(prevTS, currTS, seenPreviously) {
updates = append(updates, pl.ToProto(currTS))
}
}

for ep, prevTS := range previous {
if _, ok := current[ep]; !ok {
// This condition means the deployment was removed before we got the
// close timestamp for the endpoint. Use the current timestamp instead.
if prevTS == timestamp.InfiniteFuture {
prevTS = timestamp.Now()
}
updates = append(updates, ep.ToProto(prevTS))
}
}

return updates
}

func (m *networkFlowManager) getAllHostConnections() []*hostConnections {
// Get a snapshot of all *hostConnections. This allows us to lock the individual mutexes without having to hold
// two locks simultaneously.
Expand Down
5 changes: 3 additions & 2 deletions sensor/common/networkflow/manager/manager_impl_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ 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"
Expand Down Expand Up @@ -211,7 +212,7 @@ func (s *NetworkFlowManagerTestSuite) TestManagerOfflineMode() {
enrichTickerC := make(chan time.Time)
defer close(enrichTickerC)
defer mockCtrl.Finish()
m, mockEntity, _, mockDetector := createManager(mockCtrl, enrichTickerC)
m, mockEntity, _, mockDetector := createManager(mockCtrl, updatecomputer.NewLegacy(), enrichTickerC)
states := []struct {
testName string
notify common.SensorComponentEvent
Expand Down Expand Up @@ -374,7 +375,7 @@ func (s *NetworkFlowManagerTestSuite) TestExpireMessage() {
enrichTickerC := make(chan time.Time)
defer close(enrichTickerC)
defer mockCtrl.Finish()
m, mockEntity, _, mockDetector := createManager(mockCtrl, enrichTickerC)
m, mockEntity, _, mockDetector := createManager(mockCtrl, updatecomputer.NewLegacy(), enrichTickerC)
go m.enrichConnections(enrichTickerC)
mockEntity.EXPECT().LookupByContainerID(gomock.Any()).Times(1).DoAndReturn(func(_ any) (clusterentities.ContainerMetadata, bool, bool) {
return clusterentities.ContainerMetadata{
Expand Down
3 changes: 2 additions & 1 deletion sensor/common/networkflow/manager/manager_sending_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ 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"
Expand Down Expand Up @@ -47,7 +48,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, enrichTickerC)
b.m, b.mockEntity, _, b.mockDetector = createManager(b.mockCtrl, updatecomputer.NewLegacy(), enrichTickerC)

b.fakeTicker = make(chan time.Time)
go b.m.enrichConnections(b.fakeTicker)
Expand Down
Loading
Loading