chore(worker): remove disable consumer flag (#16)

Author: appleboy

nsq.go

@@ -19,10 +19,11 @@ var _ core.Worker = (*Worker)(nil)
 type Worker struct {
 	q         *nsq.Consumer
 	p         *nsq.Producer
+	cfg       *nsq.Config
 	stopOnce  sync.Once
+	startOnce sync.Once
 	stop      chan struct{}
 	stopFlag  int32
-	startFlag int32
 	opts      Options
 	tasks     chan *nsq.Message
 }
@@ -35,73 +36,64 @@ func NewWorker(opts ...Option) *Worker {
 		tasks: make(chan *nsq.Message),
 	}
 
-	cfg := nsq.NewConfig()
-	cfg.MaxInFlight = w.opts.maxInFlight
+	w.cfg = nsq.NewConfig()
+	w.cfg.MaxInFlight = w.opts.maxInFlight
 
-	if err := w.startProducer(cfg); err != nil {
-		panic(err)
-	}
-
-	if err := w.startConsumer(cfg); err != nil {
+	if err := w.startProducer(); err != nil {
 		panic(err)
 	}
 
 	return w
 }
 
-func (w *Worker) startProducer(cfg *nsq.Config) error {
+func (w *Worker) startProducer() error {
 	var err error
 
-	w.p, err = nsq.NewProducer(w.opts.addr, cfg)
+	w.p, err = nsq.NewProducer(w.opts.addr, w.cfg)
 
 	return err
 }
 
-func (w *Worker) startConsumer(cfg *nsq.Config) error {
-	if w.opts.disableConsumer {
-		return nil
-	}
-
-	var err error
-
-	w.q, err = nsq.NewConsumer(w.opts.topic, w.opts.channel, cfg)
-	if err != nil {
-		return err
-	}
-
-	w.q.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
-		if len(msg.Body) == 0 {
-			// Returning nil will automatically send a FIN command to NSQ to mark the message as processed.
-			// In this case, a message with an empty body is simply ignored/discarded.
-			return nil
+func (w *Worker) startConsumer() (err error) {
+	w.startOnce.Do(func() {
+		w.q, err = nsq.NewConsumer(w.opts.topic, w.opts.channel, w.cfg)
+		if err != nil {
+			return
 		}
 
-	loop:
-		for {
-			select {
-			case w.tasks <- msg:
-				break loop
-			case <-w.stop:
-				if msg != nil {
-					// re-queue the job if worker has been shutdown.
-					msg.Requeue(-1)
+		w.q.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
+			if len(msg.Body) == 0 {
+				// Returning nil will automatically send a FIN command to NSQ to mark the message as processed.
+				// In this case, a message with an empty body is simply ignored/discarded.
+				return nil
+			}
+
+		loop:
+			for {
+				select {
+				case w.tasks <- msg:
+					break loop
+				case <-w.stop:
+					if msg != nil {
+						// re-queue the job if worker has been shutdown.
+						msg.Requeue(-1)
+					}
+					break loop
+				case <-time.After(2 * time.Second):
+					msg.Touch()
 				}
-				break loop
-			case <-time.After(2 * time.Second):
-				msg.Touch()
 			}
-		}
 
-		return nil
-	}))
+			return nil
+		}))
 
-	err = w.q.ConnectToNSQD(w.opts.addr)
-	if err != nil {
-		return err
-	}
+		err = w.q.ConnectToNSQD(w.opts.addr)
+		if err != nil {
+			return
+		}
+	})
 
-	atomic.CompareAndSwapInt32(&w.startFlag, 0, 1)
-	return nil
+	return err
 }
 
 func (w *Worker) handle(job queue.Job) error {
@@ -172,12 +164,13 @@ func (w *Worker) Shutdown() error {
 		// notify shtdown event to worker and consumer
 		close(w.stop)
 		// stop producer and consumer
-		if atomic.LoadInt32(&w.startFlag) == 1 {
+		if w.q != nil {
 			w.q.ChangeMaxInFlight(0)
 			w.q.Stop()
 			<-w.q.StopChan
-			w.p.Stop()
 		}
+		w.p.Stop()
+
 		// close task channel
 		close(w.tasks)
 	})
@@ -195,6 +188,10 @@ func (w *Worker) Queue(job core.QueuedMessage) error {
 
 // Request fetch new task from queue
 func (w *Worker) Request() (core.QueuedMessage, error) {
+	if err := w.startConsumer(); err != nil {
+		return nil, err
+	}
+
 	clock := 0
 loop:
 	for {

nsq_test.go

@@ -277,7 +277,6 @@ func TestHandleTimeout(t *testing.T) {
 			time.Sleep(200 * time.Millisecond)
 			return nil
 		}),
-		WithDisableConsumer(),
 	)
 
 	err := w.handle(job)
@@ -294,7 +293,6 @@ func TestHandleTimeout(t *testing.T) {
 			time.Sleep(200 * time.Millisecond)
 			return nil
 		}),
-		WithDisableConsumer(),
 	)
 
 	done := make(chan error)
@@ -318,7 +316,6 @@ func TestJobComplete(t *testing.T) {
 		WithRunFunc(func(ctx context.Context, m core.QueuedMessage) error {
 			return errors.New("job completed")
 		}),
-		WithDisableConsumer(),
 	)
 
 	err := w.handle(job)
@@ -369,6 +366,7 @@ func TestNSQStatsinQueue(t *testing.T) {
 	assert.NoError(t, q.Queue(m))
 	assert.NoError(t, q.Queue(m))
 	q.Start()
+	time.Sleep(200 * time.Millisecond)
 	assert.Equal(t, int(1), w.Stats().Connections)
 	time.Sleep(500 * time.Millisecond)
 	assert.Equal(t, uint64(2), w.Stats().MessagesReceived)
@@ -390,14 +388,17 @@ func TestNSQStatsInWorker(t *testing.T) {
 	assert.NoError(t, w.Queue(m))
 	assert.NoError(t, w.Queue(m))
 	assert.NoError(t, w.Queue(m))
-	assert.Equal(t, int(1), w.Stats().Connections)
+	assert.Nil(t, w.Stats())
 
-	time.Sleep(50 * time.Millisecond)
+	task, err := w.Request()
+	assert.Equal(t, int(1), w.Stats().Connections)
+	assert.NotNil(t, task)
+	assert.NoError(t, err)
 
 	assert.Equal(t, uint64(1), w.Stats().MessagesReceived)
-	assert.Equal(t, uint64(0), w.Stats().MessagesFinished)
+	assert.Equal(t, uint64(1), w.Stats().MessagesFinished)
 	assert.Equal(t, uint64(0), w.Stats().MessagesRequeued)
-
+	time.Sleep(50 * time.Millisecond)
 	_ = w.Shutdown()
 	time.Sleep(50 * time.Millisecond)
 	assert.Equal(t, uint64(1), w.Stats().MessagesRequeued)

options.go

@@ -21,13 +21,12 @@ func (f OptionFunc) Apply(option *Options) {
 }
 
 type Options struct {
-	maxInFlight     int
-	addr            string
-	topic           string
-	channel         string
-	runFunc         func(context.Context, core.QueuedMessage) error
-	logger          queue.Logger
-	disableConsumer bool
+	maxInFlight int
+	addr        string
+	topic       string
+	channel     string
+	runFunc     func(context.Context, core.QueuedMessage) error
+	logger      queue.Logger
 }
 
 // WithAddr setup the addr of NSQ
@@ -72,13 +71,6 @@ func WithLogger(l queue.Logger) Option {
 	})
 }
 
-// WithDisableConsumer disable consumer
-func WithDisableConsumer() Option {
-	return OptionFunc(func(o *Options) {
-		o.disableConsumer = true
-	})
-}
-
 func newOptions(opts ...Option) Options {
 	defaultOpts := Options{
 		addr:        "127.0.0.1:4150",