Make sure that closing the graph is more secure (#19178)

Signed-off-by: SimFG <bang.fu@zilliz.com>

Signed-off-by: SimFG <bang.fu@zilliz.com>

internal/util/flowgraph/flow_graph.go

@@ -27,6 +27,7 @@ type TimeTickedFlowGraph struct {
 	nodeCtx   map[NodeName]*nodeCtx
 	stopOnce  sync.Once
 	startOnce sync.Once
+	closeWg   *sync.WaitGroup
 }
 
 // AddNode add Node into flowgraph
@@ -35,6 +36,7 @@ func (fg *TimeTickedFlowGraph) AddNode(node Node) {
 		node:                   node,
 		downstreamInputChanIdx: make(map[string]int),
 		closeCh:                make(chan struct{}),
+		closeWg:                fg.closeWg,
 	}
 	fg.nodeCtx[node.Name()] = &nodeCtx
 }
@@ -92,6 +94,7 @@ func (fg *TimeTickedFlowGraph) Close() {
 				v.Close()
 			}
 		}
+		fg.closeWg.Wait()
 	})
 }
 
@@ -99,6 +102,7 @@ func (fg *TimeTickedFlowGraph) Close() {
 func NewTimeTickedFlowGraph(ctx context.Context) *TimeTickedFlowGraph {
 	flowGraph := TimeTickedFlowGraph{
 		nodeCtx: make(map[string]*nodeCtx),
+		closeWg: &sync.WaitGroup{},
 	}
 
 	return &flowGraph

internal/util/flowgraph/node.go

@@ -59,12 +59,14 @@ type nodeCtx struct {
 	downstreamInputChanIdx map[string]int
 
 	closeCh chan struct{} // notify work to exit
+	closeWg *sync.WaitGroup
 }
 
 // Start invoke Node `Start` method and start a worker goroutine
 func (nodeCtx *nodeCtx) Start() {
 	nodeCtx.node.Start()
 
+	nodeCtx.closeWg.Add(1)
 	go nodeCtx.work()
 }
 
@@ -114,6 +116,7 @@ func (nodeCtx *nodeCtx) work() {
 			// the res decide whether the node should be closed.
 			if isCloseMsg(res) {
 				close(nodeCtx.closeCh)
+				nodeCtx.closeWg.Done()
 				nodeCtx.node.Close()
 			}
 

internal/util/flowgraph/node_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"math"
 	"os"
+	"sync"
 	"testing"
 	"time"
 
@@ -80,6 +81,7 @@ func TestNodeCtx_Start(t *testing.T) {
 		inputChannels:          make([]chan Msg, 2),
 		downstreamInputChanIdx: make(map[string]int),
 		closeCh:                make(chan struct{}),
+		closeWg:                &sync.WaitGroup{},
 	}
 
 	for i := 0; i < len(node.inputChannels); i++ {