chore: Support new pool (#33)

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

pool.go

@@ -0,0 +1,21 @@
+package queue
+
+func NewPool(size int, opts ...Option) *Queue {
+	o := []Option{
+		WithWorkerCount(size),
+		WithWorker(NewConsumer(opts...)),
+	}
+	o = append(
+		o,
+		opts...,
+	)
+
+	q, err := NewQueue(o...)
+	if err != nil {
+		panic(err)
+	}
+
+	q.Start()
+
+	return q
+}

pool_test.go

@@ -0,0 +1,34 @@
+package queue
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewPoolWithQueueTask(t *testing.T) {
+	totalN := 5
+	taskN := 100
+	rets := make(chan struct{}, taskN)
+
+	p := NewPool(totalN)
+	time.Sleep(time.Millisecond * 50)
+	assert.Equal(t, totalN, p.Workers())
+
+	for i := 0; i < taskN; i++ {
+		assert.NoError(t, p.QueueTask(func(context.Context) error {
+			rets <- struct{}{}
+			return nil
+		}))
+	}
+
+	for i := 0; i < taskN; i++ {
+		<-rets
+	}
+
+	// shutdown all, and now running worker is 0
+	p.Release()
+	assert.Equal(t, 0, p.Workers())
+}

queue.go

@@ -99,6 +99,12 @@ func (q *Queue) Shutdown() {
 	})
 }
 
+// Workers returns the numbers of workers has been created.
+func (q *Queue) Release() {
+	q.Shutdown()
+	q.Wait()
+}
+
 // Workers returns the numbers of workers has been created.
 func (q *Queue) Workers() int {
 	return int(atomic.LoadInt32(&q.runningWorkers))