chore: support NSQ worker. (#5)

Author: appleboy

.drone.yml

@@ -0,0 +1,37 @@
+---
+kind: pipeline
+name: testing
+
+platform:
+  os: linux
+  arch: amd64
+
+steps:
+- name: test
+  pull: always
+  image: golang:1.16
+  commands:
+  - go test -v -covermode=atomic -coverprofile=coverage.out ./...
+  volumes:
+  - name: gopath
+    path: /go
+
+- name: codecov
+  pull: always
+  image: robertstettner/drone-codecov
+  settings:
+    token:
+      from_secret: codecov_token
+
+services:
+- name: redis
+  image: redis
+
+- name: nsq
+  image: nsqio/nsq
+  commands:
+  - /nsqd
+
+volumes:
+- name: gopath
+  temp: {}

.github/workflows/go.yml

@@ -12,10 +12,33 @@ jobs:
   test:
     strategy:
       matrix:
-        os: [ubuntu-latest, macos-latest]
+        os: [ubuntu-latest]
         go: [1.13, 1.14, 1.15, 1.16]
     name: ${{ matrix.os }} @ Go ${{ matrix.go }}
     runs-on: ${{ matrix.os }}
+
+    # Service containers to run with `container-job`
+    services:
+      # Label used to access the service container
+      redis:
+        # Docker Hub image
+        image: redis
+        # Set health checks to wait until redis has started
+        options: >-
+          --health-cmd "redis-cli ping"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+        ports:
+          # Maps port 6379 on service container to the host
+          - 6379:6379
+
+      nsq:
+        image: nsqio/nsq
+        ports:
+          # Maps port 6379 on service container to the host
+          - 4150:4150
+
     env:
       GO111MODULE: on
       TESTTAGS: ${{ matrix.test-tags }}
@@ -36,7 +59,8 @@ jobs:
 
       - name: Run Tests
         run: |
-          go test -v -covermode=atomic -coverprofile=coverage.out
+          go test -v -covermode=atomic -coverprofile=coverage.out .
+          go test -v -covermode=atomic -coverprofile=coverage.out ./simple/...
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v1

go.mod

@@ -2,4 +2,7 @@ module github.com/appleboy/queue
 
 go 1.16
 
-require github.com/stretchr/testify v1.7.0
+require (
+	github.com/nsqio/go-nsq v1.0.8
+	github.com/stretchr/testify v1.7.0
+)

go.sum

@@ -1,5 +1,9 @@
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
+github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/nsqio/go-nsq v1.0.8 h1:3L2F8tNLlwXXlp2slDUrUWSBn2O3nMh8R1/KEDFTHPk=
+github.com/nsqio/go-nsq v1.0.8/go.mod h1:vKq36oyeVXgsS5Q8YEO7WghqidAVXQlcFxzQbQTuDEY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

nsq/README.md

@@ -0,0 +1,29 @@
+# NSQ
+
+A realtime distributed messaging platform
+
+## Setup
+
+start the NSQ lookupd
+
+```sh
+nsqlookupd
+```
+
+start the NSQ server
+
+```sh
+nsqd --lookupd-tcp-address=localhost:4160
+```
+
+start the NSQ admin dashboard
+
+```sh
+nsqadmin --lookupd-http-address localhost:4161
+```
+
+## Testing
+
+```sh
+go test -v ./...
+```

nsq/nsq.go

@@ -0,0 +1,168 @@
+package nsq
+
+import (
+	"runtime"
+	"sync"
+	"time"
+
+	"github.com/appleboy/queue"
+
+	"github.com/nsqio/go-nsq"
+)
+
+var _ queue.Worker = (*Worker)(nil)
+
+// Option for queue system
+type Option func(*Worker)
+
+// Worker for NSQ
+type Worker struct {
+	q           *nsq.Consumer
+	p           *nsq.Producer
+	startOnce   sync.Once
+	maxInFlight int
+	addr        string
+	topic       string
+	channel     string
+	runFunc     func(msg *nsq.Message) error
+}
+
+// WithAddr setup the addr of NSQ
+func WithAddr(addr string) Option {
+	return func(w *Worker) {
+		w.addr = addr
+	}
+}
+
+// WithTopic setup the topic of NSQ
+func WithTopic(topic string) Option {
+	return func(w *Worker) {
+		w.topic = topic
+	}
+}
+
+// WithChannel setup the channel of NSQ
+func WithChannel(channel string) Option {
+	return func(w *Worker) {
+		w.channel = channel
+	}
+}
+
+// WithRunFunc setup the run func of queue
+func WithRunFunc(fn func(msg *nsq.Message) error) Option {
+	return func(w *Worker) {
+		w.runFunc = fn
+	}
+}
+
+// WithMaxInFlight Maximum number of messages to allow in flight (concurrency knob)
+func WithMaxInFlight(num int) Option {
+	return func(w *Worker) {
+		w.maxInFlight = num
+	}
+}
+
+// NewWorker for struc
+func NewWorker(opts ...Option) *Worker {
+	w := &Worker{
+		addr:        "127.0.0.1:4150",
+		topic:       "gorush",
+		channel:     "ch",
+		maxInFlight: runtime.NumCPU(),
+		runFunc: 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
+			}
+
+			return nil
+		},
+	}
+
+	// Loop through each option
+	for _, opt := range opts {
+		// Call the option giving the instantiated
+		opt(w)
+	}
+
+	cfg := nsq.NewConfig()
+	cfg.MaxInFlight = w.maxInFlight
+	q, err := nsq.NewConsumer(w.topic, w.channel, cfg)
+	if err != nil {
+		panic(err)
+	}
+	w.q = q
+
+	p, err := nsq.NewProducer(w.addr, cfg)
+	if err != nil {
+		panic(err)
+	}
+	w.p = p
+
+	return w
+}
+
+// BeforeRun run script before start worker
+func (s *Worker) BeforeRun() error {
+	return nil
+}
+
+// AfterRun run script after start worker
+func (s *Worker) AfterRun() error {
+	s.startOnce.Do(func() {
+		time.Sleep(100 * time.Millisecond)
+		err := s.q.ConnectToNSQD(s.addr)
+		if err != nil {
+			panic("Could not connect nsq server: " + err.Error())
+		}
+	})
+
+	return nil
+}
+
+// Run start the worker
+func (s *Worker) Run(quit chan struct{}) error {
+	wg := &sync.WaitGroup{}
+	s.q.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
+		wg.Add(1)
+		defer wg.Done()
+		// run custom func
+		return s.runFunc(msg)
+	}))
+
+	// wait close signal
+	<-quit
+
+	// wait job completed
+	wg.Wait()
+
+	return nil
+}
+
+// Shutdown worker
+func (s *Worker) Shutdown() error {
+	s.q.Stop()
+	s.p.Stop()
+	return nil
+}
+
+// Capacity for channel
+func (s *Worker) Capacity() int {
+	return 0
+}
+
+// Usage for count of channel usage
+func (s *Worker) Usage() int {
+	return 0
+}
+
+// Queue send notification to queue
+func (s *Worker) Queue(job queue.QueuedMessage) error {
+	err := s.p.Publish(s.topic, job.Bytes())
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

nsq/nsq_test.go

@@ -0,0 +1,69 @@
+package nsq
+
+import (
+	"testing"
+	"time"
+
+	"github.com/appleboy/queue"
+
+	"github.com/nsqio/go-nsq"
+	"github.com/stretchr/testify/assert"
+)
+
+var host = "nsq"
+
+type mockMessage struct {
+	msg string
+}
+
+func (m mockMessage) Bytes() []byte {
+	return []byte(m.msg)
+}
+
+func TestShutdown(t *testing.T) {
+	w := NewWorker(
+		WithAddr(host+":4150"),
+		WithTopic("test"),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(1 * time.Second)
+	q.Shutdown()
+	// check shutdown once
+	q.Shutdown()
+	q.Wait()
+}
+
+func TestCustomFuncAndWait(t *testing.T) {
+	m := mockMessage{
+		msg: "foo",
+	}
+	w := NewWorker(
+		WithAddr(host+":4150"),
+		WithTopic("test"),
+		WithMaxInFlight(2),
+		WithRunFunc(func(msg *nsq.Message) error {
+			time.Sleep(500 * time.Millisecond)
+			return nil
+		}),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(100 * time.Millisecond)
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	time.Sleep(600 * time.Millisecond)
+	q.Shutdown()
+	q.Wait()
+	// you will see the execute time > 1000ms
+}