chore(consumer): support request timeout flag

consumer.go

@@ -25,6 +25,8 @@ type Consumer struct {
 	logger    Logger
 	stopOnce  sync.Once
 	stopFlag  int32
+
+	requestTimeout time.Duration
 }
 
 func (s *Consumer) handle(m *job.Message) error {
@@ -139,28 +141,19 @@ func (s *Consumer) Queue(task core.QueuedMessage) error {
 
 // Request a new task from channel
 func (s *Consumer) Request() (core.QueuedMessage, error) {
-	clock := 0
-loop:
-	for {
-		select {
-		case task, ok := <-s.taskQueue:
-			if !ok {
-				select {
-				case s.exit <- struct{}{}:
-				default:
-				}
-				return nil, ErrQueueHasBeenClosed
-			}
-			return task, nil
-		case <-time.After(1 * time.Second):
-			if clock == 5 {
-				break loop
+	select {
+	case task, ok := <-s.taskQueue:
+		if !ok {
+			select {
+			case s.exit <- struct{}{}:
+			default:
 			}
-			clock += 1
+			return nil, ErrQueueHasBeenClosed
 		}
+		return task, nil
+	case <-time.After(s.requestTimeout):
+		return nil, ErrNoTaskInQueue
 	}
-
-	return nil, ErrNoTaskInQueue
 }
 
 // NewConsumer for create new consumer instance
@@ -172,6 +165,8 @@ func NewConsumer(opts ...Option) *Consumer {
 		exit:      make(chan struct{}),
 		logger:    o.logger,
 		runFunc:   o.fn,
+
+		requestTimeout: o.requestTimeout,
 	}
 
 	return w

consumer_test.go

@@ -603,3 +603,16 @@ func TestCancelRetryCountWithNewMessage(t *testing.T) {
 	assert.Len(t, messages, 0)
 	assert.Equal(t, 2, count)
 }
+
+func TestRequestTimeout(t *testing.T) {
+	w := NewConsumer(
+		WithFn(func(ctx context.Context, m core.QueuedMessage) error {
+			return nil
+		}),
+		WithRequestTimeout(10*time.Millisecond),
+	)
+	task, err := w.Request()
+	assert.Nil(t, task)
+	assert.Error(t, err)
+	assert.Equal(t, ErrNoTaskInQueue, err)
+}

options.go

@@ -3,16 +3,18 @@ package queue
 import (
 	"context"
 	"runtime"
+	"time"
 
 	"github.com/golang-queue/queue/core"
 )
 
 var (
-	defaultQueueSize   = 4096
-	defaultWorkerCount = runtime.NumCPU()
-	defaultNewLogger   = NewLogger()
-	defaultFn          = func(context.Context, core.QueuedMessage) error { return nil }
-	defaultMetric      = NewMetric()
+	defaultQueueSize      = 4096
+	defaultWorkerCount    = runtime.NumCPU()
+	defaultNewLogger      = NewLogger()
+	defaultFn             = func(context.Context, core.QueuedMessage) error { return nil }
+	defaultMetric         = NewMetric()
+	defaultRequestTimeout = 5 * time.Second
 )
 
 // An Option configures a mutex.
@@ -42,6 +44,13 @@ func WithQueueSize(num int) Option {
 	})
 }
 
+// WithQueueSize set worker count
+func WithRequestTimeout(timeout time.Duration) Option {
+	return OptionFunc(func(q *Options) {
+		q.requestTimeout = timeout
+	})
+}
+
 // WithLogger set custom logger
 func WithLogger(l Logger) Option {
 	return OptionFunc(func(q *Options) {
@@ -78,6 +87,9 @@ type Options struct {
 	worker      core.Worker
 	fn          func(context.Context, core.QueuedMessage) error
 	metric      Metric
+
+	// timeout for request single task
+	requestTimeout time.Duration
 }
 
 // NewOptions initialize the default value for the options
@@ -89,6 +101,8 @@ func NewOptions(opts ...Option) *Options {
 		worker:      nil,
 		fn:          defaultFn,
 		metric:      defaultMetric,
+
+		requestTimeout: defaultRequestTimeout,
 	}
 
 	// Loop through each option