diff --git a/consumer.go b/consumer.go index c61b28b..f160933 100644 --- a/consumer.go +++ b/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 diff --git a/consumer_test.go b/consumer_test.go index 78fe3a6..0cfccda 100644 --- a/consumer_test.go +++ b/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) +} diff --git a/options.go b/options.go index ba0b51a..fc9c306 100644 --- a/options.go +++ b/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