feat: use unbuffered channel to handle new task

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

Author: appleboy

_example/server-client/client/main.go

@@ -6,8 +6,10 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/appleboy/graceful"
 	"github.com/golang-queue/nats"
 	"github.com/golang-queue/queue"
+	"github.com/golang-queue/queue/core"
 )
 
 type job struct {
@@ -26,33 +28,51 @@ func main() {
 	taskN := 10000
 	rets := make(chan string, taskN)
 
+	m := graceful.NewManager()
+
 	// define the worker
 	w := nats.NewWorker(
 		nats.WithAddr("127.0.0.1:4222"),
 		nats.WithSubj("example"),
 		nats.WithQueue("foobar"),
-		nats.WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+		nats.WithRunFunc(func(ctx context.Context, m core.QueuedMessage) error {
 			var v *job
 			if err := json.Unmarshal(m.Bytes(), &v); err != nil {
 				return err
 			}
 			rets <- v.Message
+			time.Sleep(2 * time.Second)
 			return nil
 		}),
 	)
-
 	// define the queue
 	q := queue.NewPool(
-		5,
+		1,
 		queue.WithWorker(w),
 	)
 
-	// wait until all tasks done
-	for i := 0; i < taskN; i++ {
-		fmt.Println("message:", <-rets)
-		time.Sleep(50 * time.Millisecond)
-	}
+	m.AddRunningJob(func(ctx context.Context) error {
+		for {
+			select {
+			case <-ctx.Done():
+				select {
+				case m := <-rets:
+					fmt.Println("message:", m)
+				default:
+				}
+				return nil
+			case m := <-rets:
+				fmt.Println("message:", m)
+				time.Sleep(50 * time.Millisecond)
+			}
+		}
+	})
+
+	m.AddShutdownJob(func() error {
+		// shutdown the service and notify all the worker
+		q.Release()
+		return nil
+	})
 
-	// shutdown the service and notify all the worker
-	q.Release()
+	<-m.Done()
 }

_example/server-client/go.mod

@@ -3,8 +3,9 @@ module example
 go 1.18
 
 require (
+	github.com/appleboy/graceful v0.0.4
 	github.com/golang-queue/nats v0.0.3-0.20210907015837-3e2e4b448b3d
-	github.com/golang-queue/queue v0.0.13-0.20220403053548-d431277d570f
+	github.com/golang-queue/queue v0.0.13-0.20220423025512-c4a8df54c917
 )
 
 require (

_example/server-client/go.sum

@@ -1,7 +1,10 @@
+github.com/appleboy/graceful v0.0.4 h1:Q4LCeq4DFy59qiACLtuH+mSqDERtUzwkQbCWpRaWwvQ=
+github.com/appleboy/graceful v0.0.4/go.mod h1:Q2mVx0t+N0lCDZc5MJudbcpTm6cgGM/J2gZCZIqD9dc=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/golang-queue/queue v0.0.13-0.20220403053548-d431277d570f h1:Wioq3g97ssizNPQsPwdL61DIjePabPaq+XYo7z2t2Oc=
-github.com/golang-queue/queue v0.0.13-0.20220403053548-d431277d570f/go.mod h1:KD9age1s6nk8Evz3tfKHsk8k4LwA0htxQ7MS7rJPJzA=
+github.com/golang-queue/queue v0.0.13-0.20220423025512-c4a8df54c917 h1:+khkGHxQPsad/mfmgizoTc3Jh5UVdCp1OfUkDfW+uDQ=
+github.com/golang-queue/queue v0.0.13-0.20220423025512-c4a8df54c917/go.mod h1:g1yxxDl8JMo4gUfxt11fjjU3SXU1ah61EvwshmDoSIs=
+github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
 github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
 github.com/klauspost/compress v1.15.1 h1:y9FcTHGyrebwfP0ZZqFiaxTaiDnUrGkJkI+f583BL1A=
 github.com/klauspost/compress v1.15.1/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=

_example/server-client/server/main.go

@@ -23,7 +23,7 @@ func (j *job) Bytes() []byte {
 }
 
 func main() {
-	taskN := 100
+	taskN := 10
 
 	// define the worker
 	w := nats.NewWorker(

nats.go

@@ -34,7 +34,7 @@ func NewWorker(opts ...Option) *Worker {
 		opts:  newOptions(opts...),
 		stop:  make(chan struct{}),
 		exit:  make(chan struct{}),
-		tasks: make(chan *nats.Msg, 1),
+		tasks: make(chan *nats.Msg),
 	}
 
 	w.client, err = nats.Connect(w.opts.addr)
@@ -59,13 +59,13 @@ func (w *Worker) startConsumer() error {
 		case w.tasks <- msg:
 		case <-w.stop:
 			if msg != nil {
-				// re-queue the job if worker has been shutdown.
-				w.opts.logger.Info("re-queue the old job")
+				// re-queue the task if worker has been shutdown.
+				w.opts.logger.Info("re-queue the current task")
 				if err := w.client.Publish(w.opts.subj, msg.Data); err != nil {
-					panic(err)
+					w.opts.logger.Errorf("error to re-queue the current task: %s", err.Error())
 				}
-				close(w.exit)
 			}
+			close(w.exit)
 		}
 	})