From 24cf872b31403b4a7a2d9d022b3eb3936f2b0359 Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Wed, 11 Aug 2021 11:13:29 +0800
Subject: [PATCH] chore(logx): add custom logger in queue (#627)

---
 logx/log_interface.go | 45 +++++++++++++++++++++++++++++++++++++++++++
 main.go               |  4 ++++
 router/server_test.go |  2 ++
 3 files changed, 51 insertions(+)
 create mode 100644 logx/log_interface.go

diff --git a/logx/log_interface.go b/logx/log_interface.go
new file mode 100644
index 0000000..51a564c
--- /dev/null
+++ b/logx/log_interface.go
@@ -0,0 +1,45 @@
+package logx
+
+import (
+	"fmt"
+
+	"github.com/sirupsen/logrus"
+)
+
+// QueueLogger for simple logger.
+func QueueLogger() DefaultQueueLogger {
+	return DefaultQueueLogger{
+		accessLogger: LogAccess,
+		errorLogger:  LogError,
+	}
+}
+
+// DefaultQueueLogger for queue custom logger
+type DefaultQueueLogger struct {
+	accessLogger *logrus.Logger
+	errorLogger  *logrus.Logger
+}
+
+func (l DefaultQueueLogger) Infof(format string, args ...interface{}) {
+	l.accessLogger.Printf(format, args...)
+}
+
+func (l DefaultQueueLogger) Errorf(format string, args ...interface{}) {
+	l.errorLogger.Printf(format, args...)
+}
+
+func (l DefaultQueueLogger) Fatalf(format string, args ...interface{}) {
+	l.errorLogger.Fatalf(format, args...)
+}
+
+func (l DefaultQueueLogger) Info(args ...interface{}) {
+	l.accessLogger.Println(fmt.Sprint(args...))
+}
+
+func (l DefaultQueueLogger) Error(args ...interface{}) {
+	l.errorLogger.Println(fmt.Sprint(args...))
+}
+
+func (l DefaultQueueLogger) Fatal(args ...interface{}) {
+	l.errorLogger.Println(fmt.Sprint(args...))
+}
diff --git a/main.go b/main.go
index 376c6fd..e83c490 100644
--- a/main.go
+++ b/main.go
@@ -328,6 +328,7 @@ func main() {
 		w = simple.NewWorker(
 			simple.WithQueueNum(int(cfg.Core.QueueNum)),
 			simple.WithRunFunc(notify.Run(cfg)),
+			simple.WithLogger(logx.QueueLogger()),
 		)
 	case core.NSQ:
 		w = nsq.NewWorker(
@@ -336,6 +337,7 @@ func main() {
 			nsq.WithChannel(cfg.Queue.NSQ.Channel),
 			nsq.WithMaxInFlight(int(cfg.Core.WorkerNum)),
 			nsq.WithRunFunc(notify.Run(cfg)),
+			nsq.WithLogger(logx.QueueLogger()),
 		)
 	case core.NATS:
 		w = nats.NewWorker(
@@ -343,6 +345,7 @@ func main() {
 			nats.WithSubj(cfg.Queue.NATS.Subj),
 			nats.WithQueue(cfg.Queue.NATS.Queue),
 			nats.WithRunFunc(notify.Run(cfg)),
+			nats.WithLogger(logx.QueueLogger()),
 		)
 	default:
 		logx.LogError.Fatalf("we don't support queue engine: %s", cfg.Queue.Engine)
@@ -350,6 +353,7 @@ func main() {
 
 	q, err := queue.NewQueue(
 		queue.WithWorker(w),
+		queue.WithLogger(logx.QueueLogger()),
 		queue.WithWorkerCount(int(cfg.Core.WorkerNum)),
 	)
 	if err != nil {
diff --git a/router/server_test.go b/router/server_test.go
index d0c68f4..3e83882 100644
--- a/router/server_test.go
+++ b/router/server_test.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/appleboy/gorush/config"
 	"github.com/appleboy/gorush/core"
+	"github.com/appleboy/gorush/logx"
 	"github.com/appleboy/gorush/notify"
 	"github.com/appleboy/gorush/status"
 
@@ -52,6 +53,7 @@ func TestMain(m *testing.M) {
 	q, _ = queue.NewQueue(
 		queue.WithWorker(w),
 		queue.WithWorkerCount(4),
+		queue.WithLogger(logx.QueueLogger()),
 	)
 	q.Start()