chore(queue): support job timeout setting. (#17)

Author: appleboy

logger.go

@@ -16,7 +16,8 @@ type Logger interface {
 	Fatal(args ...interface{})
 }
 
-func newLogger() Logger {
+// NewLogger for simple logger.
+func NewLogger() Logger {
 	return defaultLogger{
 		infoLogger:  log.New(os.Stderr, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile),
 		errorLogger: log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile),
@@ -53,3 +54,18 @@ func (l defaultLogger) Error(args ...interface{}) {
 func (l defaultLogger) Fatal(args ...interface{}) {
 	l.fatalLogger.Println(fmt.Sprint(args...))
 }
+
+// NewEmptyLogger for simple logger.
+func NewEmptyLogger() Logger {
+	return emptyLogger{}
+}
+
+// EmptyLogger no meesgae logger
+type emptyLogger struct{}
+
+func (l emptyLogger) Infof(format string, args ...interface{})  {}
+func (l emptyLogger) Errorf(format string, args ...interface{}) {}
+func (l emptyLogger) Fatalf(format string, args ...interface{}) {}
+func (l emptyLogger) Info(args ...interface{})                  {}
+func (l emptyLogger) Error(args ...interface{})                 {}
+func (l emptyLogger) Fatal(args ...interface{})                 {}

logger_test.go

@@ -0,0 +1,12 @@
+package queue
+
+func ExampleNewEmptyLogger() {
+	l := NewEmptyLogger()
+	l.Info("test")
+	l.Infof("test")
+	l.Error("test")
+	l.Errorf("test")
+	l.Fatal("test")
+	l.Fatalf("test")
+	// Output:
+}

queue.go

@@ -5,8 +5,12 @@ import (
 	"runtime"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
+// ErrQueueShutdown close the queue.
+var ErrQueueShutdown = errors.New("queue has been closed")
+
 type (
 	// A Queue is a message queue.
 	Queue struct {
@@ -17,9 +21,21 @@ type (
 		worker         Worker
 		stopOnce       sync.Once
 		runningWorkers int32
+		timeout        time.Duration
+	}
+
+	// Job with Timeout
+	Job struct {
+		Timeout time.Duration
+		Body    []byte
 	}
 )
 
+// Bytes get string body
+func (j Job) Bytes() []byte {
+	return j.Body
+}
+
 // Option for queue system
 type Option func(*Queue)
 
@@ -53,7 +69,8 @@ func NewQueue(opts ...Option) (*Queue, error) {
 		workerCount:  runtime.NumCPU(),
 		routineGroup: newRoutineGroup(),
 		quit:         make(chan struct{}),
-		logger:       newLogger(),
+		logger:       NewLogger(),
+		timeout:      24 * 60 * time.Minute,
 	}
 
 	// Loop through each option
@@ -106,7 +123,18 @@ func (q *Queue) Wait() {
 
 // Queue to queue all job
 func (q *Queue) Queue(job QueuedMessage) error {
-	return q.worker.Queue(job)
+	return q.worker.Queue(Job{
+		Timeout: q.timeout,
+		Body:    job.Bytes(),
+	})
+}
+
+// Queue to queue all job
+func (q *Queue) QueueWithTimeout(timeout time.Duration, job QueuedMessage) error {
+	return q.worker.Queue(Job{
+		Timeout: timeout,
+		Body:    job.Bytes(),
+	})
 }
 
 func (q *Queue) work() {
@@ -127,7 +155,7 @@ func (q *Queue) work() {
 		}()
 		q.logger.Infof("start the worker num: %d", num)
 		if err := q.worker.Run(); err != nil {
-			q.logger.Error(err)
+			q.logger.Errorf("runtime error: %s", err.Error())
 		}
 		q.logger.Infof("stop the worker num: %d", num)
 	})

queue_test.go

@@ -88,8 +88,8 @@ func TestWorkerStatus(t *testing.T) {
 
 	assert.NoError(t, q.Queue(m))
 	assert.NoError(t, q.Queue(m))
-	assert.NoError(t, q.Queue(m))
-	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.QueueWithTimeout(10*time.Millisecond, m))
+	assert.NoError(t, q.QueueWithTimeout(10*time.Millisecond, m))
 	assert.Equal(t, 100, q.Capacity())
 	assert.Equal(t, 4, q.Usage())
 	q.Start()

simple/simple.go

@@ -1,6 +1,7 @@
 package simple
 
 import (
+	"context"
 	"errors"
 	"sync"
 
@@ -18,10 +19,11 @@ var errMaxCapacity = errors.New("max capacity reached")
 
 // Worker for simple queue using channel
 type Worker struct {
-	queueNotification chan queue.QueuedMessage
-	runFunc           func(queue.QueuedMessage, <-chan struct{}) error
-	stop              chan struct{}
-	stopOnce          sync.Once
+	taskQueue chan queue.QueuedMessage
+	runFunc   func(context.Context, queue.QueuedMessage) error
+	stop      chan struct{}
+	logger    queue.Logger
+	stopOnce  sync.Once
 }
 
 // BeforeRun run script before start worker
@@ -36,36 +38,74 @@ func (s *Worker) AfterRun() error {
 
 // Run start the worker
 func (s *Worker) Run() error {
-	for notification := range s.queueNotification {
-		// run custom process function
-		_ = s.runFunc(notification, s.stop)
+	// check queue status
+	select {
+	case <-s.stop:
+		return queue.ErrQueueShutdown
+	default:
+	}
+
+	for task := range s.taskQueue {
+		done := make(chan struct{})
+		v, _ := task.(queue.Job)
+		ctx, cancel := context.WithTimeout(context.Background(), v.Timeout)
+		// vet doesn't complain if I do this
+		_ = cancel
+
+		// run the job
+		go func() {
+			// run custom process function
+			_ = s.runFunc(ctx, task)
+			close(done)
+		}()
+
+		select {
+		case <-ctx.Done(): // timeout reached
+			if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+				s.logger.Infof("job timeout: %s", v.Timeout.String())
+			}
+			// wait job
+			<-done
+		case <-s.stop: // shutdown service
+			cancel()
+			// wait job
+			<-done
+		case <-done: // job finish and continue to work
+		}
+
 	}
 	return nil
 }
 
 // Shutdown worker
 func (s *Worker) Shutdown() error {
 	s.stopOnce.Do(func() {
-		close(s.queueNotification)
 		close(s.stop)
+		close(s.taskQueue)
 	})
 	return nil
 }
 
 // Capacity for channel
 func (s *Worker) Capacity() int {
-	return cap(s.queueNotification)
+	return cap(s.taskQueue)
 }
 
 // Usage for count of channel usage
 func (s *Worker) Usage() int {
-	return len(s.queueNotification)
+	return len(s.taskQueue)
 }
 
 // Queue send notification to queue
 func (s *Worker) Queue(job queue.QueuedMessage) error {
 	select {
-	case s.queueNotification <- job:
+	case <-s.stop:
+		return queue.ErrQueueShutdown
+	default:
+	}
+
+	select {
+	case s.taskQueue <- job:
 		return nil
 	default:
 		return errMaxCapacity
@@ -75,23 +115,31 @@ func (s *Worker) Queue(job queue.QueuedMessage) error {
 // WithQueueNum setup the capcity of queue
 func WithQueueNum(num int) Option {
 	return func(w *Worker) {
-		w.queueNotification = make(chan queue.QueuedMessage, num)
+		w.taskQueue = make(chan queue.QueuedMessage, num)
 	}
 }
 
 // WithRunFunc setup the run func of queue
-func WithRunFunc(fn func(queue.QueuedMessage, <-chan struct{}) error) Option {
+func WithRunFunc(fn func(context.Context, queue.QueuedMessage) error) Option {
 	return func(w *Worker) {
 		w.runFunc = fn
 	}
 }
 
+// WithLogger set custom logger
+func WithLogger(l queue.Logger) Option {
+	return func(w *Worker) {
+		w.logger = l
+	}
+}
+
 // NewWorker for struc
 func NewWorker(opts ...Option) *Worker {
 	w := &Worker{
-		queueNotification: make(chan queue.QueuedMessage, defaultQueueSize),
-		stop:              make(chan struct{}),
-		runFunc: func(queue.QueuedMessage, <-chan struct{}) error {
+		taskQueue: make(chan queue.QueuedMessage, defaultQueueSize),
+		stop:      make(chan struct{}),
+		logger:    queue.NewLogger(),
+		runFunc: func(context.Context, queue.QueuedMessage) error {
 			return nil
 		},
 	}