feat: add example

ctx.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/oarkflow/xid"
@@ -38,7 +39,11 @@ func IsClosed(conn net.Conn) bool {
 	return false
 }
 
+var m = sync.RWMutex{}
+
 func SetHeaders(ctx context.Context, headers map[string]string) context.Context {
+	m.Lock()
+	defer m.Unlock()
 	hd, ok := GetHeaders(ctx)
 	if !ok {
 		hd = make(map[string]string)

dag/dag.go

@@ -176,7 +176,7 @@ func (d *DAG) Send(ctx context.Context, payload []byte) mq.Result {
 
 func (d *DAG) processNode(ctx context.Context, task mq.Result) mq.Result {
 	if con, ok := d.nodes[task.Topic]; ok {
-		return con.ProcessTask(ctx, mq.Task{
+		return con.ProcessTask(ctx, &mq.Task{
 			ID:      task.TaskID,
 			Payload: task.Payload,
 		})

examples/dag.go

@@ -1,5 +1,6 @@
 package main
 
+/*
 import (
 	"context"
 	"encoding/json"
@@ -46,13 +47,13 @@ func main() {
 	}()
 
 	time.Sleep(10 * time.Second)
-	/*d.Prepare()
+	d.Prepare()
 	http.HandleFunc("POST /publish", requestHandler("publish"))
 	http.HandleFunc("POST /request", requestHandler("request"))
 	err := d.Start(context.TODO(), ":8083")
 	if err != nil {
 		panic(err)
-	}*/
+	}
 }
 
 func requestHandler(requestType string) func(w http.ResponseWriter, r *http.Request) {
@@ -89,3 +90,4 @@ func requestHandler(requestType string) func(w http.ResponseWriter, r *http.Requ
 		json.NewEncoder(w).Encode(result)
 	}
 }
+*/

examples/dag_v2.go

@@ -4,19 +4,21 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
+	"net/http"
 
 	"github.com/oarkflow/mq"
 	"github.com/oarkflow/mq/v2"
 )
 
 func handler1(ctx context.Context, task *mq.Task) mq.Result {
-	return mq.Result{Payload: task.Payload, Ctx: ctx}
+	return mq.Result{Payload: task.Payload}
 }
 
 func handler2(ctx context.Context, task *mq.Task) mq.Result {
 	var user map[string]any
 	json.Unmarshal(task.Payload, &user)
-	return mq.Result{Payload: task.Payload, Ctx: ctx}
+	return mq.Result{Payload: task.Payload}
 }
 
 func handler3(ctx context.Context, task *mq.Task) mq.Result {
@@ -29,7 +31,7 @@ func handler3(ctx context.Context, task *mq.Task) mq.Result {
 	}
 	user["status"] = status
 	resultPayload, _ := json.Marshal(user)
-	return mq.Result{Payload: resultPayload, Status: status, Ctx: ctx}
+	return mq.Result{Payload: resultPayload, Status: status}
 }
 
 func handler4(ctx context.Context, task *mq.Task) mq.Result {
@@ -37,7 +39,7 @@ func handler4(ctx context.Context, task *mq.Task) mq.Result {
 	json.Unmarshal(task.Payload, &user)
 	user["final"] = "D"
 	resultPayload, _ := json.Marshal(user)
-	return mq.Result{Payload: resultPayload, Ctx: ctx}
+	return mq.Result{Payload: resultPayload}
 }
 
 func handler5(ctx context.Context, task *mq.Task) mq.Result {
@@ -45,34 +47,73 @@ func handler5(ctx context.Context, task *mq.Task) mq.Result {
 	json.Unmarshal(task.Payload, &user)
 	user["salary"] = "E"
 	resultPayload, _ := json.Marshal(user)
-	return mq.Result{Payload: resultPayload, Ctx: ctx}
+	return mq.Result{Payload: resultPayload}
 }
 
 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 mq.Result{Payload: resultPayload, Ctx: ctx}
+	return mq.Result{Payload: resultPayload}
 }
 
+var (
+	d = v2.NewDAG(mq.WithSyncMode(true))
+)
+
 func main() {
-	dag := v2.NewDAG()
+	d.AddNode("A", handler1)
-	dag.AddNode("A", handler1)
+	d.AddNode("B", handler2)
-	dag.AddNode("B", handler2)
+	d.AddNode("C", handler3)
-	dag.AddNode("C", handler3)
+	d.AddNode("D", handler4)
-	dag.AddNode("D", handler4)
+	d.AddNode("E", handler5)
-	dag.AddNode("E", handler5)
+	d.AddNode("F", handler6)
-	dag.AddNode("F", handler6)
+	d.AddEdge("A", "B", v2.LoopEdge)
-	dag.AddEdge("A", "B", v2.LoopEdge)
+	d.AddCondition("C", map[string]string{"PASS": "D", "FAIL": "E"})
-	dag.AddCondition("C", map[string]string{"PASS": "D", "FAIL": "E"})
+	d.AddEdge("B", "C")
-	dag.AddEdge("B", "C")
+	d.AddEdge("D", "F")
-	dag.AddEdge("D", "F")
+	d.AddEdge("E", "F")
-	dag.AddEdge("E", "F")
 
 	initialPayload, _ := json.Marshal([]map[string]any{
 		{"user_id": 1, "age": 12},
 		{"user_id": 2, "age": 34},
 	})
-	rs := dag.ProcessTask(context.Background(), "A", initialPayload)
+	rs := d.ProcessTask(context.Background(), "A", initialPayload)
 	fmt.Println(string(rs.Payload))
+	http.HandleFunc("POST /publish", requestHandler("publish"))
+	http.HandleFunc("POST /request", requestHandler("request"))
+	err := d.Start(context.TODO(), ":8083")
+	if err != nil {
+		panic(err)
+	}
+}
+
+func requestHandler(requestType string) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			http.Error(w, "Invalid request method", http.StatusMethodNotAllowed)
+			return
+		}
+		var payload []byte
+		if r.Body != nil {
+			defer r.Body.Close()
+			var err error
+			payload, err = io.ReadAll(r.Body)
+			if err != nil {
+				http.Error(w, "Failed to read request body", http.StatusBadRequest)
+				return
+			}
+		} else {
+			http.Error(w, "Empty request body", http.StatusBadRequest)
+			return
+		}
+		rs := d.ProcessTask(context.Background(), "A", payload)
+		w.Header().Set("Content-Type", "application/json")
+		result := map[string]any{
+			"message_id": rs.TaskID,
+			"payload":    string(rs.Payload),
+			"error":      rs.Error,
+		}
+		json.NewEncoder(w).Encode(result)
+	}
 }

options.go

@@ -8,7 +8,6 @@ import (
 )
 
 type Result struct {
-	Ctx     context.Context
 	Payload json.RawMessage `json:"payload"`
 	Topic   string          `json:"topic"`
 	TaskID  string          `json:"task_id"`
@@ -38,7 +37,6 @@ func HandleError(ctx context.Context, err error, status ...string) Result {
 	return Result{
 		Status: st,
 		Error:  err,
-		Ctx:    ctx,
 	}
 }
 
@@ -50,7 +48,6 @@ func (r Result) WithData(status string, data []byte) Result {
 		Status:  status,
 		Payload: data,
 		Error:   nil,
-		Ctx:     r.Ctx,
 	}
 }
 

v2/dag.go

@@ -3,6 +3,8 @@ package v2
 import (
 	"context"
 	"encoding/json"
+	"log"
+	"net/http"
 	"sync"
 
 	"github.com/oarkflow/xid"
@@ -44,17 +46,51 @@ type Edge struct {
 
 type DAG struct {
 	Nodes       map[string]*Node
+	server      *mq.Broker
 	taskContext map[string]*TaskManager
 	conditions  map[string]map[string]string
 	mu          sync.RWMutex
 }
 
-func NewDAG() *DAG {
+func NewDAG(opts ...mq.Option) *DAG {
-	return &DAG{
+	d := &DAG{
 		Nodes:       make(map[string]*Node),
 		taskContext: make(map[string]*TaskManager),
 		conditions:  make(map[string]map[string]string),
 	}
+	opts = append(opts, mq.WithCallback(d.onTaskCallback))
+	d.server = mq.NewBroker(opts...)
+	return d
+}
+
+func (tm *DAG) onTaskCallback(ctx context.Context, result mq.Result) mq.Result {
+	if taskContext, ok := tm.taskContext[result.TaskID]; ok {
+		return taskContext.handleCallback(ctx, result)
+	}
+	return mq.Result{}
+}
+
+func (tm *DAG) Start(ctx context.Context, addr string) error {
+	if tm.server.SyncMode() {
+		return nil
+	}
+	go func() {
+		err := tm.server.Start(ctx)
+		if err != nil {
+			panic(err)
+		}
+	}()
+	for _, con := range tm.Nodes {
+		go func(con *Node) {
+			con.consumer.Consume(ctx)
+		}(con)
+	}
+	log.Printf("HTTP server started on %s", addr)
+	config := tm.server.TLSConfig()
+	if config.UseTLS {
+		return http.ListenAndServeTLS(addr, config.CertPath, config.KeyPath, nil)
+	}
+	return http.ListenAndServe(addr, nil)
 }
 
 func (tm *DAG) AddNode(key string, handler mq.Handler) {

v2/task_manager.go

@@ -46,13 +46,18 @@ func (tm *TaskManager) processTask(ctx context.Context, nodeID string, task *mq.
 		tm.mutex.Lock()
 		defer tm.mutex.Unlock()
 		if len(tm.results) == 1 {
-			return tm.callback(ctx, tm.results[0])
+			return tm.handleResult(ctx, tm.results[0])
 		}
-		return tm.callback(ctx, tm.results)
+		return tm.handleResult(ctx, tm.results)
 	}
 }
 
-func (tm *TaskManager) callback(ctx context.Context, results any) mq.Result {
+func (tm *TaskManager) handleCallback(ctx context.Context, result mq.Result) mq.Result {
+	fmt.Println(string(result.Payload), result.Topic, result.TaskID)
+	return mq.Result{}
+}
+
+func (tm *TaskManager) handleResult(ctx context.Context, results any) mq.Result {
 	var rs mq.Result
 	switch res := results.(type) {
 	case []mq.Result:
@@ -101,17 +106,21 @@ func (tm *TaskManager) processNode(ctx context.Context, node *Node, task *mq.Tas
 		return
 	default:
 		ctx = mq.SetHeaders(ctx, map[string]string{consts.QueueKey: node.Key})
-		result = node.consumer.ProcessTask(ctx, task)
+		if tm.dag.server.SyncMode() {
-		result.Topic = node.Key
+			result = node.consumer.ProcessTask(ctx, task)
-		if result.Error != nil {
+			result.Topic = node.Key
-			tm.appendFinalResult(result)
+			if result.Error != nil {
-			return
+				tm.appendFinalResult(result)
+				return
+			}
+		} else {
+			err := tm.dag.server.Publish(ctx, *task, node.Key)
+			if err != nil {
+				tm.appendFinalResult(mq.Result{Error: err})
+				return
+			}
 		}
 	}
-	if result.Ctx == nil {
-		result.Ctx = ctx
-	}
-	ctx = result.Ctx
 	tm.mutex.Lock()
 	task.Results[node.Key] = result
 	tm.mutex.Unlock()