feat: add example

2025-10-16 01:00:36 +08:00 · 2024-10-08 16:24:52 +05:45
parent a4dff9c77b
commit 612eb535ec
4 changed files with 87 additions and 99 deletions
--- a/examples/dag_v2.go
+++ b/examples/dag_v2.go
@@ -5,20 +5,21 @@ import (
 	"encoding/json"
 	"fmt"

+	"github.com/oarkflow/mq"
 	"github.com/oarkflow/mq/v2"
 )

-func handler1(ctx context.Context, task *v2.Task) v2.Result {
-	return v2.Result{Payload: task.Payload, Ctx: ctx}
+func handler1(ctx context.Context, task *mq.Task) mq.Result {
+	return mq.Result{Payload: task.Payload, Ctx: ctx}
 }

-func handler2(ctx context.Context, task *v2.Task) v2.Result {
+func handler2(ctx context.Context, task *mq.Task) mq.Result {
 	var user map[string]any
 	json.Unmarshal(task.Payload, &user)
-	return v2.Result{Payload: task.Payload, Ctx: ctx}
+	return mq.Result{Payload: task.Payload, Ctx: ctx}
 }

-func handler3(ctx context.Context, task *v2.Task) v2.Result {
+func handler3(ctx context.Context, task *mq.Task) mq.Result {
 	var user map[string]any
 	json.Unmarshal(task.Payload, &user)
 	age := int(user["age"].(float64))
@@ -28,30 +29,30 @@ func handler3(ctx context.Context, task *v2.Task) v2.Result {
 	}
 	user["status"] = status
 	resultPayload, _ := json.Marshal(user)
-	return v2.Result{Payload: resultPayload, Status: status, Ctx: ctx}
+	return mq.Result{Payload: resultPayload, Status: status, Ctx: ctx}
 }

-func handler4(ctx context.Context, task *v2.Task) v2.Result {
+func handler4(ctx context.Context, task *mq.Task) mq.Result {
 	var user map[string]any
 	json.Unmarshal(task.Payload, &user)
 	user["final"] = "D"
 	resultPayload, _ := json.Marshal(user)
-	return v2.Result{Payload: resultPayload, Ctx: ctx}
+	return mq.Result{Payload: resultPayload, Ctx: ctx}
 }

-func handler5(ctx context.Context, task *v2.Task) v2.Result {
+func handler5(ctx context.Context, task *mq.Task) mq.Result {
 	var user map[string]any
 	json.Unmarshal(task.Payload, &user)
 	user["salary"] = "E"
 	resultPayload, _ := json.Marshal(user)
-	return v2.Result{Payload: resultPayload, Ctx: ctx}
+	return mq.Result{Payload: resultPayload, Ctx: ctx}
 }

-func handler6(ctx context.Context, task *v2.Task) v2.Result {
+func handler6(ctx context.Context, task *mq.Task) mq.Result {
 	var user map[string]any
 	json.Unmarshal(task.Payload, &user)
 	resultPayload, _ := json.Marshal(map[string]any{"storage": user})
-	return v2.Result{Payload: resultPayload, Ctx: ctx}
+	return mq.Result{Payload: resultPayload, Ctx: ctx}
 }

 func main() {
--- a/options.go
+++ b/options.go
@@ -3,10 +3,12 @@ package mq
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"time"
 )

 type Result struct {
+	Ctx     context.Context
 	Payload json.RawMessage `json:"payload"`
 	Topic   string          `json:"topic"`
 	TaskID  string          `json:"task_id"`
@@ -14,6 +16,44 @@ type Result struct {
 	Status  string          `json:"status"`
 }

+func (r Result) Unmarshal(data any) error {
+	if r.Payload == nil {
+		return fmt.Errorf("payload is nil")
+	}
+	return json.Unmarshal(r.Payload, data)
+}
+
+func (r Result) String() string {
+	return string(r.Payload)
+}
+
+func HandleError(ctx context.Context, err error, status ...string) Result {
+	st := "Failed"
+	if len(status) > 0 {
+		st = status[0]
+	}
+	if err == nil {
+		return Result{}
+	}
+	return Result{
+		Status: st,
+		Error:  err,
+		Ctx:    ctx,
+	}
+}
+
+func (r Result) WithData(status string, data []byte) Result {
+	if r.Error != nil {
+		return r
+	}
+	return Result{
+		Status:  status,
+		Payload: data,
+		Error:   nil,
+		Ctx:     r.Ctx,
+	}
+}
+
 type TLSConfig struct {
 	UseTLS   bool
 	CertPath string
--- a/v2/dag.go
+++ b/v2/dag.go
@@ -3,83 +3,28 @@ package v2
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"sync"

 	"github.com/oarkflow/xid"
+
+	"github.com/oarkflow/mq"
 )

-type Handler func(ctx context.Context, task *Task) Result
-
-type Result struct {
-	Ctx     context.Context
-	TaskID  string          `json:"task_id"`
-	Payload json.RawMessage `json:"payload"`
-	Status  string          `json:"status"`
-	Error   error           `json:"error"`
-	nodeKey string
-}
-
-func (r Result) Unmarshal(data any) error {
-	if r.Payload == nil {
-		return fmt.Errorf("payload is nil")
-	}
-	return json.Unmarshal(r.Payload, data)
-}
-
-func (r Result) String() string {
-	return string(r.Payload)
-}
-
-func HandleError(ctx context.Context, err error, status ...string) Result {
-	st := "Failed"
-	if len(status) > 0 {
-		st = status[0]
-	}
-	if err == nil {
-		return Result{}
-	}
-	return Result{
-		Status: st,
-		Error:  err,
-		Ctx:    ctx,
-	}
-}
-
-func (r Result) WithData(status string, data []byte) Result {
-	if r.Error != nil {
-		return r
-	}
-	return Result{
-		Status:  status,
-		Payload: data,
-		Error:   nil,
-		Ctx:     r.Ctx,
-	}
-}
-
-type Task struct {
-	ID      string            `json:"id"`
-	NodeKey string            `json:"node_key"`
-	Payload json.RawMessage   `json:"payload"`
-	Results map[string]Result `json:"results"`
-}
-
-func NewTask(id string, payload json.RawMessage, nodeKey string, results ...map[string]Result) *Task {
+func NewTask(id string, payload json.RawMessage, nodeKey string, results ...map[string]mq.Result) *mq.Task {
 	if id == "" {
 		id = xid.New().String()
 	}
-	result := make(map[string]Result)
+	result := make(map[string]mq.Result)
 	if len(results) > 0 && results[0] != nil {
 		result = results[0]
 	}
-	return &Task{ID: id, Payload: payload, NodeKey: nodeKey, Results: result}
+	return &mq.Task{ID: id, Payload: payload, Topic: nodeKey, Results: result}
 }

 type Node struct {
 	Key     string
 	Edges   []Edge
-	handler Handler
+	handler mq.Handler
 }

 type EdgeType int
@@ -112,7 +57,7 @@ func NewDAG() *DAG {
 	}
 }

-func (tm *DAG) AddNode(key string, handler Handler) {
+func (tm *DAG) AddNode(key string, handler mq.Handler) {
 	tm.mu.Lock()
 	defer tm.mu.Unlock()
 	tm.Nodes[key] = &Node{
@@ -145,11 +90,11 @@ func (tm *DAG) AddEdge(from, to string, edgeTypes ...EdgeType) {
 	fromNode.Edges = append(fromNode.Edges, edge)
 }

-func (tm *DAG) ProcessTask(ctx context.Context, node string, payload []byte) Result {
+func (tm *DAG) ProcessTask(ctx context.Context, node string, payload []byte) mq.Result {
 	tm.mu.Lock()
 	defer tm.mu.Unlock()
 	taskID := xid.New().String()
-	task := NewTask(taskID, payload, node, make(map[string]Result))
+	task := NewTask(taskID, payload, node, make(map[string]mq.Result))
 	manager := NewTaskManager(tm)
 	tm.taskContext[taskID] = manager
 	return manager.processTask(ctx, node, task)
--- a/v2/task_manager.go
+++ b/v2/task_manager.go
@@ -5,30 +5,32 @@ import (
 	"encoding/json"
 	"fmt"
 	"sync"
+
+	"github.com/oarkflow/mq"
 )

 type TaskManager struct {
 	dag         *DAG
 	wg          sync.WaitGroup
 	mutex       sync.Mutex
-	results     []Result
-	nodeResults map[string]Result
+	results     []mq.Result
+	nodeResults map[string]mq.Result
 	done        chan struct{}
 }

 func NewTaskManager(d *DAG) *TaskManager {
 	return &TaskManager{
 		dag:         d,
-		nodeResults: make(map[string]Result),
-		results:     make([]Result, 0),
+		nodeResults: make(map[string]mq.Result),
+		results:     make([]mq.Result, 0),
 		done:        make(chan struct{}),
 	}
 }

-func (tm *TaskManager) processTask(ctx context.Context, nodeID string, task *Task) Result {
+func (tm *TaskManager) processTask(ctx context.Context, nodeID string, task *mq.Task) mq.Result {
 	node, ok := tm.dag.Nodes[nodeID]
 	if !ok {
-		return Result{Error: fmt.Errorf("nodeID %s not found", nodeID)}
+		return mq.Result{Error: fmt.Errorf("nodeID %s not found", nodeID)}
 	}
 	tm.wg.Add(1)
 	go tm.processNode(ctx, node, task, nil)
@@ -38,7 +40,7 @@ func (tm *TaskManager) processTask(ctx context.Context, nodeID string, task *Tas
 	}()
 	select {
 	case <-ctx.Done():
-		return Result{Error: ctx.Err()}
+		return mq.Result{Error: ctx.Err()}
 	case <-tm.done:
 		tm.mutex.Lock()
 		defer tm.mutex.Unlock()
@@ -49,10 +51,10 @@ func (tm *TaskManager) processTask(ctx context.Context, nodeID string, task *Tas
 	}
 }

-func (tm *TaskManager) callback(ctx context.Context, results any) Result {
-	var rs Result
+func (tm *TaskManager) callback(ctx context.Context, results any) mq.Result {
+	var rs mq.Result
 	switch res := results.(type) {
-	case []Result:
+	case []mq.Result:
 		aggregatedOutput := make([]json.RawMessage, 0)
 		for i, result := range res {
 			if i == 0 {
@@ -62,43 +64,43 @@ func (tm *TaskManager) callback(ctx context.Context, results any) Result {
 			var item json.RawMessage
 			err := json.Unmarshal(result.Payload, &item)
 			if err != nil {
-				return HandleError(ctx, err)
+				return mq.HandleError(ctx, err)
 			}
 			aggregatedOutput = append(aggregatedOutput, item)
 		}
 		finalOutput, err := json.Marshal(aggregatedOutput)
-		return HandleError(ctx, err).WithData(rs.Status, finalOutput)
-	case Result:
+		return mq.HandleError(ctx, err).WithData(rs.Status, finalOutput)
+	case mq.Result:
 		rs.TaskID = res.TaskID
 		var item json.RawMessage
 		err := json.Unmarshal(res.Payload, &item)
 		if err != nil {
-			return HandleError(ctx, err)
+			return mq.HandleError(ctx, err)
 		}
 		finalOutput, err := json.Marshal(item)
-		return HandleError(ctx, err).WithData(res.Status, finalOutput)
+		return mq.HandleError(ctx, err).WithData(res.Status, finalOutput)
 	}
 	return rs
 }

-func (tm *TaskManager) appendFinalResult(result Result) {
+func (tm *TaskManager) appendFinalResult(result mq.Result) {
 	tm.mutex.Lock()
 	tm.results = append(tm.results, result)
-	tm.nodeResults[result.nodeKey] = result
+	tm.nodeResults[result.Topic] = result
 	tm.mutex.Unlock()
 }

-func (tm *TaskManager) processNode(ctx context.Context, node *Node, task *Task, parentNode *Node) {
+func (tm *TaskManager) processNode(ctx context.Context, node *Node, task *mq.Task, parentNode *Node) {
 	defer tm.wg.Done()
-	var result Result
+	var result mq.Result
 	select {
 	case <-ctx.Done():
-		result = Result{TaskID: task.ID, nodeKey: node.Key, Error: ctx.Err()}
+		result = mq.Result{TaskID: task.ID, Topic: node.Key, Error: ctx.Err()}
 		tm.appendFinalResult(result)
 		return
 	default:
 		result = node.handler(ctx, task)
-		result.nodeKey = node.Key
+		result.Topic = node.Key
 		if result.Error != nil {
 			tm.appendFinalResult(result)
 			return
@@ -113,7 +115,7 @@ func (tm *TaskManager) processNode(ctx context.Context, node *Node, task *Task,
 	edges := make([]Edge, len(node.Edges))
 	copy(edges, node.Edges)
 	if result.Status != "" {
-		if conditions, ok := tm.dag.conditions[result.nodeKey]; ok {
+		if conditions, ok := tm.dag.conditions[result.Topic]; ok {
 			if targetNodeKey, ok := conditions[result.Status]; ok {
 				if targetNode, ok := tm.dag.Nodes[targetNodeKey]; ok {
 					edges = append(edges, Edge{From: node, To: targetNode})
@@ -133,7 +135,7 @@ func (tm *TaskManager) processNode(ctx context.Context, node *Node, task *Task,
 			var items []json.RawMessage
 			err := json.Unmarshal(result.Payload, &items)
 			if err != nil {
-				tm.appendFinalResult(Result{TaskID: task.ID, nodeKey: node.Key, Error: err})
+				tm.appendFinalResult(mq.Result{TaskID: task.ID, Topic: node.Key, Error: err})
 				return
 			}
 			for _, item := range items {