add queue package

core/queue/balancedqueuepusher.go

@@ -0,0 +1,44 @@
+package queue
+
+import (
+	"errors"
+	"sync/atomic"
+
+	"github.com/tal-tech/go-zero/core/logx"
+)
+
+var ErrNoAvailablePusher = errors.New("no available pusher")
+
+type BalancedQueuePusher struct {
+	name    string
+	pushers []Pusher
+	index   uint64
+}
+
+func NewBalancedQueuePusher(pushers []Pusher) Pusher {
+	return &BalancedQueuePusher{
+		name:    generateName(pushers),
+		pushers: pushers,
+	}
+}
+
+func (pusher *BalancedQueuePusher) Name() string {
+	return pusher.name
+}
+
+func (pusher *BalancedQueuePusher) Push(message string) error {
+	size := len(pusher.pushers)
+
+	for i := 0; i < size; i++ {
+		index := atomic.AddUint64(&pusher.index, 1) % uint64(size)
+		target := pusher.pushers[index]
+
+		if err := target.Push(message); err != nil {
+			logx.Error(err)
+		} else {
+			return nil
+		}
+	}
+
+	return ErrNoAvailablePusher
+}

core/queue/balancedqueuepusher_test.go

@@ -0,0 +1,43 @@
+package queue
+
+import (
+	"fmt"
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBalancedQueuePusher(t *testing.T) {
+	const numPushers = 100
+	var pushers []Pusher
+	var mockedPushers []*mockedPusher
+	for i := 0; i < numPushers; i++ {
+		p := &mockedPusher{
+			name: "pusher:" + strconv.Itoa(i),
+		}
+		pushers = append(pushers, p)
+		mockedPushers = append(mockedPushers, p)
+	}
+
+	pusher := NewBalancedQueuePusher(pushers)
+	assert.True(t, len(pusher.Name()) > 0)
+
+	for i := 0; i < numPushers*1000; i++ {
+		assert.Nil(t, pusher.Push("item"))
+	}
+
+	var counts []int
+	for _, p := range mockedPushers {
+		counts = append(counts, p.count)
+	}
+	mean := calcMean(counts)
+	variance := calcVariance(mean, counts)
+	assert.True(t, variance < 100, fmt.Sprintf("too big variance - %.2f", variance))
+}
+
+func TestBalancedQueuePusher_NoAvailable(t *testing.T) {
+	pusher := NewBalancedQueuePusher(nil)
+	assert.True(t, len(pusher.Name()) == 0)
+	assert.Equal(t, ErrNoAvailablePusher, pusher.Push("item"))
+}

core/queue/consumer.go

@@ -0,0 +1,10 @@
+package queue
+
+type (
+	Consumer interface {
+		Consume(string) error
+		OnEvent(event interface{})
+	}
+
+	ConsumerFactory func() (Consumer, error)
+)

core/queue/messagequeue.go

@@ -0,0 +1,6 @@
+package queue
+
+type MessageQueue interface {
+	Start()
+	Stop()
+}

core/queue/multiqueuepusher.go

@@ -0,0 +1,31 @@
+package queue
+
+import "github.com/tal-tech/go-zero/core/errorx"
+
+type MultiQueuePusher struct {
+	name    string
+	pushers []Pusher
+}
+
+func NewMultiQueuePusher(pushers []Pusher) Pusher {
+	return &MultiQueuePusher{
+		name:    generateName(pushers),
+		pushers: pushers,
+	}
+}
+
+func (pusher *MultiQueuePusher) Name() string {
+	return pusher.name
+}
+
+func (pusher *MultiQueuePusher) Push(message string) error {
+	var batchError errorx.BatchError
+
+	for _, each := range pusher.pushers {
+		if err := each.Push(message); err != nil {
+			batchError.Add(err)
+		}
+	}
+
+	return batchError.Err()
+}

core/queue/multiqueuepusher_test.go

@@ -0,0 +1,39 @@
+package queue
+
+import (
+	"fmt"
+	"math"
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMultiQueuePusher(t *testing.T) {
+	const numPushers = 100
+	var pushers []Pusher
+	var mockedPushers []*mockedPusher
+	for i := 0; i < numPushers; i++ {
+		p := &mockedPusher{
+			name: "pusher:" + strconv.Itoa(i),
+		}
+		pushers = append(pushers, p)
+		mockedPushers = append(mockedPushers, p)
+	}
+
+	pusher := NewMultiQueuePusher(pushers)
+	assert.True(t, len(pusher.Name()) > 0)
+
+	for i := 0; i < 1000; i++ {
+		_ = pusher.Push("item")
+	}
+
+	var counts []int
+	for _, p := range mockedPushers {
+		counts = append(counts, p.count)
+	}
+	mean := calcMean(counts)
+	variance := calcVariance(mean, counts)
+	assert.True(t, math.Abs(mean-1000*(1-failProba)) < 10)
+	assert.True(t, variance < 100, fmt.Sprintf("too big variance - %.2f", variance))
+}

core/queue/producer.go

@@ -0,0 +1,15 @@
+package queue
+
+type (
+	Producer interface {
+		AddListener(listener ProduceListener)
+		Produce() (string, bool)
+	}
+
+	ProduceListener interface {
+		OnProducerPause()
+		OnProducerResume()
+	}
+
+	ProducerFactory func() (Producer, error)
+)