package/mq
1
0
Fork 0
mirror of https://github.com/oarkflow/mq.git synced 2025-10-06 00:16:49 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

init: publisher

Browse Source
This commit is contained in:
sujit
2024-09-29 16:37:52 +05:45
parent 32b042be67
commit 43f06156bb
1 changed files with 11 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
dag/dag.go
View File

@@ -20,7 +20,7 @@ type DAG struct {
server *mq.Broker server *mq.Broker
nodes map[string]*mq.Consumer nodes map[string]*mq.Consumer
edges map[string][]string edges map[string][]string
loopEdges map[string]string loopEdges map[string][]string
taskChMap map[string]chan mq.Result taskChMap map[string]chan mq.Result
taskResults map[string]map[string]*taskContext taskResults map[string]map[string]*taskContext
mu sync.Mutex mu sync.Mutex
@@ -30,7 +30,7 @@ func New(opts ...mq.Option) *DAG {
d := &DAG{ d := &DAG{
nodes: make(map[string]*mq.Consumer), nodes: make(map[string]*mq.Consumer),
edges: make(map[string][]string), edges: make(map[string][]string),
loopEdges: make(map[string]string), loopEdges: make(map[string][]string),
taskChMap: make(map[string]chan mq.Result), taskChMap: make(map[string]chan mq.Result),
taskResults: make(map[string]map[string]*taskContext), taskResults: make(map[string]map[string]*taskContext),
} }
@@ -49,7 +49,7 @@ func (d *DAG) AddEdge(fromNode string, toNodes ...string) {
d.edges[fromNode] = toNodes d.edges[fromNode] = toNodes
} }
func (d *DAG) AddLoop(fromNode string, toNode string) { func (d *DAG) AddLoop(fromNode string, toNode ...string) {
d.loopEdges[fromNode] = toNode d.loopEdges[fromNode] = toNode
} }
@@ -119,7 +119,7 @@ func (d *DAG) TaskCallback(ctx context.Context, task *mq.Task) error {
} else { } else {
payload = task.Result payload = task.Result
} }
if loopNode, exists := d.loopEdges[task.CurrentQueue]; exists { if loopNodes, exists := d.loopEdges[task.CurrentQueue]; exists {
var items []json.RawMessage var items []json.RawMessage
if err := json.Unmarshal(payload, &items); err != nil { if err := json.Unmarshal(payload, &items); err != nil {
return err return err
@@ -132,12 +132,15 @@ func (d *DAG) TaskCallback(ctx context.Context, task *mq.Task) error {
} }
ctx = mq.SetHeaders(ctx, map[string]string{mq.TriggerNode: task.CurrentQueue}) ctx = mq.SetHeaders(ctx, map[string]string{mq.TriggerNode: task.CurrentQueue})
for _, item := range items { for _, loopNode := range loopNodes {
_, err := d.PublishTask(ctx, item, loopNode, task.ID) for _, item := range items {
if err != nil { _, err := d.PublishTask(ctx, item, loopNode, task.ID)
return err if err != nil {
return err
}
} }
} }
return nil return nil
} }
if nodeType == "loop" && completed { if nodeType == "loop" && completed {