feat: add example

v2/dag.go

@@ -3,6 +3,7 @@ package v2
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"log"
 	"net/http"
 	"sync"
@@ -20,12 +21,6 @@ func NewTask(id string, payload json.RawMessage, nodeKey string) *mq.Task {
 	return &mq.Task{ID: id, Payload: payload, Topic: nodeKey}
 }
 
-type Node struct {
-	Key      string
-	Edges    []Edge
-	consumer *mq.Consumer
-}
-
 type EdgeType int
 
 func (c EdgeType) IsValid() bool { return c >= SimpleEdge && c <= LoopEdge }
@@ -35,6 +30,12 @@ const (
 	LoopEdge
 )
 
+type Node struct {
+	Key      string
+	Edges    []Edge
+	consumer *mq.Consumer
+}
+
 type Edge struct {
 	From *Node
 	To   *Node
@@ -42,6 +43,7 @@ type Edge struct {
 }
 
 type DAG struct {
+	FirstNode   string
 	Nodes       map[string]*Node
 	server      *mq.Broker
 	taskContext map[string]*TaskManager
@@ -68,21 +70,21 @@ func (tm *DAG) onTaskCallback(ctx context.Context, result mq.Result) 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)
+	if !tm.server.SyncMode() {
+		go func() {
+			err := tm.server.Start(ctx)
+			if err != nil {
+				panic(err)
+			}
+		}()
+		for _, con := range tm.Nodes {
+			go func(con *Node) {
+				time.Sleep(1 * time.Second)
+				con.consumer.Consume(ctx)
+			}(con)
 		}
-	}()
-	for _, con := range tm.Nodes {
-		go func(con *Node) {
-			time.Sleep(1 * time.Second)
-			con.consumer.Consume(ctx)
-		}(con)
 	}
+
 	log.Printf("HTTP server started on %s", addr)
 	config := tm.server.TLSConfig()
 	if config.UseTLS {
@@ -91,7 +93,7 @@ func (tm *DAG) Start(ctx context.Context, addr string) error {
 	return http.ListenAndServe(addr, nil)
 }
 
-func (tm *DAG) AddNode(key string, handler mq.Handler) {
+func (tm *DAG) AddNode(key string, handler mq.Handler, firstNode ...bool) {
 	tm.mu.Lock()
 	defer tm.mu.Unlock()
 	con := mq.NewConsumer(key, key, handler)
@@ -99,6 +101,9 @@ func (tm *DAG) AddNode(key string, handler mq.Handler) {
 		Key:      key,
 		consumer: con,
 	}
+	if len(firstNode) > 0 && firstNode[0] {
+		tm.FirstNode = key
+	}
 }
 
 func (tm *DAG) AddCondition(fromNode string, conditions map[string]string) {
@@ -125,11 +130,51 @@ 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) mq.Result {
+func (tm *DAG) ProcessTask(ctx context.Context, payload []byte) mq.Result {
+	val := ctx.Value("initial_node")
+	initialNode, ok := val.(string)
+	if !ok {
+		if tm.FirstNode == "" {
+			firstNode := tm.FindInitialNode()
+			if firstNode != nil {
+				tm.FirstNode = firstNode.Key
+			}
+		}
+		if tm.FirstNode == "" {
+			return mq.Result{Error: fmt.Errorf("initial node not found")}
+		}
+		initialNode = tm.FirstNode
+	}
 	tm.mu.Lock()
 	defer tm.mu.Unlock()
 	taskID := xid.New().String()
 	manager := NewTaskManager(tm, taskID)
 	tm.taskContext[taskID] = manager
-	return manager.processTask(ctx, node, payload)
+	return manager.processTask(ctx, initialNode, payload)
+}
+
+func (tm *DAG) FindInitialNode() *Node {
+	incomingEdges := make(map[string]bool)
+	connectedNodes := make(map[string]bool)
+	for _, node := range tm.Nodes {
+		for _, edge := range node.Edges {
+			if edge.Type.IsValid() {
+				connectedNodes[node.Key] = true
+				connectedNodes[edge.To.Key] = true
+				incomingEdges[edge.To.Key] = true
+			}
+		}
+		if cond, ok := tm.conditions[node.Key]; ok {
+			for _, target := range cond {
+				connectedNodes[target] = true
+				incomingEdges[target] = true
+			}
+		}
+	}
+	for nodeID, node := range tm.Nodes {
+		if !incomingEdges[nodeID] && connectedNodes[nodeID] {
+			return node
+		}
+	}
+	return nil
 }