dag.go

package godag

import (
	"context"
	"sync"
	"time"
	// "fmt"
)

// default go impl, can replace by other go routine pool impl
var Go = func (f func()) {
	go f()
}

type Op interface {
	// global: global variable to pass through DAG
	// input: the output generated by parent (the order is the same as node.prev)
	Process(ctx context.Context, global interface{}, input ...interface{}) interface{} // pass input with the order of prev
}

type StateKey string
const NodeID = "__nodeID__"

type DAG struct {
	startNode   *Node
	mu          sync.Mutex
	activeNum   int
	taskChan    chan *Node
	doneChan    chan struct{}
	stateKeeper StateKeeper
}

func (p *DAG) Init(startNode *Node, stateKeeper StateKeeper) bool {
	p.startNode = startNode
	if stateKeeper == nil {
		p.stateKeeper = NewDefaultStateKeeper()
	} else {
		p.stateKeeper = stateKeeper
	}
	p.activeNum = 1
	p.taskChan = make(chan *Node)
	p.doneChan = make(chan struct{})
	Go(func() {
		p.taskChan <- startNode
	})
	return true
}

func (p *DAG) Execute(ctx context.Context) {
	for {
		select {
		case node := <-p.taskChan:
			Go(func (){
				p.processNode(ctx, node)
			})
		case <-p.doneChan:
			return
		}
	}
}

func (p *DAG) processNode(ctx context.Context, node *Node) {
	if ctx == nil {
		ctx = context.Background()
	}
	
	if node.timeout > 0 {
		// fmt.Println("node timeout = ", node.id, node.timeout)
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, node.timeout)
		defer cancel()
	}

	doneChan := make(chan struct{})
	startTime := time.Now()
	Go(func() {
		if node.op != nil {
			args := make([]interface{}, len(node.prev))
			for idx := range node.prev {
				// NOTE: the order of prev will result the order of args passed to op
				args[idx] = p.stateKeeper.GetInput(node.prev[idx].id, node.id) // will get the parent output as input of current
			}
			ctx = context.WithValue(ctx, StateKey(NodeID), node.id)
			global := p.stateKeeper.GetGlobal()
			output := node.op.Process(ctx, global, args...)
			if !node.isCanceled { // if timeout, no need to save output
				p.stateKeeper.SetOutput(node.id, output)
			}
		}
		close(doneChan) // close can make chan readable
	})
	if node.timeout > 0 {
		select {
		case <-ctx.Done():
			// fmt.Println("timeout", node, ctx)
			node.isCanceled = true
		case <-doneChan:
			node.isCanceled = false
		}
	} else {
		<- doneChan
		node.isCanceled = false
	}
	
	node.costTime = time.Now().Sub(startTime)
	Go(func() {
		for _, nextOne := range node.next {
			p.mu.Lock()
			nextOne.indegree--
			indegree := nextOne.indegree
			if indegree == 0 {
				p.activeNum++ // should add before chan put
			}
			p.mu.Unlock()
			if indegree == 0 {
				p.taskChan <- nextOne
			}
		}
		p.mu.Lock()
		p.activeNum--
		activeNum := p.activeNum
		p.mu.Unlock()
		if activeNum == 0 {
			close(p.doneChan)
		}
	})
}

func (d *DAG) GetStateKeeper() StateKeeper {
	return d.stateKeeper
}