initial import

dq/config.go

@@ -0,0 +1,15 @@
+package dq
+
+import "zero/core/stores/redis"
+
+type (
+	Beanstalk struct {
+		Endpoint string
+		Tube     string
+	}
+
+	DqConf struct {
+		Beanstalks []Beanstalk
+		Redis      redis.RedisConf
+	}
+)

dq/connection.go

@@ -0,0 +1,65 @@
+package dq
+
+import (
+	"sync"
+
+	"github.com/beanstalkd/beanstalk"
+)
+
+type connection struct {
+	lock     sync.RWMutex
+	endpoint string
+	tube     string
+	conn     *beanstalk.Conn
+}
+
+func newConnection(endpint, tube string) *connection {
+	return &connection{
+		endpoint: endpint,
+		tube:     tube,
+	}
+}
+
+func (c *connection) Close() error {
+	c.lock.Lock()
+	conn := c.conn
+	c.conn = nil
+	defer c.lock.Unlock()
+
+	if conn != nil {
+		return conn.Close()
+	}
+
+	return nil
+}
+
+func (c *connection) get() (*beanstalk.Conn, error) {
+	c.lock.RLock()
+	conn := c.conn
+	c.lock.RUnlock()
+	if conn != nil {
+		return conn, nil
+	}
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	var err error
+	c.conn, err = beanstalk.Dial("tcp", c.endpoint)
+	if err != nil {
+		return nil, err
+	}
+
+	c.conn.Tube.Name = c.tube
+	return c.conn, err
+}
+
+func (c *connection) reset() {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	if c.conn != nil {
+		c.conn.Close()
+		c.conn = nil
+	}
+}

dq/consumer.go

@@ -0,0 +1,100 @@
+package dq
+
+import (
+	"strconv"
+	"time"
+
+	"zero/core/hash"
+	"zero/core/logx"
+	"zero/core/service"
+	"zero/core/stores/redis"
+)
+
+const (
+	expiration = 3600 // seconds
+	guardValue = "1"
+	tolerance  = time.Minute * 30
+)
+
+var maxCheckBytes = getMaxTimeLen()
+
+type (
+	Consume func(body []byte)
+
+	Consumer interface {
+		Consume(consume Consume)
+	}
+
+	consumerCluster struct {
+		nodes []*consumerNode
+		red   *redis.Redis
+	}
+)
+
+func NewConsumer(c DqConf) Consumer {
+	var nodes []*consumerNode
+	for _, node := range c.Beanstalks {
+		nodes = append(nodes, newConsumerNode(node.Endpoint, node.Tube))
+	}
+	return &consumerCluster{
+		nodes: nodes,
+		red:   c.Redis.NewRedis(),
+	}
+}
+
+func (c *consumerCluster) Consume(consume Consume) {
+	guardedConsume := func(body []byte) {
+		key := hash.Md5Hex(body)
+		body, ok := c.unwrap(body)
+		if !ok {
+			logx.Errorf("discarded: %q", string(body))
+			return
+		}
+
+		ok, err := c.red.SetnxEx(key, guardValue, expiration)
+		if err != nil {
+			logx.Error(err)
+		} else if ok {
+			consume(body)
+		}
+	}
+
+	group := service.NewServiceGroup()
+	for _, node := range c.nodes {
+		group.Add(consumeService{
+			c:       node,
+			consume: guardedConsume,
+		})
+	}
+	group.Start()
+}
+
+func (c *consumerCluster) unwrap(body []byte) ([]byte, bool) {
+	var pos = -1
+	for i := 0; i < maxCheckBytes; i++ {
+		if body[i] == timeSep {
+			pos = i
+			break
+		}
+	}
+	if pos < 0 {
+		return nil, false
+	}
+
+	val, err := strconv.ParseInt(string(body[:pos]), 10, 64)
+	if err != nil {
+		logx.Error(err)
+		return nil, false
+	}
+
+	t := time.Unix(0, val)
+	if t.Add(tolerance).Before(time.Now()) {
+		return nil, false
+	}
+
+	return body[pos+1:], true
+}
+
+func getMaxTimeLen() int {
+	return len(strconv.FormatInt(time.Now().UnixNano(), 10)) + 2
+}

dq/consumernode.go

@@ -0,0 +1,95 @@
+package dq
+
+import (
+	"time"
+
+	"zero/core/logx"
+	"zero/core/syncx"
+
+	"github.com/beanstalkd/beanstalk"
+)
+
+type (
+	consumerNode struct {
+		conn *connection
+		tube string
+		on   *syncx.AtomicBool
+	}
+
+	consumeService struct {
+		c       *consumerNode
+		consume Consume
+	}
+)
+
+func newConsumerNode(endpoint, tube string) *consumerNode {
+	return &consumerNode{
+		conn: newConnection(endpoint, tube),
+		tube: tube,
+		on:   syncx.ForAtomicBool(true),
+	}
+}
+
+func (c *consumerNode) dispose() {
+	c.on.Set(false)
+}
+
+func (c *consumerNode) consumeEvents(consume Consume) {
+	for c.on.True() {
+		conn, err := c.conn.get()
+		if err != nil {
+			logx.Error(err)
+			time.Sleep(time.Second)
+			continue
+		}
+
+		// because getting conn takes at most one second, reserve tasks at most 5 seconds,
+		// if don't check on/off here, the conn might not be closed due to
+		// graceful shutdon waits at most 5.5 seconds.
+		if !c.on.True() {
+			break
+		}
+
+		conn.Tube.Name = c.tube
+		conn.TubeSet.Name[c.tube] = true
+		id, body, err := conn.Reserve(reserveTimeout)
+		if err == nil {
+			conn.Delete(id)
+			consume(body)
+			continue
+		}
+
+		// the error can only be beanstalk.NameError or beanstalk.ConnError
+		switch cerr := err.(type) {
+		case beanstalk.ConnError:
+			switch cerr.Err {
+			case beanstalk.ErrTimeout:
+				// timeout error on timeout, just continue the loop
+			case beanstalk.ErrBadChar, beanstalk.ErrBadFormat, beanstalk.ErrBuried, beanstalk.ErrDeadline,
+				beanstalk.ErrDraining, beanstalk.ErrEmpty, beanstalk.ErrInternal, beanstalk.ErrJobTooBig,
+				beanstalk.ErrNoCRLF, beanstalk.ErrNotFound, beanstalk.ErrNotIgnored, beanstalk.ErrTooLong:
+				// won't reset
+				logx.Error(err)
+			default:
+				// beanstalk.ErrOOM, beanstalk.ErrUnknown and other errors
+				logx.Error(err)
+				c.conn.reset()
+				time.Sleep(time.Second)
+			}
+		default:
+			logx.Error(err)
+		}
+	}
+
+	if err := c.conn.Close(); err != nil {
+		logx.Error(err)
+	}
+}
+
+func (cs consumeService) Start() {
+	cs.c.consumeEvents(cs.consume)
+}
+
+func (cs consumeService) Stop() {
+	cs.c.dispose()
+}

dq/producer.go

@@ -0,0 +1,156 @@
+package dq
+
+import (
+	"bytes"
+	"log"
+	"math/rand"
+	"strconv"
+	"strings"
+	"time"
+
+	"zero/core/errorx"
+	"zero/core/fx"
+	"zero/core/logx"
+)
+
+const (
+	replicaNodes    = 3
+	minWrittenNodes = 2
+)
+
+type (
+	Producer interface {
+		At(body []byte, at time.Time) (string, error)
+		Close() error
+		Delay(body []byte, delay time.Duration) (string, error)
+		Revoke(ids string) error
+	}
+
+	producerCluster struct {
+		nodes []Producer
+	}
+)
+
+func init() {
+	rand.Seed(time.Now().UnixNano())
+}
+
+func NewProducer(beanstalks []Beanstalk) Producer {
+	if len(beanstalks) < minWrittenNodes {
+		log.Fatalf("nodes must be equal or greater than %d", minWrittenNodes)
+	}
+
+	var nodes []Producer
+	for _, node := range beanstalks {
+		nodes = append(nodes, NewProducerNode(node.Endpoint, node.Tube))
+	}
+	return &producerCluster{nodes: nodes}
+}
+
+func (p *producerCluster) At(body []byte, at time.Time) (string, error) {
+	return p.insert(func(node Producer) (string, error) {
+		return node.At(p.wrap(body, at), at)
+	})
+}
+
+func (p *producerCluster) Close() error {
+	var be errorx.BatchError
+	for _, node := range p.nodes {
+		if err := node.Close(); err != nil {
+			be.Add(err)
+		}
+	}
+	return be.Err()
+}
+
+func (p *producerCluster) Delay(body []byte, delay time.Duration) (string, error) {
+	return p.insert(func(node Producer) (string, error) {
+		return node.Delay(p.wrap(body, time.Now().Add(delay)), delay)
+	})
+}
+
+func (p *producerCluster) Revoke(ids string) error {
+	var be errorx.BatchError
+
+	fx.From(func(source chan<- interface{}) {
+		for _, node := range p.nodes {
+			source <- node
+		}
+	}).Map(func(item interface{}) interface{} {
+		node := item.(Producer)
+		return node.Revoke(ids)
+	}).ForEach(func(item interface{}) {
+		if item != nil {
+			be.Add(item.(error))
+		}
+	})
+
+	return be.Err()
+}
+
+func (p *producerCluster) cloneNodes() []Producer {
+	return append([]Producer(nil), p.nodes...)
+}
+
+func (p *producerCluster) getWriteNodes() []Producer {
+	if len(p.nodes) <= replicaNodes {
+		return p.nodes
+	}
+
+	nodes := p.cloneNodes()
+	rand.Shuffle(len(nodes), func(i, j int) {
+		nodes[i], nodes[j] = nodes[j], nodes[i]
+	})
+	return nodes[:replicaNodes]
+}
+
+func (p *producerCluster) insert(fn func(node Producer) (string, error)) (string, error) {
+	type idErr struct {
+		id  string
+		err error
+	}
+	var ret []idErr
+	fx.From(func(source chan<- interface{}) {
+		for _, node := range p.getWriteNodes() {
+			source <- node
+		}
+	}).Map(func(item interface{}) interface{} {
+		node := item.(Producer)
+		id, err := fn(node)
+		return idErr{
+			id:  id,
+			err: err,
+		}
+	}).ForEach(func(item interface{}) {
+		ret = append(ret, item.(idErr))
+	})
+
+	var ids []string
+	var be errorx.BatchError
+	for _, val := range ret {
+		if val.err != nil {
+			be.Add(val.err)
+		} else {
+			ids = append(ids, val.id)
+		}
+	}
+
+	jointId := strings.Join(ids, idSep)
+	if len(ids) >= minWrittenNodes {
+		return jointId, nil
+	}
+
+	if err := p.Revoke(jointId); err != nil {
+		logx.Error(err)
+	}
+
+	return "", be.Err()
+}
+
+func (p *producerCluster) wrap(body []byte, at time.Time) []byte {
+	var builder bytes.Buffer
+	builder.WriteString(strconv.FormatInt(at.UnixNano(), 10))
+	builder.WriteByte(timeSep)
+	builder.Write(body)
+	return builder.Bytes()
+}

dq/producernode.go

@@ -0,0 +1,98 @@
+package dq
+
+import (
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/beanstalkd/beanstalk"
+)
+
+var ErrTimeBeforeNow = errors.New("can't schedule task to past time")
+
+type producerNode struct {
+	endpoint string
+	tube     string
+	conn     *connection
+}
+
+func NewProducerNode(endpoint, tube string) Producer {
+	return &producerNode{
+		endpoint: endpoint,
+		tube:     tube,
+		conn:     newConnection(endpoint, tube),
+	}
+}
+
+func (p *producerNode) At(body []byte, at time.Time) (string, error) {
+	now := time.Now()
+	if at.Before(now) {
+		return "", ErrTimeBeforeNow
+	}
+
+	duration := at.Sub(now)
+	return p.Delay(body, duration)
+}
+
+func (p *producerNode) Close() error {
+	return p.conn.Close()
+}
+
+func (p *producerNode) Delay(body []byte, delay time.Duration) (string, error) {
+	conn, err := p.conn.get()
+	if err != nil {
+		return "", err
+	}
+
+	id, err := conn.Put(body, PriNormal, delay, defaultTimeToRun)
+	if err == nil {
+		return fmt.Sprintf("%s/%s/%d", p.endpoint, p.tube, id), nil
+	}
+
+	// the error can only be beanstalk.NameError or beanstalk.ConnError
+	// just return when the error is beanstalk.NameError, don't reset
+	switch cerr := err.(type) {
+	case beanstalk.ConnError:
+		switch cerr.Err {
+		case beanstalk.ErrBadChar, beanstalk.ErrBadFormat, beanstalk.ErrBuried, beanstalk.ErrDeadline,
+			beanstalk.ErrDraining, beanstalk.ErrEmpty, beanstalk.ErrInternal, beanstalk.ErrJobTooBig,
+			beanstalk.ErrNoCRLF, beanstalk.ErrNotFound, beanstalk.ErrNotIgnored, beanstalk.ErrTooLong:
+			// won't reset
+		default:
+			// beanstalk.ErrOOM, beanstalk.ErrTimeout, beanstalk.ErrUnknown and other errors
+			p.conn.reset()
+		}
+	}
+
+	return "", err
+}
+
+func (p *producerNode) Revoke(jointId string) error {
+	ids := strings.Split(jointId, idSep)
+	for _, id := range ids {
+		fields := strings.Split(id, "/")
+		if len(fields) < 3 {
+			continue
+		}
+		if fields[0] != p.endpoint || fields[1] != p.tube {
+			continue
+		}
+
+		conn, err := p.conn.get()
+		if err != nil {
+			return err
+		}
+
+		n, err := strconv.ParseUint(fields[2], 10, 64)
+		if err != nil {
+			return err
+		}
+
+		return conn.Delete(n)
+	}
+
+	// if not in this beanstalk, ignore
+	return nil
+}

dq/vars.go

@@ -0,0 +1,15 @@
+package dq
+
+import "time"
+
+const (
+	PriHigh   = 1
+	PriNormal = 2
+	PriLow    = 3
+
+	defaultTimeToRun = time.Second * 5
+	reserveTimeout   = time.Second * 5
+
+	idSep   = ","
+	timeSep = '/'
+)