mirror of
https://github.com/oarkflow/mq.git
synced 2025-10-06 16:36:53 +08:00
650 lines
14 KiB
Go
650 lines
14 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"os/signal"
|
|
"sync"
|
|
"sync/atomic"
|
|
"syscall"
|
|
"time"
|
|
)
|
|
|
|
// ---------------------- Public API Interfaces ----------------------
|
|
|
|
type Processor func(ctx context.Context, in any) (any, error)
|
|
|
|
type Node interface {
|
|
ID() string
|
|
Start(ctx context.Context, in <-chan any) <-chan any
|
|
}
|
|
|
|
type Pipeline interface {
|
|
Start(ctx context.Context, inputs <-chan any) (<-chan any, error)
|
|
}
|
|
|
|
// ---------------------- Processor Registry ----------------------
|
|
|
|
var procRegistry = map[string]Processor{}
|
|
|
|
func RegisterProcessor(name string, p Processor) {
|
|
procRegistry[name] = p
|
|
}
|
|
|
|
func GetProcessor(name string) (Processor, bool) {
|
|
p, ok := procRegistry[name]
|
|
return p, ok
|
|
}
|
|
|
|
// ---------------------- Ring Buffer (SPSC lock-free) ----------------------
|
|
|
|
type RingBuffer struct {
|
|
buf []any
|
|
mask uint64
|
|
head uint64
|
|
tail uint64
|
|
}
|
|
|
|
func NewRingBuffer(size uint64) *RingBuffer {
|
|
if size == 0 || (size&(size-1)) != 0 {
|
|
panic("ring size must be power of two")
|
|
}
|
|
return &RingBuffer{buf: make([]any, size), mask: size - 1}
|
|
}
|
|
|
|
func (r *RingBuffer) Push(v any) bool {
|
|
t := atomic.LoadUint64(&r.tail)
|
|
h := atomic.LoadUint64(&r.head)
|
|
if t-h == uint64(len(r.buf)) {
|
|
return false
|
|
}
|
|
r.buf[t&r.mask] = v
|
|
atomic.AddUint64(&r.tail, 1)
|
|
return true
|
|
}
|
|
|
|
func (r *RingBuffer) Pop() (any, bool) {
|
|
h := atomic.LoadUint64(&r.head)
|
|
t := atomic.LoadUint64(&r.tail)
|
|
if t == h {
|
|
return nil, false
|
|
}
|
|
v := r.buf[h&r.mask]
|
|
atomic.AddUint64(&r.head, 1)
|
|
return v, true
|
|
}
|
|
|
|
// ---------------------- Node Implementations ----------------------
|
|
|
|
type ChannelNode struct {
|
|
id string
|
|
processor Processor
|
|
buf int
|
|
workers int
|
|
}
|
|
|
|
func NewChannelNode(id string, proc Processor, buf int, workers int) *ChannelNode {
|
|
if buf <= 0 {
|
|
buf = 64
|
|
}
|
|
if workers <= 0 {
|
|
workers = 1
|
|
}
|
|
return &ChannelNode{id: id, processor: proc, buf: buf, workers: workers}
|
|
}
|
|
|
|
func (c *ChannelNode) ID() string { return c.id }
|
|
|
|
func (c *ChannelNode) Start(ctx context.Context, in <-chan any) <-chan any {
|
|
out := make(chan any, c.buf)
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < c.workers; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case v, ok := <-in:
|
|
if !ok {
|
|
return
|
|
}
|
|
res, err := c.processor(ctx, v)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "processor %s error: %v\n", c.id, err)
|
|
continue
|
|
}
|
|
select {
|
|
case out <- res:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
go func() {
|
|
wg.Wait()
|
|
close(out)
|
|
}()
|
|
return out
|
|
}
|
|
|
|
type PageNode struct {
|
|
id string
|
|
processor Processor
|
|
buf int
|
|
workers int
|
|
}
|
|
|
|
func NewPageNode(id string, proc Processor, buf int, workers int) *PageNode {
|
|
if buf <= 0 {
|
|
buf = 64
|
|
}
|
|
if workers <= 0 {
|
|
workers = 1
|
|
}
|
|
return &PageNode{id: id, processor: proc, buf: buf, workers: workers}
|
|
}
|
|
|
|
func (c *PageNode) ID() string { return c.id }
|
|
|
|
func (c *PageNode) Start(ctx context.Context, in <-chan any) <-chan any {
|
|
out := make(chan any, c.buf)
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < c.workers; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case v, ok := <-in:
|
|
if !ok {
|
|
return
|
|
}
|
|
res, err := c.processor(ctx, v)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "processor %s error: %v\n", c.id, err)
|
|
continue
|
|
}
|
|
select {
|
|
case out <- res:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
go func() {
|
|
wg.Wait()
|
|
close(out)
|
|
}()
|
|
return out
|
|
}
|
|
|
|
type RingNode struct {
|
|
id string
|
|
processor Processor
|
|
size uint64
|
|
}
|
|
|
|
func NewRingNode(id string, proc Processor, size uint64) *RingNode {
|
|
if size == 0 {
|
|
size = 1024
|
|
}
|
|
n := uint64(1)
|
|
for n < size {
|
|
n <<= 1
|
|
}
|
|
return &RingNode{id: id, processor: proc, size: n}
|
|
}
|
|
|
|
func (r *RingNode) ID() string { return r.id }
|
|
|
|
func (r *RingNode) Start(ctx context.Context, in <-chan any) <-chan any {
|
|
out := make(chan any, 64)
|
|
ring := NewRingBuffer(r.size)
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case v, ok := <-in:
|
|
if !ok {
|
|
return
|
|
}
|
|
for !ring.Push(v) {
|
|
time.Sleep(time.Microsecond)
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
default:
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
go func() {
|
|
defer close(out)
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-done:
|
|
// process remaining items in ring
|
|
for {
|
|
v, ok := ring.Pop()
|
|
if !ok {
|
|
return
|
|
}
|
|
res, err := r.processor(ctx, v)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "processor %s error: %v\n", r.id, err)
|
|
continue
|
|
}
|
|
select {
|
|
case out <- res:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
default:
|
|
v, ok := ring.Pop()
|
|
if !ok {
|
|
time.Sleep(time.Microsecond)
|
|
continue
|
|
}
|
|
res, err := r.processor(ctx, v)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "processor %s error: %v\n", r.id, err)
|
|
continue
|
|
}
|
|
select {
|
|
case out <- res:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
return out
|
|
}
|
|
|
|
// ---------------------- DAG Pipeline ----------------------
|
|
|
|
type NodeSpec struct {
|
|
ID string `json:"id"`
|
|
Type string `json:"type"`
|
|
Processor string `json:"processor"`
|
|
Buf int `json:"buf,omitempty"`
|
|
Workers int `json:"workers,omitempty"`
|
|
RingSize uint64 `json:"ring_size,omitempty"`
|
|
}
|
|
|
|
type EdgeSpec struct {
|
|
Source string `json:"source"`
|
|
Targets []string `json:"targets"`
|
|
Type string `json:"type,omitempty"`
|
|
}
|
|
|
|
type PipelineSpec struct {
|
|
Nodes []NodeSpec `json:"nodes"`
|
|
Edges []EdgeSpec `json:"edges"`
|
|
EntryIDs []string `json:"entry_ids,omitempty"`
|
|
Conditions map[string]map[string]string `json:"conditions,omitempty"`
|
|
}
|
|
|
|
type DAGPipeline struct {
|
|
nodes map[string]Node
|
|
edges map[string][]EdgeSpec
|
|
rev map[string][]string
|
|
entry []string
|
|
conditions map[string]map[string]string
|
|
}
|
|
|
|
func NewDAGPipeline() *DAGPipeline {
|
|
return &DAGPipeline{
|
|
nodes: map[string]Node{},
|
|
edges: map[string][]EdgeSpec{},
|
|
rev: map[string][]string{},
|
|
conditions: map[string]map[string]string{},
|
|
}
|
|
}
|
|
|
|
func (d *DAGPipeline) AddNode(n Node) {
|
|
d.nodes[n.ID()] = n
|
|
}
|
|
|
|
func (d *DAGPipeline) AddEdge(from string, tos []string, typ string) {
|
|
if typ == "" {
|
|
typ = "simple"
|
|
}
|
|
e := EdgeSpec{Source: from, Targets: tos, Type: typ}
|
|
d.edges[from] = append(d.edges[from], e)
|
|
for _, to := range tos {
|
|
d.rev[to] = append(d.rev[to], from)
|
|
}
|
|
}
|
|
|
|
func (d *DAGPipeline) AddCondition(id string, cond map[string]string) {
|
|
d.conditions[id] = cond
|
|
for _, to := range cond {
|
|
d.rev[to] = append(d.rev[to], id)
|
|
}
|
|
}
|
|
|
|
func (d *DAGPipeline) Start(ctx context.Context, inputs <-chan any) (<-chan any, error) {
|
|
nCh := map[string]chan any{}
|
|
outCh := map[string]<-chan any{}
|
|
wgMap := map[string]*sync.WaitGroup{}
|
|
for id := range d.nodes {
|
|
nCh[id] = make(chan any, 128)
|
|
wgMap[id] = &sync.WaitGroup{}
|
|
}
|
|
if len(d.entry) == 0 {
|
|
for id := range d.nodes {
|
|
if len(d.rev[id]) == 0 {
|
|
d.entry = append(d.entry, id)
|
|
}
|
|
}
|
|
}
|
|
for id, node := range d.nodes {
|
|
in := nCh[id]
|
|
out := node.Start(ctx, in)
|
|
outCh[id] = out
|
|
if cond, ok := d.conditions[id]; ok {
|
|
go func(o <-chan any, cond map[string]string) {
|
|
for v := range o {
|
|
if m, ok := v.(map[string]any); ok {
|
|
if status, ok := m["condition_status"].(string); ok {
|
|
if target, ok := cond[status]; ok {
|
|
wgMap[target].Add(1)
|
|
go func(c chan any, v any, wg *sync.WaitGroup) {
|
|
defer wg.Done()
|
|
select {
|
|
case c <- v:
|
|
case <-ctx.Done():
|
|
}
|
|
}(nCh[target], v, wgMap[target])
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}(out, cond)
|
|
} else {
|
|
for _, e := range d.edges[id] {
|
|
for _, dep := range e.Targets {
|
|
if e.Type == "iterator" {
|
|
go func(o <-chan any, c chan any, wg *sync.WaitGroup) {
|
|
for v := range o {
|
|
if arr, ok := v.([]any); ok {
|
|
for _, item := range arr {
|
|
wg.Add(1)
|
|
go func(item any) {
|
|
defer wg.Done()
|
|
select {
|
|
case c <- item:
|
|
case <-ctx.Done():
|
|
}
|
|
}(item)
|
|
}
|
|
}
|
|
}
|
|
}(out, nCh[dep], wgMap[dep])
|
|
} else {
|
|
wgMap[dep].Add(1)
|
|
go func(o <-chan any, c chan any, wg *sync.WaitGroup) {
|
|
defer wg.Done()
|
|
for v := range o {
|
|
select {
|
|
case c <- v:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}(out, nCh[dep], wgMap[dep])
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
for _, id := range d.entry {
|
|
wgMap[id].Add(1)
|
|
}
|
|
go func() {
|
|
defer func() {
|
|
for _, id := range d.entry {
|
|
wgMap[id].Done()
|
|
}
|
|
}()
|
|
for v := range inputs {
|
|
for _, id := range d.entry {
|
|
select {
|
|
case nCh[id] <- v:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
for id, wg := range wgMap {
|
|
go func(id string, wg *sync.WaitGroup, ch chan any) {
|
|
time.Sleep(time.Millisecond)
|
|
wg.Wait()
|
|
close(ch)
|
|
}(id, wg, nCh[id])
|
|
}
|
|
finalOut := make(chan any, 128)
|
|
var wg sync.WaitGroup
|
|
for id := range d.nodes {
|
|
if len(d.edges[id]) == 0 && len(d.conditions[id]) == 0 {
|
|
wg.Add(1)
|
|
go func(o <-chan any) {
|
|
defer wg.Done()
|
|
for v := range o {
|
|
select {
|
|
case finalOut <- v:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}(outCh[id])
|
|
}
|
|
}
|
|
go func() {
|
|
wg.Wait()
|
|
close(finalOut)
|
|
}()
|
|
return finalOut, nil
|
|
}
|
|
|
|
func BuildDAGFromSpec(spec PipelineSpec) (*DAGPipeline, error) {
|
|
d := NewDAGPipeline()
|
|
for _, ns := range spec.Nodes {
|
|
proc, ok := GetProcessor(ns.Processor)
|
|
if !ok {
|
|
return nil, fmt.Errorf("processor %s not registered", ns.Processor)
|
|
}
|
|
var node Node
|
|
switch ns.Type {
|
|
case "channel":
|
|
node = NewChannelNode(ns.ID, proc, ns.Buf, ns.Workers)
|
|
case "ring":
|
|
node = NewRingNode(ns.ID, proc, ns.RingSize)
|
|
case "page":
|
|
node = NewPageNode(ns.ID, proc, ns.Buf, ns.Workers)
|
|
default:
|
|
return nil, fmt.Errorf("unknown node type %s", ns.Type)
|
|
}
|
|
d.AddNode(node)
|
|
}
|
|
for _, e := range spec.Edges {
|
|
if _, ok := d.nodes[e.Source]; !ok {
|
|
return nil, fmt.Errorf("edge source %s not found", e.Source)
|
|
}
|
|
for _, tgt := range e.Targets {
|
|
if _, ok := d.nodes[tgt]; !ok {
|
|
return nil, fmt.Errorf("edge target %s not found", tgt)
|
|
}
|
|
}
|
|
d.AddEdge(e.Source, e.Targets, e.Type)
|
|
}
|
|
if len(spec.EntryIDs) > 0 {
|
|
d.entry = spec.EntryIDs
|
|
}
|
|
if spec.Conditions != nil {
|
|
for id, cond := range spec.Conditions {
|
|
d.AddCondition(id, cond)
|
|
}
|
|
}
|
|
return d, nil
|
|
}
|
|
|
|
// ---------------------- Example Processors ----------------------
|
|
|
|
func doubleProc(ctx context.Context, in any) (any, error) {
|
|
switch v := in.(type) {
|
|
case int:
|
|
return v * 2, nil
|
|
case float64:
|
|
return v * 2, nil
|
|
default:
|
|
return nil, errors.New("unsupported type for double")
|
|
}
|
|
}
|
|
|
|
func incProc(ctx context.Context, in any) (any, error) {
|
|
if n, ok := in.(int); ok {
|
|
return n + 1, nil
|
|
}
|
|
return nil, errors.New("inc: not int")
|
|
}
|
|
|
|
func printProc(ctx context.Context, in any) (any, error) {
|
|
fmt.Printf("OUTPUT: %#v\n", in)
|
|
return in, nil
|
|
}
|
|
|
|
func getDataProc(ctx context.Context, in any) (any, error) {
|
|
return in, nil
|
|
}
|
|
|
|
func loopProc(ctx context.Context, in any) (any, error) {
|
|
return in, nil
|
|
}
|
|
|
|
func validateAgeProc(ctx context.Context, in any) (any, error) {
|
|
m, ok := in.(map[string]any)
|
|
if !ok {
|
|
return nil, errors.New("not map")
|
|
}
|
|
age, ok := m["age"].(float64)
|
|
if !ok {
|
|
return nil, errors.New("no age")
|
|
}
|
|
status := "default"
|
|
if age >= 18 {
|
|
status = "pass"
|
|
}
|
|
m["condition_status"] = status
|
|
return m, nil
|
|
}
|
|
|
|
func validateGenderProc(ctx context.Context, in any) (any, error) {
|
|
m, ok := in.(map[string]any)
|
|
if !ok {
|
|
return nil, errors.New("not map")
|
|
}
|
|
gender, ok := m["gender"].(string)
|
|
if !ok {
|
|
return nil, errors.New("no gender")
|
|
}
|
|
m["female_voter"] = gender == "female"
|
|
return m, nil
|
|
}
|
|
|
|
func finalProc(ctx context.Context, in any) (any, error) {
|
|
m, ok := in.(map[string]any)
|
|
if !ok {
|
|
return nil, errors.New("not map")
|
|
}
|
|
m["done"] = true
|
|
return m, nil
|
|
}
|
|
|
|
// ---------------------- Main Demo ----------------------
|
|
|
|
func main() {
|
|
RegisterProcessor("double", doubleProc)
|
|
RegisterProcessor("inc", incProc)
|
|
RegisterProcessor("print", printProc)
|
|
RegisterProcessor("getData", getDataProc)
|
|
RegisterProcessor("loop", loopProc)
|
|
RegisterProcessor("validateAge", validateAgeProc)
|
|
RegisterProcessor("validateGender", validateGenderProc)
|
|
RegisterProcessor("final", finalProc)
|
|
|
|
jsonSpec := `{
|
|
"nodes": [
|
|
{"id":"getData","type":"channel","processor":"getData"},
|
|
{"id":"loop","type":"channel","processor":"loop"},
|
|
{"id":"validateAge","type":"channel","processor":"validateAge"},
|
|
{"id":"validateGender","type":"channel","processor":"validateGender"},
|
|
{"id":"final","type":"channel","processor":"final"}
|
|
],
|
|
"edges": [
|
|
{"source":"getData","targets":["loop"],"type":"simple"},
|
|
{"source":"loop","targets":["validateAge"],"type":"iterator"},
|
|
{"source":"validateGender","targets":["final"],"type":"simple"}
|
|
],
|
|
"entry_ids":["getData"],
|
|
"conditions": {
|
|
"validateAge": {"pass": "validateGender", "default": "final"}
|
|
}
|
|
}`
|
|
|
|
var spec PipelineSpec
|
|
if err := json.Unmarshal([]byte(jsonSpec), &spec); err != nil {
|
|
panic(err)
|
|
}
|
|
dag, err := BuildDAGFromSpec(spec)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
in := make(chan any)
|
|
ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
|
defer cancel()
|
|
|
|
out, err := dag.Start(ctx, in)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
go func() {
|
|
data := []any{
|
|
map[string]any{"age": 15.0, "gender": "female"},
|
|
map[string]any{"age": 18.0, "gender": "male"},
|
|
}
|
|
in <- data
|
|
close(in)
|
|
}()
|
|
|
|
var results []any
|
|
for r := range out {
|
|
results = append(results, r)
|
|
}
|
|
|
|
fmt.Println("Final results:", results)
|
|
fmt.Println("pipeline finished")
|
|
}
|