mirror of
https://github.com/oarkflow/mq.git
synced 2025-11-02 20:04:02 +08:00
feat: separate broker
This commit is contained in:
Oarkflow
@@ -4,272 +4,58 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
v2 "github.com/oarkflow/mq/v2"
|
||||
)
|
||||
|
||||
type Task struct {
|
||||
ID string `json:"id"`
|
||||
Payload json.RawMessage `json:"payload"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
ProcessedAt time.Time `json:"processed_at"`
|
||||
Status string `json:"status"`
|
||||
Error error `json:"error"`
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
Payload json.RawMessage `json:"payload"`
|
||||
Queue string `json:"queue"`
|
||||
MessageID string `json:"message_id"`
|
||||
Error error `json:"error,omitempty"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
const (
|
||||
SimpleEdge = iota
|
||||
LoopEdge
|
||||
ConditionEdge
|
||||
)
|
||||
|
||||
type Edge struct {
|
||||
edgeType int
|
||||
to string
|
||||
conditions map[string]string
|
||||
}
|
||||
|
||||
type Node struct {
|
||||
key string
|
||||
handler func(context.Context, Task) Result
|
||||
edges []Edge
|
||||
}
|
||||
|
||||
type RadixTrie struct {
|
||||
children map[rune]*RadixTrie
|
||||
node *Node
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewRadixTrie() *RadixTrie {
|
||||
return &RadixTrie{
|
||||
children: make(map[rune]*RadixTrie),
|
||||
func handler1(ctx context.Context, task *v2.Task) v2.Result {
|
||||
return v2.Result{
|
||||
TaskID: task.ID,
|
||||
NodeKey: "A",
|
||||
Payload: task.Payload,
|
||||
Status: "success",
|
||||
}
|
||||
}
|
||||
|
||||
func (trie *RadixTrie) Insert(key string, node *Node) {
|
||||
trie.mu.Lock()
|
||||
defer trie.mu.Unlock()
|
||||
|
||||
current := trie
|
||||
for _, char := range key {
|
||||
if _, exists := current.children[char]; !exists {
|
||||
current.children[char] = NewRadixTrie()
|
||||
}
|
||||
current = current.children[char]
|
||||
}
|
||||
current.node = node
|
||||
}
|
||||
|
||||
func (trie *RadixTrie) Search(key string) (*Node, bool) {
|
||||
trie.mu.RLock()
|
||||
defer trie.mu.RUnlock()
|
||||
current := trie
|
||||
for _, char := range key {
|
||||
if _, exists := current.children[char]; !exists {
|
||||
return nil, false
|
||||
}
|
||||
current = current.children[char]
|
||||
}
|
||||
if current.node != nil {
|
||||
return current.node, true
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
type DAG struct {
|
||||
trie *RadixTrie
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewDAG() *DAG {
|
||||
return &DAG{
|
||||
trie: NewRadixTrie(),
|
||||
func handler2(ctx context.Context, task *v2.Task) v2.Result {
|
||||
var user map[string]any
|
||||
json.Unmarshal(task.Payload, &user)
|
||||
return v2.Result{
|
||||
TaskID: task.ID,
|
||||
NodeKey: "B",
|
||||
Payload: task.Payload,
|
||||
Status: "success",
|
||||
}
|
||||
}
|
||||
|
||||
func (d *DAG) AddNode(key string, handler func(context.Context, Task) Result, isRoot ...bool) {
|
||||
node := &Node{key: key, handler: handler}
|
||||
d.trie.Insert(key, node)
|
||||
}
|
||||
|
||||
func (d *DAG) AddEdge(fromKey string, toKey string) {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
node, exists := d.trie.Search(fromKey)
|
||||
if !exists {
|
||||
fmt.Printf("Node %s not found to add edge.\n", fromKey)
|
||||
return
|
||||
func handler3(ctx context.Context, task *v2.Task) v2.Result {
|
||||
var user map[string]any
|
||||
json.Unmarshal(task.Payload, &user)
|
||||
age := int(user["age"].(float64))
|
||||
status := "FAIL"
|
||||
if age > 20 {
|
||||
status = "PASS"
|
||||
}
|
||||
edge := Edge{edgeType: SimpleEdge, to: toKey}
|
||||
node.edges = append(node.edges, edge)
|
||||
}
|
||||
|
||||
func (d *DAG) AddLoop(fromKey string, toKey string) {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
node, exists := d.trie.Search(fromKey)
|
||||
if !exists {
|
||||
fmt.Printf("Node %s not found to add loop edge.\n", fromKey)
|
||||
return
|
||||
user["status"] = status
|
||||
resultPayload, _ := json.Marshal(user)
|
||||
return v2.Result{
|
||||
TaskID: task.ID,
|
||||
NodeKey: "C",
|
||||
Payload: resultPayload,
|
||||
Status: status,
|
||||
}
|
||||
edge := Edge{edgeType: LoopEdge, to: toKey}
|
||||
node.edges = append(node.edges, edge)
|
||||
}
|
||||
|
||||
func (d *DAG) AddCondition(fromKey string, conditions map[string]string) {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
node, exists := d.trie.Search(fromKey)
|
||||
if !exists {
|
||||
fmt.Printf("Node %s not found to add condition edge.\n", fromKey)
|
||||
return
|
||||
}
|
||||
edge := Edge{edgeType: ConditionEdge, conditions: conditions}
|
||||
node.edges = append(node.edges, edge)
|
||||
}
|
||||
|
||||
type ProcessCallback func(ctx context.Context, key string, result Result) string
|
||||
|
||||
func (d *DAG) ProcessTask(ctx context.Context, key string, task Task) {
|
||||
node, exists := d.trie.Search(key)
|
||||
if !exists {
|
||||
fmt.Printf("Node %s not found.\n", key)
|
||||
return
|
||||
}
|
||||
result := node.handler(ctx, task)
|
||||
nextKey := d.callback(ctx, key, result)
|
||||
if nextKey != "" {
|
||||
d.ProcessTask(ctx, nextKey, task)
|
||||
}
|
||||
}
|
||||
|
||||
func (d *DAG) ProcessLoop(ctx context.Context, key string, task Task) {
|
||||
_, exists := d.trie.Search(key)
|
||||
if !exists {
|
||||
fmt.Printf("Node %s not found.\n", key)
|
||||
return
|
||||
}
|
||||
var items []json.RawMessage
|
||||
err := json.Unmarshal(task.Payload, &items)
|
||||
if err != nil {
|
||||
fmt.Printf("Error unmarshaling payload as slice: %v\n", err)
|
||||
return
|
||||
}
|
||||
for _, item := range items {
|
||||
newTask := Task{
|
||||
ID: task.ID,
|
||||
Payload: item,
|
||||
}
|
||||
|
||||
d.ProcessTask(ctx, key, newTask)
|
||||
}
|
||||
}
|
||||
|
||||
func (d *DAG) callback(ctx context.Context, currentKey string, result Result) string {
|
||||
fmt.Printf("Callback received result from %s: %s\n", currentKey, string(result.Payload))
|
||||
node, exists := d.trie.Search(currentKey)
|
||||
if !exists {
|
||||
return ""
|
||||
}
|
||||
for _, edge := range node.edges {
|
||||
switch edge.edgeType {
|
||||
case SimpleEdge:
|
||||
return edge.to
|
||||
case LoopEdge:
|
||||
|
||||
d.ProcessLoop(ctx, edge.to, Task{Payload: result.Payload})
|
||||
return ""
|
||||
case ConditionEdge:
|
||||
if nextKey, conditionMet := edge.conditions[result.Status]; conditionMet {
|
||||
return nextKey
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func Node1(ctx context.Context, task Task) Result {
|
||||
return Result{Payload: task.Payload, MessageID: task.ID}
|
||||
}
|
||||
|
||||
func Node2(ctx context.Context, task Task) Result {
|
||||
return Result{Payload: task.Payload, MessageID: task.ID}
|
||||
}
|
||||
|
||||
func Node3(ctx context.Context, task Task) Result {
|
||||
var data map[string]any
|
||||
err := json.Unmarshal(task.Payload, &data)
|
||||
if err != nil {
|
||||
return Result{Error: err}
|
||||
}
|
||||
data["salary"] = fmt.Sprintf("12000%v", data["user_id"])
|
||||
bt, _ := json.Marshal(data)
|
||||
return Result{Payload: bt, MessageID: task.ID}
|
||||
}
|
||||
|
||||
func Node4(ctx context.Context, task Task) Result {
|
||||
var data []map[string]any
|
||||
err := json.Unmarshal(task.Payload, &data)
|
||||
if err != nil {
|
||||
return Result{Error: err}
|
||||
}
|
||||
payload := map[string]any{"storage": data}
|
||||
bt, _ := json.Marshal(payload)
|
||||
return Result{Payload: bt, MessageID: task.ID}
|
||||
}
|
||||
|
||||
func CheckCondition(ctx context.Context, task Task) Result {
|
||||
var data map[string]any
|
||||
err := json.Unmarshal(task.Payload, &data)
|
||||
if err != nil {
|
||||
return Result{Error: err}
|
||||
}
|
||||
var status string
|
||||
if data["user_id"].(float64) == 2 {
|
||||
status = "pass"
|
||||
} else {
|
||||
status = "fail"
|
||||
}
|
||||
return Result{Status: status, Payload: task.Payload, MessageID: task.ID}
|
||||
}
|
||||
|
||||
func Pass(ctx context.Context, task Task) Result {
|
||||
fmt.Println("Pass")
|
||||
return Result{Payload: task.Payload}
|
||||
}
|
||||
|
||||
func Fail(ctx context.Context, task Task) Result {
|
||||
fmt.Println("Fail")
|
||||
return Result{Payload: []byte(`{"test2": "asdsa"}`)}
|
||||
}
|
||||
|
||||
func main() {
|
||||
dag := NewDAG()
|
||||
dag.AddNode("queue1", Node1, true)
|
||||
dag.AddNode("queue2", Node2)
|
||||
dag.AddNode("queue3", Node3)
|
||||
dag.AddNode("queue4", Node4)
|
||||
dag.AddNode("queue5", CheckCondition)
|
||||
dag.AddNode("queue6", Pass)
|
||||
dag.AddNode("queue7", Fail)
|
||||
dag.AddEdge("queue1", "queue2")
|
||||
dag.AddEdge("queue2", "queue4")
|
||||
dag.AddEdge("queue3", "queue5")
|
||||
dag.AddLoop("queue2", "queue3")
|
||||
dag.AddCondition("queue5", map[string]string{"pass": "queue6", "fail": "queue7"})
|
||||
ctx := context.Background()
|
||||
task := Task{
|
||||
ID: "task1",
|
||||
Payload: []byte(`[{"user_id": 1}, {"user_id": 2}]`),
|
||||
}
|
||||
dag.ProcessTask(ctx, "queue1", task)
|
||||
dag := v2.NewDAG()
|
||||
dag.AddNode("A", handler1)
|
||||
dag.AddNode("B", handler2)
|
||||
dag.AddNode("C", handler3)
|
||||
dag.AddEdge("A", "B", v2.LoopEdge)
|
||||
dag.AddEdge("B", "C", v2.SimpleEdge)
|
||||
initialPayload, _ := json.Marshal([]map[string]any{
|
||||
{"user_id": 1, "age": 12},
|
||||
{"user_id": 2, "age": 34},
|
||||
})
|
||||
rs := dag.ProcessTask(context.Background(), "A", initialPayload)
|
||||
fmt.Println(string(rs.Payload))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user