initial import

Author: kevwan

Diff for: .gitattributes

@@ -0,0 +1 @@
+* text=auto eol=lf

Diff for: .gitignore

@@ -0,0 +1,22 @@
+# Ignore all
+*
+
+# Unignore all with extensions
+!*.*
+!**/Dockerfile
+
+# Unignore all dirs
+!*/
+!api
+
+.idea
+**/.DS_Store
+**/logs
+!Makefile
+
+# gitlab ci
+.cache
+
+# vim auto backup file
+*~
+!OWNERS

Diff for: dq/config.go

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

Diff for: dq/connection.go

@@ -0,0 +1,65 @@
+package dq
+
+import (
+	"sync"
+
+	"github.com/beanstalkd/go-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
+	}
+}

Diff for: dq/consumer.go

@@ -0,0 +1,100 @@
+package dq
+
+import (
+	"strconv"
+	"time"
+
+	"github.com/tal-tech/go-zero/core/hash"
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/service"
+	"github.com/tal-tech/go-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
+}

Diff for: dq/consumernode.go

@@ -0,0 +1,94 @@
+package dq
+
+import (
+	"time"
+
+	"github.com/beanstalkd/go-beanstalk"
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/syncx"
+)
+
+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()
+}