netflow: add slow path to Flow func

for case where flow has not been yet synced.

follow up on #85

Author: makasim

netflow/registry.go

@@ -20,6 +20,7 @@ type Registry struct {
 
 	hostsFlowsMux sync.Mutex
 	hostFlows     map[string]*Flow
+	notFoundFlows map[string]struct{}
 	closeCh       chan struct{}
 	closedCh      chan struct{}
 
@@ -43,7 +44,20 @@ func NewRegistry(httpHost string, d flowstate.Driver, l *slog.Logger) *Registry
 }
 
 func (fr *Registry) Flow(id flowstate.TransitionID) (flowstate.Flow, error) {
-	return fr.fr.Flow(id)
+	f, err := fr.fr.Flow(id)
+
+	// slow path, flow not found locally, might not synced yet
+	if errors.Is(err, flowstate.ErrFlowNotFound) {
+		stateCtx := &flowstate.StateCtx{}
+		if err := fr.d.GetStateByID(flowstate.GetStateByID(stateCtx, flowStateID(id), 0)); err != nil {
+			return nil, err
+		}
+
+		fr.setFlow(stateCtx.Current)
+		return fr.fr.Flow(id)
+	}
+
+	return f, nil
 }
 
 func (fr *Registry) SetFlow(id flowstate.TransitionID, flow flowstate.Flow) error {
@@ -162,52 +176,54 @@ func (fr *Registry) sync() bool {
 
 	for _, state := range res.States {
 		fr.sinceRev = state.Rev
+		fr.setFlow(state)
+	}
 
-		if state.Annotations[`flowstate.flow.http_host`] == `` {
-			fr.l.Warn("state has no 'flowstate.flow.http_host' annotation set, skipping", "state_id", state.ID, "state_rev", state.Rev)
-			continue
-		}
-		httpHost := state.Annotations[`flowstate.flow.http_host`]
+	return res.More
+}
 
-		// local flow, skip
-		if httpHost == fr.httpHost {
-			continue
-		}
+func (fr *Registry) setFlow(state flowstate.State) {
+	if state.Annotations[`flowstate.flow.http_host`] == `` {
+		fr.l.Warn("state has no 'flowstate.flow.http_host' annotation set, skipping", "state_id", state.ID, "state_rev", state.Rev)
+	}
+	httpHost := state.Annotations[`flowstate.flow.http_host`]
 
-		if state.Annotations[`flowstate.flow.transition_id`] == `` {
-			fr.l.Warn("flow state has no 'flowstate.flow.transition_id' annotation set, skipping", "state_id", state.ID, "state_rev", state.Rev)
-			continue
-		}
-		tsID := flowstate.TransitionID(state.Annotations[`flowstate.flow.transition_id`])
+	// local flow, skip
+	if httpHost == fr.httpHost {
+		return
+	}
 
-		if flowstate.Ended(state) {
-			if err := fr.fr.UnsetFlow(tsID); err != nil {
-				fr.l.Warn("flow registry: unset flow failed", "error", err, "transition_id", tsID, "state_id", state.ID, "state_rev", state.Rev)
-			}
+	if state.Annotations[`flowstate.flow.transition_id`] == `` {
+		fr.l.Warn("flow state has no 'flowstate.flow.transition_id' annotation set, skipping", "state_id", state.ID, "state_rev", state.Rev)
+		return
+	}
+	tsID := flowstate.TransitionID(state.Annotations[`flowstate.flow.transition_id`])
 
-			continue
+	if flowstate.Ended(state) {
+		if err := fr.fr.UnsetFlow(tsID); err != nil {
+			fr.l.Warn("flow registry: unset flow failed", "error", err, "transition_id", tsID, "state_id", state.ID, "state_rev", state.Rev)
 		}
 
-		fr.hostsFlowsMux.Lock()
-
-		if fr.hostFlows == nil {
-			fr.hostFlows = make(map[string]*Flow)
-		}
+		return
+	}
 
-		f, ok := fr.hostFlows[httpHost]
-		if !ok {
-			f = New(httpHost)
-			fr.hostFlows[httpHost] = f
-		}
+	fr.hostsFlowsMux.Lock()
 
-		fr.hostsFlowsMux.Unlock()
+	if fr.hostFlows == nil {
+		fr.hostFlows = make(map[string]*Flow)
+	}
 
-		if err := fr.fr.SetFlow(tsID, f); err != nil {
-			fr.l.Warn("flow registry: set flow failed", "error", err, "transition_id", tsID, "state_id", state.ID, "state_rev", state.Rev)
-		}
+	f, ok := fr.hostFlows[httpHost]
+	if !ok {
+		f = New(httpHost)
+		fr.hostFlows[httpHost] = f
 	}
 
-	return res.More
+	fr.hostsFlowsMux.Unlock()
+
+	if err := fr.fr.SetFlow(tsID, f); err != nil {
+		fr.l.Warn("flow registry: set flow failed", "error", err, "transition_id", tsID, "state_id", state.ID, "state_rev", state.Rev)
+	}
 }
 
 func flowStateID(tsID flowstate.TransitionID) flowstate.StateID {

netflow/registry_synctest_test.go

@@ -44,22 +44,7 @@ func TestRegistry(t *testing.T) {
 		})); err != nil {
 			t.Fatalf("expected no error, got %v", err)
 		}
-
-		// only local flows should be available
-
-		if _, err := firstFR.Flow(`aFlowOnFirstFR`); err != nil {
-			t.Fatalf("expected no error, got %v", err)
-		}
-		if _, err := firstFR.Flow(`aFlowOnSecondFR`); err == nil {
-			t.Fatalf("expected no error, got %v", err)
-		}
-		if _, err := secondFR.Flow(`aFlowOnSecondFR`); err != nil {
-			t.Fatalf("expected no error, got %v", err)
-		}
-		if _, err := secondFR.Flow(`aFlowOnFirstFR`); err == nil {
-			t.Fatalf("expected no error, got %v", err)
-		}
-
+		
 		time.Sleep(time.Second * 7)
 
 		// now we should be able to get all flows
@@ -143,3 +128,36 @@ func TestRegistry(t *testing.T) {
 		}
 	})
 }
+
+func TestRegistry_Flow_SlowPath(t *testing.T) {
+	synctest.Run(func() {
+		lh := slogassert.New(t, slog.LevelDebug, nil)
+		l := slog.New(slogassert.New(t, slog.LevelDebug, lh))
+		l = slog.Default()
+
+		d := memdriver.New(l)
+
+		fr := netflow.NewRegistry(`http://aHost:8080`, d, l)
+		defer fr.Close()
+
+		// make sure we skipped the first round of Registry.watchFlows call
+		time.Sleep(time.Second * 5)
+
+		// simulate a flow registered by someone else, but not yet available in the registry
+		stateCtx := &flowstate.StateCtx{
+			Current: flowstate.State{
+				ID: `flowstate.flow.aFlowID`,
+			},
+		}
+		stateCtx.Current.SetLabel(`flow.type`, `remote`)
+		stateCtx.Current.SetAnnotation(`flowstate.flow.transition_id`, `aFlowID`)
+		stateCtx.Current.SetAnnotation(`flowstate.flow.http_host`, `http://anotherHost:8080`)
+		if err := d.Commit(flowstate.Commit(flowstate.Pause(stateCtx))); err != nil {
+			t.Fatalf("expected no error, got %v", err)
+		}
+
+		if _, err := fr.Flow(`aFlowID`); err != nil {
+			t.Fatalf("expected no error, got %v", err)
+		}
+	})
+}