package/mq
1
0
Fork 0
mirror of https://github.com/oarkflow/mq.git synced 2025-10-06 00:16:49 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: add example

Browse Source
This commit is contained in:
sujit
2024-10-08 16:18:53 +05:45
parent a226982aa2
commit a4dff9c77b
8 changed files with 84 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
codec/codec.go
View File

@@ -31,7 +31,7 @@ func NewMessage(cmd consts.CMD, payload json.RawMessage, queue string, headers m
func (m *Message) Serialize(aesKey, hmacKey []byte, encrypt bool) ([]byte, string, error) { func (m *Message) Serialize(aesKey, hmacKey []byte, encrypt bool) ([]byte, string, error) {
var buf bytes.Buffer var buf bytes.Buffer
// Serialize Headers, Queue, Command, Payload, and Metadata // Serialize Headers, Topic, Command, Payload, and Metadata
if err := writeLengthPrefixedJSON(&buf, m.Headers); err != nil { if err := writeLengthPrefixedJSON(&buf, m.Headers); err != nil {
return nil, "", fmt.Errorf("error serializing headers: %v", err) return nil, "", fmt.Errorf("error serializing headers: %v", err)
} }
@@ -62,7 +62,7 @@ func Deserialize(data, aesKey, hmacKey []byte, receivedHMAC string, decrypt bool
buf := bytes.NewReader(data) buf := bytes.NewReader(data)
// Deserialize Headers, Queue, Command, Payload, and Metadata // Deserialize Headers, Topic, Command, Payload, and Metadata
headers := make(map[string]string) headers := make(map[string]string)
if err := readLengthPrefixedJSON(buf, &headers); err != nil { if err := readLengthPrefixedJSON(buf, &headers); err != nil {
return nil, fmt.Errorf("error deserializing headers: %v", err) return nil, fmt.Errorf("error deserializing headers: %v", err)

2
consts/constants.go
View File

@@ -54,7 +54,7 @@ var (
PublisherKey = "Publisher-Key" PublisherKey = "Publisher-Key"
ContentType = "Content-Type" ContentType = "Content-Type"
AwaitResponseKey = "Await-Response" AwaitResponseKey = "Await-Response"
QueueKey = "Queue" QueueKey = "Topic"
TypeJson = "application/json" TypeJson = "application/json"
HeaderKey = "headers" HeaderKey = "headers"
TriggerNode = "triggerNode" TriggerNode = "triggerNode"

8
consumer.go
View File

@@ -89,9 +89,9 @@ func (c *Consumer) OnMessage(ctx context.Context, msg *codec.Message, conn net.C
return return
} }
ctx = SetHeaders(ctx, map[string]string{consts.QueueKey: msg.Queue}) ctx = SetHeaders(ctx, map[string]string{consts.QueueKey: msg.Queue})
result := c.ProcessTask(ctx, task) result := c.ProcessTask(ctx, &task)
result.MessageID = task.ID result.TaskID = task.ID
result.Queue = msg.Queue result.Topic = msg.Queue
if result.Status == "" { if result.Status == "" {
if result.Error != nil { if result.Error != nil {
result.Status = "FAILED" result.Status = "FAILED"
@@ -107,7 +107,7 @@ func (c *Consumer) OnMessage(ctx context.Context, msg *codec.Message, conn net.C
} }
// ProcessTask handles a received task message and invokes the appropriate handler. // ProcessTask handles a received task message and invokes the appropriate handler.
func (c *Consumer) ProcessTask(ctx context.Context, msg Task) Result { func (c *Consumer) ProcessTask(ctx context.Context, msg *Task) Result {
queue, _ := GetQueue(ctx) queue, _ := GetQueue(ctx)
handler, exists := c.handlers[queue] handler, exists := c.handlers[queue]
if !exists { if !exists {

16
ctx.go
View File

@@ -16,15 +16,17 @@ import (
) )
type Task struct { type Task struct {
ID string `json:"id"` ID string `json:"id"`
Payload json.RawMessage `json:"payload"` Results map[string]Result `json:"results"`
CreatedAt time.Time `json:"created_at"` Topic string `json:"topic"`
ProcessedAt time.Time `json:"processed_at"` Payload json.RawMessage `json:"payload"`
Status string `json:"status"` CreatedAt time.Time `json:"created_at"`
Error error `json:"error"` ProcessedAt time.Time `json:"processed_at"`
Status string `json:"status"`
Error error `json:"error"`
} }
type Handler func(context.Context, Task) Result type Handler func(context.Context, *Task) Result
func IsClosed(conn net.Conn) bool { func IsClosed(conn net.Conn) bool {
_, err := conn.Read(make([]byte, 1)) _, err := conn.Read(make([]byte, 1))

96
dag/dag.go
View File

@@ -124,9 +124,9 @@ func (d *DAG) PublishTask(ctx context.Context, payload json.RawMessage, taskID .
return mq.Result{Error: err} return mq.Result{Error: err}
} }
return mq.Result{ return mq.Result{
Payload: payload, Payload: payload,
Queue: queue, Topic: queue,
MessageID: id, TaskID: id,
} }
} }
@@ -168,37 +168,37 @@ func (d *DAG) Send(ctx context.Context, payload []byte) mq.Result {
return result return result
} }
d.mu.Lock() d.mu.Lock()
d.taskChMap[result.MessageID] = resultCh d.taskChMap[result.TaskID] = resultCh
d.mu.Unlock() d.mu.Unlock()
finalResult := <-resultCh finalResult := <-resultCh
return finalResult return finalResult
} }
func (d *DAG) processNode(ctx context.Context, task mq.Result) mq.Result { func (d *DAG) processNode(ctx context.Context, task mq.Result) mq.Result {
if con, ok := d.nodes[task.Queue]; ok { if con, ok := d.nodes[task.Topic]; ok {
return con.ProcessTask(ctx, mq.Task{ return con.ProcessTask(ctx, mq.Task{
ID: task.MessageID, ID: task.TaskID,
Payload: task.Payload, Payload: task.Payload,
}) })
} }
return mq.Result{Error: fmt.Errorf("no consumer to process %s", task.Queue)} return mq.Result{Error: fmt.Errorf("no consumer to process %s", task.Topic)}
} }
func (d *DAG) sendSync(ctx context.Context, task mq.Result) mq.Result { func (d *DAG) sendSync(ctx context.Context, task mq.Result) mq.Result {
if task.MessageID == "" { if task.TaskID == "" {
task.MessageID = mq.NewID() task.TaskID = mq.NewID()
} }
if task.Queue == "" { if task.Topic == "" {
task.Queue = d.FirstNode task.Topic = d.FirstNode
} }
ctx = mq.SetHeaders(ctx, map[string]string{ ctx = mq.SetHeaders(ctx, map[string]string{
consts.QueueKey: task.Queue, consts.QueueKey: task.Topic,
}) })
result := d.processNode(ctx, task) result := d.processNode(ctx, task)
if result.Error != nil { if result.Error != nil {
return result return result
} }
for _, target := range d.loopEdges[task.Queue] { for _, target := range d.loopEdges[task.Topic] {
var items, results []json.RawMessage var items, results []json.RawMessage
if err := json.Unmarshal(result.Payload, &items); err != nil { if err := json.Unmarshal(result.Payload, &items); err != nil {
return mq.Result{Error: err} return mq.Result{Error: err}
@@ -208,9 +208,9 @@ func (d *DAG) sendSync(ctx context.Context, task mq.Result) mq.Result {
consts.QueueKey: target, consts.QueueKey: target,
}) })
result = d.sendSync(ctx, mq.Result{ result = d.sendSync(ctx, mq.Result{
Payload: item, Payload: item,
Queue: target, Topic: target,
MessageID: result.MessageID, TaskID: result.TaskID,
}) })
if result.Error != nil { if result.Error != nil {
return result return result
@@ -223,29 +223,29 @@ func (d *DAG) sendSync(ctx context.Context, task mq.Result) mq.Result {
} }
result.Payload = bt result.Payload = bt
} }
if conditions, ok := d.conditions[task.Queue]; ok { if conditions, ok := d.conditions[task.Topic]; ok {
if target, exists := conditions[result.Status]; exists { if target, exists := conditions[result.Status]; exists {
ctx = mq.SetHeaders(ctx, map[string]string{ ctx = mq.SetHeaders(ctx, map[string]string{
consts.QueueKey: target, consts.QueueKey: target,
}) })
result = d.sendSync(ctx, mq.Result{ result = d.sendSync(ctx, mq.Result{
Payload: result.Payload, Payload: result.Payload,
Queue: target, Topic: target,
MessageID: result.MessageID, TaskID: result.TaskID,
}) })
if result.Error != nil { if result.Error != nil {
return result return result
} }
} }
} }
if target, ok := d.edges[task.Queue]; ok { if target, ok := d.edges[task.Topic]; ok {
ctx = mq.SetHeaders(ctx, map[string]string{ ctx = mq.SetHeaders(ctx, map[string]string{
consts.QueueKey: target, consts.QueueKey: target,
}) })
result = d.sendSync(ctx, mq.Result{ result = d.sendSync(ctx, mq.Result{
Payload: result.Payload, Payload: result.Payload,
Queue: target, Topic: target,
MessageID: result.MessageID, TaskID: result.TaskID,
}) })
if result.Error != nil { if result.Error != nil {
return result return result
@@ -260,7 +260,7 @@ func (d *DAG) getCompletedResults(task mq.Result, ok bool, triggeredNode string)
completed := false completed := false
multipleResults := false multipleResults := false
if ok && triggeredNode != "" { if ok && triggeredNode != "" {
taskResults, ok := d.taskResults[task.MessageID] taskResults, ok := d.taskResults[task.TaskID]
if ok { if ok {
nodeResult, exists := taskResults[triggeredNode] nodeResult, exists := taskResults[triggeredNode]
if exists { if exists {
@@ -300,25 +300,25 @@ func (d *DAG) TaskCallback(ctx context.Context, task mq.Result) mq.Result {
} }
triggeredNode, ok := mq.GetTriggerNode(ctx) triggeredNode, ok := mq.GetTriggerNode(ctx)
payload, completed, multipleResults := d.getCompletedResults(task, ok, triggeredNode) payload, completed, multipleResults := d.getCompletedResults(task, ok, triggeredNode)
if loopNodes, exists := d.loopEdges[task.Queue]; exists { if loopNodes, exists := d.loopEdges[task.Topic]; exists {
var items []json.RawMessage var items []json.RawMessage
if err := json.Unmarshal(payload, &items); err != nil { if err := json.Unmarshal(payload, &items); err != nil {
return mq.Result{Error: task.Error} return mq.Result{Error: task.Error}
} }
d.taskResults[task.MessageID] = map[string]*taskContext{ d.taskResults[task.TaskID] = map[string]*taskContext{
task.Queue: { task.Topic: {
totalItems: len(items), totalItems: len(items),
multipleResults: true, multipleResults: true,
}, },
} }
ctx = mq.SetHeaders(ctx, map[string]string{consts.TriggerNode: task.Queue}) ctx = mq.SetHeaders(ctx, map[string]string{consts.TriggerNode: task.Topic})
for _, loopNode := range loopNodes { for _, loopNode := range loopNodes {
for _, item := range items { for _, item := range items {
ctx = mq.SetHeaders(ctx, map[string]string{ ctx = mq.SetHeaders(ctx, map[string]string{
consts.QueueKey: loopNode, consts.QueueKey: loopNode,
}) })
result := d.PublishTask(ctx, item, task.MessageID) result := d.PublishTask(ctx, item, task.TaskID)
if result.Error != nil { if result.Error != nil {
return result return result
} }
@@ -328,51 +328,51 @@ func (d *DAG) TaskCallback(ctx context.Context, task mq.Result) mq.Result {
return task return task
} }
if multipleResults && completed { if multipleResults && completed {
task.Queue = triggeredNode task.Topic = triggeredNode
} }
if conditions, ok := d.conditions[task.Queue]; ok { if conditions, ok := d.conditions[task.Topic]; ok {
if target, exists := conditions[task.Status]; exists { if target, exists := conditions[task.Status]; exists {
d.taskResults[task.MessageID] = map[string]*taskContext{ d.taskResults[task.TaskID] = map[string]*taskContext{
task.Queue: { task.Topic: {
totalItems: len(conditions), totalItems: len(conditions),
}, },
} }
ctx = mq.SetHeaders(ctx, map[string]string{ ctx = mq.SetHeaders(ctx, map[string]string{
consts.QueueKey: target, consts.QueueKey: target,
consts.TriggerNode: task.Queue, consts.TriggerNode: task.Topic,
}) })
result := d.PublishTask(ctx, payload, task.MessageID) result := d.PublishTask(ctx, payload, task.TaskID)
if result.Error != nil { if result.Error != nil {
return result return result
} }
} }
} else { } else {
ctx = mq.SetHeaders(ctx, map[string]string{consts.TriggerNode: task.Queue}) ctx = mq.SetHeaders(ctx, map[string]string{consts.TriggerNode: task.Topic})
edge, exists := d.edges[task.Queue] edge, exists := d.edges[task.Topic]
if exists { if exists {
d.taskResults[task.MessageID] = map[string]*taskContext{ d.taskResults[task.TaskID] = map[string]*taskContext{
task.Queue: { task.Topic: {
totalItems: 1, totalItems: 1,
}, },
} }
ctx = mq.SetHeaders(ctx, map[string]string{ ctx = mq.SetHeaders(ctx, map[string]string{
consts.QueueKey: edge, consts.QueueKey: edge,
}) })
result := d.PublishTask(ctx, payload, task.MessageID) result := d.PublishTask(ctx, payload, task.TaskID)
if result.Error != nil { if result.Error != nil {
return result return result
} }
} else if completed { } else if completed {
d.mu.Lock() d.mu.Lock()
if resultCh, ok := d.taskChMap[task.MessageID]; ok { if resultCh, ok := d.taskChMap[task.TaskID]; ok {
resultCh <- mq.Result{ resultCh <- mq.Result{
Payload: payload, Payload: payload,
Queue: task.Queue, Topic: task.Topic,
MessageID: task.MessageID, TaskID: task.TaskID,
Status: "done", Status: "done",
} }
delete(d.taskChMap, task.MessageID) delete(d.taskChMap, task.TaskID)
delete(d.taskResults, task.MessageID) delete(d.taskResults, task.TaskID)
} }
d.mu.Unlock() d.mu.Unlock()
} }

2
examples/dag.go
View File

@@ -82,7 +82,7 @@ func requestHandler(requestType string) func(w http.ResponseWriter, r *http.Requ
} }
w.Header().Set("Content-Type", "application/json") w.Header().Set("Content-Type", "application/json")
result := map[string]any{ result := map[string]any{
"message_id": rs.MessageID, "message_id": rs.TaskID,
"payload": string(rs.Payload), "payload": string(rs.Payload),
"error": rs.Error, "error": rs.Error,
} }

27
examples/tasks/tasks.go
View File

@@ -4,18 +4,19 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"github.com/oarkflow/mq" "github.com/oarkflow/mq"
) )
func Node1(ctx context.Context, task mq.Task) mq.Result { func Node1(ctx context.Context, task *mq.Task) mq.Result {
return mq.Result{Payload: task.Payload, MessageID: task.ID} return mq.Result{Payload: task.Payload, TaskID: task.ID}
} }
func Node2(ctx context.Context, task mq.Task) mq.Result { func Node2(ctx context.Context, task *mq.Task) mq.Result {
return mq.Result{Payload: task.Payload, MessageID: task.ID} return mq.Result{Payload: task.Payload, TaskID: task.ID}
} }
func Node3(ctx context.Context, task mq.Task) mq.Result { func Node3(ctx context.Context, task *mq.Task) mq.Result {
var data map[string]any var data map[string]any
err := json.Unmarshal(task.Payload, &data) err := json.Unmarshal(task.Payload, &data)
if err != nil { if err != nil {
@@ -23,10 +24,10 @@ func Node3(ctx context.Context, task mq.Task) mq.Result {
} }
data["salary"] = fmt.Sprintf("12000%v", data["user_id"]) data["salary"] = fmt.Sprintf("12000%v", data["user_id"])
bt, _ := json.Marshal(data) bt, _ := json.Marshal(data)
return mq.Result{Payload: bt, MessageID: task.ID} return mq.Result{Payload: bt, TaskID: task.ID}
} }
func Node4(ctx context.Context, task mq.Task) mq.Result { func Node4(ctx context.Context, task *mq.Task) mq.Result {
var data []map[string]any var data []map[string]any
err := json.Unmarshal(task.Payload, &data) err := json.Unmarshal(task.Payload, &data)
if err != nil { if err != nil {
@@ -34,10 +35,10 @@ func Node4(ctx context.Context, task mq.Task) mq.Result {
} }
payload := map[string]any{"storage": data} payload := map[string]any{"storage": data}
bt, _ := json.Marshal(payload) bt, _ := json.Marshal(payload)
return mq.Result{Payload: bt, MessageID: task.ID} return mq.Result{Payload: bt, TaskID: task.ID}
} }
func CheckCondition(ctx context.Context, task mq.Task) mq.Result { func CheckCondition(ctx context.Context, task *mq.Task) mq.Result {
var data map[string]any var data map[string]any
err := json.Unmarshal(task.Payload, &data) err := json.Unmarshal(task.Payload, &data)
if err != nil { if err != nil {
@@ -49,20 +50,20 @@ func CheckCondition(ctx context.Context, task mq.Task) mq.Result {
} else { } else {
status = "fail" status = "fail"
} }
return mq.Result{Status: status, Payload: task.Payload, MessageID: task.ID} return mq.Result{Status: status, Payload: task.Payload, TaskID: task.ID}
} }
func Pass(ctx context.Context, task mq.Task) mq.Result { func Pass(ctx context.Context, task *mq.Task) mq.Result {
fmt.Println("Pass") fmt.Println("Pass")
return mq.Result{Payload: task.Payload} return mq.Result{Payload: task.Payload}
} }
func Fail(ctx context.Context, task mq.Task) mq.Result { func Fail(ctx context.Context, task *mq.Task) mq.Result {
fmt.Println("Fail") fmt.Println("Fail")
return mq.Result{Payload: []byte(`{"test2": "asdsa"}`)} return mq.Result{Payload: []byte(`{"test2": "asdsa"}`)}
} }
func Callback(ctx context.Context, task mq.Result) mq.Result { func Callback(ctx context.Context, task mq.Result) mq.Result {
fmt.Println("Received task", task.MessageID, "Payload", string(task.Payload), task.Error, task.Queue) fmt.Println("Received task", task.TaskID, "Payload", string(task.Payload), task.Error, task.Topic)
return mq.Result{} return mq.Result{}
} }

10
options.go
View File

@@ -7,11 +7,11 @@ import (
) )
type Result struct { type Result struct {
Payload json.RawMessage `json:"payload"` Payload json.RawMessage `json:"payload"`
Queue string `json:"queue"` Topic string `json:"topic"`
MessageID string `json:"message_id"` TaskID string `json:"task_id"`
Error error `json:"error,omitempty"` Error error `json:"error,omitempty"`
Status string `json:"status"` Status string `json:"status"`
} }
type TLSConfig struct { type TLSConfig struct {