Improve the close method in the graph (#19100)

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

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

internal/util/flowgraph/flow_graph.go

@@ -89,13 +89,6 @@ func (fg *TimeTickedFlowGraph) Close() {
 	fg.stopOnce.Do(func() {
 		for _, v := range fg.nodeCtx {
 			if v.node.IsInputNode() {
-				// close inputNode first
-				v.Close()
-			}
-		}
-		for _, v := range fg.nodeCtx {
-			if !v.node.IsInputNode() {
-				// close other nodes
 				v.Close()
 			}
 		}

internal/util/flowgraph/input_node.go

@@ -30,6 +30,7 @@ type InputNode struct {
 	BaseNode
 	inStream     msgstream.MsgStream
 	name         string
+	closeMsgChan chan struct{}
 }
 
 // IsInputNode returns whether Node is InputNode
@@ -44,11 +45,16 @@ func (inNode *InputNode) Start() {
 
 // Close implements node
 func (inNode *InputNode) Close() {
-	inNode.inStream.Close()
+	select {
+	case <-inNode.closeMsgChan:
+		return
+	default:
+		close(inNode.closeMsgChan)
 		log.Debug("message stream closed",
 			zap.String("node name", inNode.name),
 		)
 	}
+}
 
 // Name returns node name
 func (inNode *InputNode) Name() string {
@@ -62,7 +68,13 @@ func (inNode *InputNode) InStream() msgstream.MsgStream {
 
 // Operate consume a message pack from msgstream and return
 func (inNode *InputNode) Operate(in []Msg) []Msg {
-	msgPack, ok := <-inNode.inStream.Chan()
+	select {
+	case <-inNode.closeMsgChan:
+		inNode.inStream.Close()
+		return []Msg{&MsgStreamMsg{
+			isCloseMsg: true,
+		}}
+	case msgPack, ok := <-inNode.inStream.Chan():
 		if !ok {
 			log.Warn("MsgStream closed", zap.Any("input node", inNode.Name()))
 			return []Msg{}
@@ -94,6 +106,7 @@ func (inNode *InputNode) Operate(in []Msg) []Msg {
 
 		return []Msg{msgStreamMsg}
 	}
+}
 
 // NewInputNode composes an InputNode with provided MsgStream, name and parameters
 func NewInputNode(inStream msgstream.MsgStream, nodeName string, maxQueueLength int32, maxParallelism int32) *InputNode {
@@ -105,5 +118,6 @@ func NewInputNode(inStream msgstream.MsgStream, nodeName string, maxQueueLength
 		BaseNode:     baseNode,
 		inStream:     inStream,
 		name:         nodeName,
+		closeMsgChan: make(chan struct{}),
 	}
 }

internal/util/flowgraph/input_node_test.go

@@ -41,10 +41,7 @@ func TestInputNode(t *testing.T) {
 	produceStream.Produce(&msgPack)
 
 	nodeName := "input_node"
-	inputNode := &InputNode{
+	inputNode := NewInputNode(msgStream, nodeName, 100, 100)
-		inStream: msgStream,
-		name:     nodeName,
-	}
 	defer inputNode.Close()
 
 	isInputNode := inputNode.IsInputNode()

internal/util/flowgraph/message.go

@@ -32,6 +32,7 @@ type MsgStreamMsg struct {
 	timestampMax   Timestamp
 	startPositions []*MsgPosition
 	endPositions   []*MsgPosition
+	isCloseMsg     bool
 }
 
 // GenerateMsgStreamMsg is used to create a new MsgStreamMsg object

internal/util/flowgraph/node.go

@@ -59,23 +59,28 @@ type nodeCtx struct {
 	downstreamInputChanIdx map[string]int
 
 	closeCh chan struct{} // notify work to exit
-	closeWg sync.WaitGroup // block Close until work exit
 }
 
 // Start invoke Node `Start` method and start a worker goroutine
 func (nodeCtx *nodeCtx) Start() {
 	nodeCtx.node.Start()
 
-	nodeCtx.closeWg.Add(1)
 	go nodeCtx.work()
 }
 
+func isCloseMsg(msgs []Msg) bool {
+	if len(msgs) == 1 {
+		msg, ok := msgs[0].(*MsgStreamMsg)
+		return ok && msg.isCloseMsg
+	}
+	return false
+}
+
 // work handles node work spinning
 // 1. collectMessage from upstream or just produce Msg from InputNode
 // 2. invoke node.Operate
 // 3. deliver the Operate result to downstream nodes
 func (nodeCtx *nodeCtx) work() {
-	defer nodeCtx.closeWg.Done()
 	name := fmt.Sprintf("nodeCtxTtChecker-%s", nodeCtx.node.Name())
 	var checker *timerecord.GroupChecker
 	if enableTtChecker {
@@ -98,8 +103,19 @@ func (nodeCtx *nodeCtx) work() {
 				nodeCtx.collectInputMessages()
 				inputs = nodeCtx.inputMessages
 			}
+			// the input message decides whether the operate method is executed
+			if isCloseMsg(inputs) {
+				res = inputs
+			}
+			if len(res) == 0 {
 				n := nodeCtx.node
 				res = n.Operate(inputs)
+			}
+			// the res decide whether the node should be closed.
+			if isCloseMsg(res) {
+				close(nodeCtx.closeCh)
+				nodeCtx.node.Close()
+			}
 
 			if enableTtChecker {
 				checker.Check(name)
@@ -127,13 +143,7 @@ func (nodeCtx *nodeCtx) work() {
 // Close handles cleanup logic and notify worker to quit
 func (nodeCtx *nodeCtx) Close() {
 	if nodeCtx.node.IsInputNode() {
-		nodeCtx.node.Close() // close input msgStream
+		nodeCtx.node.Close()
-		close(nodeCtx.closeCh)
-		nodeCtx.closeWg.Wait()
-	} else {
-		close(nodeCtx.closeCh)
-		nodeCtx.closeWg.Wait()
-		nodeCtx.node.Close() // close output msgStream, and etc...
 	}
 }
 
@@ -146,10 +156,7 @@ func (nodeCtx *nodeCtx) deliverMsg(wg *sync.WaitGroup, msg Msg, inputChanIdx int
 			log.Warn(fmt.Sprintln(err))
 		}
 	}()
-	select {
+	nodeCtx.inputChannels[inputChanIdx] <- msg
-	case <-nodeCtx.closeCh:
-	case nodeCtx.inputChannels[inputChanIdx] <- msg:
-	}
 }
 
 func (nodeCtx *nodeCtx) collectInputMessages() {
@@ -161,10 +168,7 @@ func (nodeCtx *nodeCtx) collectInputMessages() {
 	// and move them to inputMessages.
 	for i := 0; i < inputsNum; i++ {
 		channel := nodeCtx.inputChannels[i]
-		select {
+		msg, ok := <-channel
-		case <-nodeCtx.closeCh:
-			return
-		case msg, ok := <-channel:
 		if !ok {
 			// TODO: add status
 			log.Warn("input channel closed")
@@ -172,7 +176,6 @@ func (nodeCtx *nodeCtx) collectInputMessages() {
 		}
 		nodeCtx.inputMessages[i] = msg
 	}
-	}
 
 	// timeTick alignment check
 	if len(nodeCtx.inputMessages) > 1 {
@@ -191,10 +194,7 @@ func (nodeCtx *nodeCtx) collectInputMessages() {
 				for nodeCtx.inputMessages[i].TimeTick() != latestTime {
 					log.Debug("Try to align timestamp", zap.Uint64("t1", latestTime), zap.Uint64("t2", nodeCtx.inputMessages[i].TimeTick()))
 					channel := nodeCtx.inputChannels[i]
-					select {
+					msg, ok := <-channel
-					case <-nodeCtx.closeCh:
-						return
-					case msg, ok := <-channel:
 					if !ok {
 						log.Warn("input channel closed")
 						return
@@ -202,7 +202,6 @@ func (nodeCtx *nodeCtx) collectInputMessages() {
 					nodeCtx.inputMessages[i] = msg
 				}
 			}
-			}
 			sign <- struct{}{}
 		}()
 
@@ -210,7 +209,6 @@ func (nodeCtx *nodeCtx) collectInputMessages() {
 		case <-time.After(10 * time.Second):
 			panic("Fatal, misaligned time tick, please restart pulsar")
 		case <-sign:
-		case <-nodeCtx.closeCh:
 		}
 	}
 }

internal/util/flowgraph/node_test.go

@@ -73,10 +73,7 @@ func TestNodeCtx_Start(t *testing.T) {
 	produceStream.Produce(&msgPack)
 
 	nodeName := "input_node"
-	inputNode := &InputNode{
+	inputNode := NewInputNode(msgStream, nodeName, 100, 100)
-		inStream: msgStream,
-		name:     nodeName,
-	}
 
 	node := &nodeCtx{
 		node:                   inputNode,