Files
mq/examples/dag.go
2024-09-26 19:10:36 +05:45

81 lines
2.4 KiB
Go

package main
import (
"context"
"encoding/json"
"fmt"
"github.com/oarkflow/mq"
"time"
)
func handleNode1(_ context.Context, task mq.Task) mq.Result {
result := []map[string]string{
{"field": "facility", "item": "item1"},
{"field": "facility", "item": "item2"},
{"field": "facility", "item": "item3"},
}
var payload string
err := json.Unmarshal(task.Payload, &payload)
if err != nil {
return mq.Result{Status: "fail", Payload: json.RawMessage(`{"field": "node1", "item": "error"}`)}
}
fmt.Printf("Processing task at node1: %s\n", string(task.Payload))
bt, _ := json.Marshal(result)
return mq.Result{Status: "completed", Payload: bt}
}
func handleNode2(_ context.Context, task mq.Task) mq.Result {
var payload map[string]string
err := json.Unmarshal(task.Payload, &payload)
if err != nil {
return mq.Result{Status: "fail", Payload: json.RawMessage(`{"field": "node2", "item": "error"}`)}
}
status := "fail"
if payload["item"] == "item2" {
status = "pass"
}
fmt.Printf("Processing task at node2: %s %s\n", payload, status)
bt, _ := json.Marshal(payload)
return mq.Result{Status: status, Payload: bt}
}
func handleNode3(_ context.Context, task mq.Task) mq.Result {
result := `{"field": "node3", "item": %s}`
fmt.Printf("Processing task at node3: %s\n", string(task.Payload))
return mq.Result{Status: "completed", Payload: json.RawMessage(fmt.Sprintf(result, string(task.Payload)))}
}
func handleNode4(_ context.Context, task mq.Task) mq.Result {
result := `{"field": "node4", "item": %s}`
fmt.Printf("Processing task at node4: %s\n", string(task.Payload))
return mq.Result{Status: "completed", Payload: json.RawMessage(fmt.Sprintf(result, string(task.Payload)))}
}
func main() {
ctx := context.Background()
d := mq.NewDAG(":8082", false)
d.AddNode("node1", handleNode1, true)
d.AddNode("node2", handleNode2)
d.AddNode("node3", handleNode3)
d.AddNode("node4", handleNode4)
d.AddCondition("node2", map[string]string{"pass": "node3", "fail": "node4"})
err := d.AddLoop("node1", "node2")
if err != nil {
panic(err)
}
err = d.Prepare(ctx)
if err != nil {
panic(err)
}
// Start the DAG and process the task
go func() {
if err := d.Start(ctx); err != nil {
fmt.Println("Error starting DAG:", err)
}
}()
result := d.ProcessTask(ctx, mq.Task{Payload: []byte(`"Start processing"`)})
fmt.Println(string(result.Payload))
time.Sleep(50 * time.Second)
}