Unblock processor shutdown process if processor is waiting for semaphore

token

processor.go

@@ -21,14 +21,19 @@ type processor struct {
 	// in case of a program shutdown or additon of a new queue.
 	dequeueTimeout time.Duration
 
+	// running represents the state of the "processor" goroutine
+	mu      sync.Mutex
+	running bool
+
 	// sema is a counting semaphore to ensure the number of active workers
 	// does not exceed the limit.
 	sema chan struct{}
 
 	// channel to communicate back to the long running "processor" goroutine.
-	// once is used to send value to the channel only once.
 	done chan struct{}
-	once sync.Once
+
+	// shutdown channel is closed when the shutdown of the "processor" goroutine starts.
+	shutdown chan struct{}
 
 	// quit channel communicates to the in-flight worker goroutines to stop.
 	quit chan struct{}
@@ -38,9 +43,10 @@ func newProcessor(r *rdb.RDB, numWorkers int, handler Handler) *processor {
 	return &processor{
 		rdb:            r,
 		handler:        handler,
-		dequeueTimeout: 5 * time.Second,
+		dequeueTimeout: 2 * time.Second,
 		sema:           make(chan struct{}, numWorkers),
 		done:           make(chan struct{}),
+		shutdown:       make(chan struct{}),
 		quit:           make(chan struct{}),
 	}
 }
@@ -48,12 +54,18 @@ func newProcessor(r *rdb.RDB, numWorkers int, handler Handler) *processor {
 // Note: stops only the "processor" goroutine, does not stop workers.
 // It's safe to call this method multiple times.
 func (p *processor) stop() {
-	p.once.Do(func() {
-		log.Println("[INFO] Processor shutting down...")
-		// Signal the processor goroutine to stop processing tasks
-		// from the queue.
-		p.done <- struct{}{}
-	})
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if !p.running {
+		return
+	}
+	log.Println("[INFO] Processor shutting down...")
+	// Unblock if processor is waiting for sema token.
+	close(p.shutdown)
+	// Signal the processor goroutine to stop processing tasks
+	// from the queue.
+	p.done <- struct{}{}
+	p.running = false
 }
 
 // NOTE: once terminated, processor cannot be re-started.
@@ -73,6 +85,11 @@ func (p *processor) terminate() {
 }
 
 func (p *processor) start() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.running {
+		return
+	}
 	// NOTE: The call to "restore" needs to complete before starting
 	// the processor goroutine.
 	p.restore()
@@ -87,6 +104,7 @@ func (p *processor) start() {
 			}
 		}
 	}()
+	p.running = true
 }
 
 // exec pulls a task out of the queue and starts a worker goroutine to
@@ -102,37 +120,43 @@ func (p *processor) exec() {
 		return
 	}
 
-	p.sema <- struct{}{} // acquire token
-	go func() {
-		defer func() { <-p.sema /* release token */ }()
-
-		resCh := make(chan error, 1)
-		task := &Task{Type: msg.Type, Payload: msg.Payload}
+	select {
+	case <-p.shutdown:
+		// shutdown is starting, return immediately after requeuing the message.
+		p.requeue(msg)
+		return
+	case p.sema <- struct{}{}: // acquire token
 		go func() {
-			resCh <- perform(p.handler, task)
-		}()
+			defer func() { <-p.sema /* release token */ }()
 
-		select {
-		case <-p.quit:
-			// time is up, quit this worker goroutine.
-			return
-		case resErr := <-resCh:
-			// Note: One of three things should happen.
-			// 1) Done  -> Removes the message from InProgress
-			// 2) Retry -> Removes the message from InProgress & Adds the message to Retry
-			// 3) Kill  -> Removes the message from InProgress & Adds the message to Dead
-			if resErr != nil {
-				if msg.Retried >= msg.Retry {
-					p.kill(msg, resErr.Error())
+			resCh := make(chan error, 1)
+			task := &Task{Type: msg.Type, Payload: msg.Payload}
+			go func() {
+				resCh <- perform(p.handler, task)
+			}()
+
+			select {
+			case <-p.quit:
+				// time is up, quit this worker goroutine.
+				return
+			case resErr := <-resCh:
+				// Note: One of three things should happen.
+				// 1) Done  -> Removes the message from InProgress
+				// 2) Retry -> Removes the message from InProgress & Adds the message to Retry
+				// 3) Kill  -> Removes the message from InProgress & Adds the message to Dead
+				if resErr != nil {
+					if msg.Retried >= msg.Retry {
+						p.kill(msg, resErr.Error())
+						return
+					}
+					p.retry(msg, resErr.Error())
 					return
 				}
-				p.retry(msg, resErr.Error())
+				p.markAsDone(msg)
 				return
 			}
-			p.markAsDone(msg)
-			return
-		}
-	}()
+		}()
+	}
 }
 
 // restore moves all tasks from "in-progress" back to queue
@@ -144,6 +168,13 @@ func (p *processor) restore() {
 	}
 }
 
+func (p *processor) requeue(msg *rdb.TaskMessage) {
+	err := p.rdb.Requeue(msg)
+	if err != nil {
+		log.Printf("[ERROR] Could not move task from InProgress back to queue: %v\n", err)
+	}
+}
+
 func (p *processor) markAsDone(msg *rdb.TaskMessage) {
 	err := p.rdb.Done(msg)
 	if err != nil {