Merge pull request #2 from RyouZhang/dev

Dev
RyouZhang · Jul 11, 2024 · bb8da63 · bb8da63
2 parents bf64af6 + 509b413
commit bb8da63
Show file tree

Hide file tree

Showing 12 changed files with 121 additions and 38 deletions.
diff --git a/base.go b/base.go
@@ -2,6 +2,7 @@ package flow
 
 type IBrick interface {
 	Name() string
+	AddLifeCycle(l ILifeCycle)
 }
 
 type IInput interface {
@@ -24,3 +25,9 @@ type IError interface {
 type IRoute interface {
 	RouteOutput(func(*Message) bool) <-chan *Message
 }
+
+type ILifeCycle interface {
+	Add(int)
+	Done()
+	Wait()
+}
diff --git a/base_brick.go b/base_brick.go
@@ -5,7 +5,9 @@ import (
 )
 
 type BaseBrick struct {
-	name     string
+	name string
+	lc   ILifeCycle
+
 	kernal   func(*Message) (*Message, error)
 	errQueue chan error
 	outQueue chan *Message
@@ -15,6 +17,10 @@ func (b *BaseBrick) Name() string {
 	return b.name
 }
 
+func (b *BaseBrick) AddLifeCycle(lc ILifeCycle) {
+	b.lc = lc
+}
+
 func (b *BaseBrick) Linked(inQueue <-chan *Message) {
 	b.loop(inQueue)
 }
@@ -31,6 +37,7 @@ func (b *BaseBrick) loop(inQueue <-chan *Message) {
 	defer func() {
 		close(b.errQueue)
 		close(b.outQueue)
+		b.lc.Done()
 	}()
 	for msg := range inQueue {
 		res, err := async.Safety(func() (interface{}, error) {

diff --git a/board.go b/board.go
@@ -8,15 +8,27 @@ import (
 
 type Board struct {
 	name       string
-	wg         sync.WaitGroup
 	bricks     map[string]IBrick
 	errHandler func(string, error)
+	lc         ILifeCycle
 }
 
 func NewBoard(name string) *Board {
 	return &Board{
 		name:   name,
 		bricks: make(map[string]IBrick),
+		lc:     &sync.WaitGroup{},
+	}
+}
+
+func NewBoardWithLifeCycle(name string, lc ILifeCycle) *Board {
+	if lc == nil {
+		return nil
+	}
+	return &Board{
+		name:   name,
+		bricks: make(map[string]IBrick),
+		lc:     lc,
 	}
 }
 
@@ -29,6 +41,9 @@ func (b *Board) Add(bricks ...IBrick) *Board {
 		_, ok := b.bricks[brick.Name()]
 		if false == ok {
 			b.bricks[brick.Name()] = brick
+			// add life cycle
+			b.lc.Add(1)
+			brick.AddLifeCycle(b.lc)
 
 			if _, ok := brick.(IError); ok {
 				go b.onError(brick.(IBrick).Name(), brick.(IError).Errors())
@@ -49,19 +64,10 @@ func (b *Board) Connect(outName string, inName string) *Board {
 	if false == ok {
 		panic(errors.New(fmt.Sprintf("Invalid Brick %s", inName)))
 	}
-	b.connectBrick(out.(IOutput), in.(IInput))
+	in.(IInput).Linked(out.(IOutput).Output())
 	return b
 }
 
-func (b *Board) connectBrick(out IOutput, in IInput) {
-	b.wg.Add(1)
-	target := out.Output()
-	go func() {
-		defer b.wg.Done()
-		in.Linked(target)
-	}()
-}
-
 func (b *Board) RouteConnect(outName string, inName string, method func(*Message) bool) *Board {
 	out, ok := b.bricks[outName]
 	if false == ok {
@@ -71,26 +77,15 @@ func (b *Board) RouteConnect(outName string, inName string, method func(*Message
 	if false == ok {
 		panic(errors.New(fmt.Sprintf("Invalid Brick %s", inName)))
 	}
-	b.routeConnectBrick(out.(IRoute), in.(IInput), method)
+	in.(IInput).Linked(out.(IRoute).RouteOutput(method))
 	return b
 }
 
-func (b *Board) routeConnectBrick(out IRoute, in IInput, method func(*Message) bool) {
-	b.wg.Add(1)
-	target := out.RouteOutput(method)
-	go func() {
-		defer b.wg.Done()
-		in.Linked(target)
-	}()
-}
-
 func (b *Board) Start() {
 	for _, brick := range b.bricks {
 		ob, ok := brick.(IEntry)
 		if ok {
-			b.wg.Add(1)
 			go func() {
-				defer b.wg.Done()
 				ob.Start()
 			}()
 		}
@@ -104,7 +99,7 @@ func (b *Board) Stop() {
 			ob.Stop()
 		}
 	}
-	b.wg.Wait()
+	b.lc.Wait()
 }
 
 func (b *Board) onError(name string, inQueue <-chan error) {

diff --git a/input_brick.go b/input_brick.go
@@ -5,7 +5,9 @@ import (
 )
 
 type InputBrick struct {
-	name     string
+	name string
+	lc   ILifeCycle
+
 	kernal   func(chan<- *Message, chan<- error, <-chan bool)
 	shutdown chan bool
 
@@ -21,6 +23,10 @@ func (b *InputBrick) Name() string {
 	return b.name
 }
 
+func (b *InputBrick) AddLifeCycle(lc ILifeCycle) {
+	b.lc = lc
+}
+
 func (b *InputBrick) Output() <-chan *Message {
 	b.useDefaultOut = true
 	return b.outQueue
@@ -44,6 +50,9 @@ func (b *InputBrick) Start() {
 }
 
 func (b *InputBrick) Stop() {
+	defer func() {
+		b.lc.Done()
+	}()
 	b.shutdown <- true
 	close(b.shutdown)
 }

diff --git a/logic_brick.go b/logic_brick.go
@@ -7,7 +7,9 @@ import (
 )
 
 type LogicBrick struct {
-	name    string
+	name string
+	lc   ILifeCycle
+
 	workers chan bool
 	wg      sync.WaitGroup
 	kernal  func(*Message, chan<- *Message) error
@@ -30,6 +32,10 @@ func (b *LogicBrick) Name() string {
 	return b.name
 }
 
+func (b *LogicBrick) AddLifeCycle(lc ILifeCycle) {
+	b.lc = lc
+}
+
 func (b *LogicBrick) Linked(queue <-chan *Message) {
 	b.inMux.Add(1)
 	go func() {
@@ -114,6 +120,7 @@ func (b *LogicBrick) pump() {
 	}
 	close(b.outQueue)
 	close(b.errQueue)
+	b.lc.Done()
 }
 
 func NewLogicBrick(

diff --git a/merge_brick.go b/merge_brick.go
@@ -5,7 +5,9 @@ import (
 )
 
 type MergeBrick struct {
-	name     string
+	name string
+	lc   ILifeCycle
+
 	once     sync.Once
 	wg       sync.WaitGroup
 	outQueue chan *Message
@@ -15,14 +17,21 @@ func (b *MergeBrick) Name() string {
 	return b.name
 }
 
+func (b *MergeBrick) AddLifeCycle(lc ILifeCycle) {
+	b.lc = lc
+}
+
 func (b *MergeBrick) Output() <-chan *Message {
 	return b.outQueue
 }
 
 func (b *MergeBrick) Linked(inQueue <-chan *Message) {
 	b.wg.Add(1)
 	go func() {
-		defer b.wg.Done()
+		defer func() {
+			b.wg.Done()
+			b.lc.Done()
+		}()
 		for msg := range inQueue {
 			b.outQueue <- msg
 		}

diff --git a/output_brick.go b/output_brick.go
@@ -5,7 +5,9 @@ import (
 )
 
 type OutputBrick struct {
-	name         string
+	name string
+	lc   ILifeCycle
+
 	kernal       func(*Message) error
 	gracefulStop func()
 	errQueue     chan error
@@ -15,8 +17,12 @@ func (b *OutputBrick) Name() string {
 	return b.name
 }
 
+func (b *OutputBrick) AddLifeCycle(lc ILifeCycle) {
+	b.lc = lc
+}
+
 func (b *OutputBrick) Linked(inQueue <-chan *Message) {
-	b.loop(inQueue)
+	go b.loop(inQueue)
 }
 
 func (b *OutputBrick) Errors() <-chan error {
@@ -26,6 +32,7 @@ func (b *OutputBrick) Errors() <-chan error {
 func (b *OutputBrick) loop(inQueue <-chan *Message) {
 	defer func() {
 		close(b.errQueue)
+		b.lc.Done()
 	}()
 	for msg := range inQueue {
 		_, err := async.Safety(func() (interface{}, error) {

diff --git a/parallel_brick.go b/parallel_brick.go
@@ -7,7 +7,9 @@ import (
 )
 
 type ParallelBrick struct {
-	name      string
+	name string
+	lc   ILifeCycle
+
 	hash      func(*Message) int
 	subQueues []chan *Message
 	wg        sync.WaitGroup
@@ -31,6 +33,10 @@ func (b *ParallelBrick) Name() string {
 	return b.name
 }
 
+func (b *ParallelBrick) AddLifeCycle(lc ILifeCycle) {
+	b.lc = lc
+}
+
 func (b *ParallelBrick) Linked(queue <-chan *Message) {
 	b.inMux.Add(1)
 	go func() {
@@ -107,6 +113,7 @@ func (b *ParallelBrick) pump() {
 	}
 	close(b.outQueue)
 	close(b.errQueue)
+	b.lc.Done()
 }
 
 func (b *ParallelBrick) worker(queue <-chan *Message) {

diff --git a/priority_brick.go b/priority_brick.go
@@ -8,7 +8,9 @@ import (
 )
 
 type PriorityBrick struct {
-	name    string
+	name string
+	lc   ILifeCycle
+
 	workers chan bool
 	wg      sync.WaitGroup
 	kernal  func(*Message, chan<- *Message) error
@@ -31,6 +33,10 @@ func (b *PriorityBrick) Name() string {
 	return b.name
 }
 
+func (b *PriorityBrick) AddLifeCycle(lc ILifeCycle) {
+	b.lc = lc
+}
+
 func (b *PriorityBrick) Linked(queue <-chan *Message) {
 	b.inQueueMux.Lock()
 	defer b.inQueueMux.Unlock()
@@ -86,6 +92,9 @@ func (b *PriorityBrick) handler(queue <-chan *Message) error {
 }
 
 func (b *PriorityBrick) loop() {
+	defer func() {
+		close(b.resQueue)
+	}()
 	flagKey := 0
 	for {
 	PULL:
@@ -113,6 +122,9 @@ func (b *PriorityBrick) loop() {
 }
 
 func (b *PriorityBrick) pump() {
+	defer func() {
+		b.lc.Done()
+	}()
 	for msg := range b.resQueue {
 		flag := false
 		for _, item := range b.outQueues {

diff --git a/reduce_brick.go b/reduce_brick.go
@@ -7,7 +7,9 @@ import (
 )
 
 type ReduceBrick struct {
-	name          string
+	name string
+	lc   ILifeCycle
+
 	windowSeconds int //sec
 	msgKeys       []string
 	msgDic        map[string][]*Message
@@ -22,8 +24,12 @@ func (b *ReduceBrick) Name() string {
 	return b.name
 }
 
+func (b *ReduceBrick) AddLifeCycle(lc ILifeCycle) {
+	b.lc = lc
+}
+
 func (b *ReduceBrick) Linked(inQueue <-chan *Message) {
-	b.loop(inQueue)
+	go b.loop(inQueue)
 }
 
 func (b *ReduceBrick) Output() <-chan *Message {
@@ -130,6 +136,7 @@ End:
 	}
 	close(b.outQueue)
 	close(b.errQueue)
+	b.lc.Done()
 }
 
 func NewReduceBrick(