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

@@ -28,8 +28,9 @@ import (
 // InputNode is the entry point of flowgragh
 type InputNode struct {
 	BaseNode
-	inStream msgstream.MsgStream
+	inStream     msgstream.MsgStream
-	name     string
+	name         string
+	closeMsgChan chan struct{}
 }
 
 // IsInputNode returns whether Node is InputNode
@@ -44,10 +45,15 @@ func (inNode *InputNode) Start() {
 
 // Close implements node
 func (inNode *InputNode) Close() {
-	inNode.inStream.Close()
+	select {
-	log.Debug("message stream closed",
+	case <-inNode.closeMsgChan:
-		zap.String("node name", inNode.name),
+		return
-	)
+	default:
+		close(inNode.closeMsgChan)
+		log.Debug("message stream closed",
+			zap.String("node name", inNode.name),
+		)
+	}
 }
 
 // Name returns node name
@@ -62,37 +68,44 @@ 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 {
-	if !ok {
+	case <-inNode.closeMsgChan:
-		log.Warn("MsgStream closed", zap.Any("input node", inNode.Name()))
+		inNode.inStream.Close()
-		return []Msg{}
+		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{}
+		}
 
-	// TODO: add status
+		// TODO: add status
-	if msgPack == nil {
+		if msgPack == nil {
-		return nil
+			return nil
-	}
+		}
-	var spans []opentracing.Span
+		var spans []opentracing.Span
-	for _, msg := range msgPack.Msgs {
+		for _, msg := range msgPack.Msgs {
-		sp, ctx := trace.StartSpanFromContext(msg.TraceCtx())
+			sp, ctx := trace.StartSpanFromContext(msg.TraceCtx())
-		sp.LogFields(oplog.String("input_node name", inNode.Name()))
+			sp.LogFields(oplog.String("input_node name", inNode.Name()))
-		spans = append(spans, sp)
+			spans = append(spans, sp)
-		msg.SetTraceCtx(ctx)
+			msg.SetTraceCtx(ctx)
-	}
+		}
 
-	var msgStreamMsg Msg = &MsgStreamMsg{
+		var msgStreamMsg Msg = &MsgStreamMsg{
-		tsMessages:     msgPack.Msgs,
+			tsMessages:     msgPack.Msgs,
-		timestampMin:   msgPack.BeginTs,
+			timestampMin:   msgPack.BeginTs,
-		timestampMax:   msgPack.EndTs,
+			timestampMax:   msgPack.EndTs,
-		startPositions: msgPack.StartPositions,
+			startPositions: msgPack.StartPositions,
-		endPositions:   msgPack.EndPositions,
+			endPositions:   msgPack.EndPositions,
-	}
+		}
 
-	for _, span := range spans {
+		for _, span := range spans {
-		span.Finish()
+			span.Finish()
-	}
+		}
 
-	return []Msg{msgStreamMsg}
+		return []Msg{msgStreamMsg}
+	}
 }
 
 // NewInputNode composes an InputNode with provided MsgStream, name and parameters
@@ -102,8 +115,9 @@ func NewInputNode(inStream msgstream.MsgStream, nodeName string, maxQueueLength
 	baseNode.SetMaxParallelism(maxParallelism)
 
 	return &InputNode{
-		BaseNode: baseNode,
+		BaseNode:     baseNode,
-		inStream: inStream,
+		inStream:     inStream,
-		name:     nodeName,
+		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

@@ -58,24 +58,29 @@ type nodeCtx struct {
 	downstream             []*nodeCtx
 	downstreamInputChanIdx map[string]int
 
-	closeCh chan struct{}  // notify work to exit
+	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
 			}
-			n := nodeCtx.node
+			// the input message decides whether the operate method is executed
-			res = n.Operate(inputs)
+			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,17 +168,13 @@ 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:
+		if !ok {
+			// TODO: add status
+			log.Warn("input channel closed")
 			return
-		case msg, ok := <-channel:
-			if !ok {
-				// TODO: add status
-				log.Warn("input channel closed")
-				return
-			}
-			nodeCtx.inputMessages[i] = msg
 		}
+		nodeCtx.inputMessages[i] = msg
 	}
 
 	// timeTick alignment check
@@ -191,16 +194,12 @@ 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:
+					if !ok {
+						log.Warn("input channel closed")
 						return
-					case msg, ok := <-channel:
-						if !ok {
-							log.Warn("input channel closed")
-							return
-						}
-						nodeCtx.inputMessages[i] = msg
 					}
+					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,