Skip to content

Commit

Permalink
chore(worker): remove disable consumer flag (#13)
Browse files Browse the repository at this point in the history
  • Loading branch information
appleboy authored Apr 25, 2022
1 parent 476fc66 commit 0cbfbd1
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 49 deletions.
32 changes: 16 additions & 16 deletions nats.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ type Worker struct {
exit chan struct{}
stopFlag int32
stopOnce sync.Once
startOnce sync.Once
opts options
subscription *nats.Subscription
tasks chan *nats.Msg
Expand Down Expand Up @@ -49,24 +50,22 @@ func NewWorker(opts ...Option) *Worker {
return w
}

func (w *Worker) startConsumer() error {
if w.opts.disableConsumer {
return nil
}
var err error
w.subscription, err = w.client.QueueSubscribe(w.opts.subj, w.opts.queue, func(msg *nats.Msg) {
select {
case w.tasks <- msg:
case <-w.stop:
if msg != nil {
// 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 {
w.opts.logger.Errorf("error to re-queue the current task: %s", err.Error())
func (w *Worker) startConsumer() (err error) {
w.startOnce.Do(func() {
w.subscription, err = w.client.QueueSubscribe(w.opts.subj, w.opts.queue, func(msg *nats.Msg) {
select {
case w.tasks <- msg:
case <-w.stop:
if msg != nil {
// 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 {
w.opts.logger.Errorf("error to re-queue the current task: %s", err.Error())
}
}
close(w.exit)
}
close(w.exit)
}
})
})

return err
Expand Down Expand Up @@ -165,6 +164,7 @@ func (w *Worker) Queue(job core.QueuedMessage) error {

// Request a new task
func (w *Worker) Request() (core.QueuedMessage, error) {
_ = w.startConsumer()
clock := 0
loop:
for {
Expand Down
20 changes: 0 additions & 20 deletions nats_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -372,23 +372,3 @@ func TestReQueueTaskInWorkerBeforeShutdown(t *testing.T) {
// see "re-queue the current task" message
assert.NoError(t, w.Shutdown())
}

func TestWithDisableConsumer(t *testing.T) {
job := queue.Job{
Payload: []byte("foo"),
}
w := NewWorker(
WithAddr(host+":4222"),
WithSubj("test02"),
WithQueue("test02"),
WithDisableConsumer(),
)

assert.NoError(t, w.Queue(job))
assert.NoError(t, w.Queue(job))
assert.NoError(t, w.Queue(job))
time.Sleep(100 * time.Millisecond)
assert.Equal(t, 0, len(w.tasks))
// see "re-queue the old job" message
assert.NoError(t, w.Shutdown())
}
18 changes: 5 additions & 13 deletions options.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,11 @@ import (
type Option func(*options)

type options struct {
runFunc func(context.Context, core.QueuedMessage) error
logger queue.Logger
addr string
subj string
queue string
disableConsumer bool
runFunc func(context.Context, core.QueuedMessage) error
logger queue.Logger
addr string
subj string
queue string
}

// WithAddr setup the addr of NATS
Expand Down Expand Up @@ -54,13 +53,6 @@ func WithLogger(l queue.Logger) Option {
}
}

// WithDisableConsumer disable consumer
func WithDisableConsumer() Option {
return func(w *options) {
w.disableConsumer = true
}
}

func newOptions(opts ...Option) options {
defaultOpts := options{
addr: "127.0.0.1:4222",
Expand Down

0 comments on commit 0cbfbd1

Please sign in to comment.