fix: try to process all jobs before shutdown

pool.go

@@ -114,43 +114,37 @@ func (p *Pool) startWorker(w Worker) {
 	defer func() {
 		if r := recover(); r != nil {
 			if err, ok := r.(error); ok && err.Error() == "send on closed channel" {
-				w.Close()
 				return
 			}
+
 			log.Error("[BridgeWorker][addToQueue] recover from panic", "message", r, "trace", string(debug.Stack()))
 		}
 	}()
 	for {
 		// push worker chan into queue if worker has not closed yet
 		p.Queue <- w.Channel()
-		select {
-		case job := <-w.Channel():
-			if job == nil {
-				continue
-			}
-			if p.isClosed.Load().(bool) {
-				// update job to db
-				if err := job.Save(stores.STATUS_PENDING); err != nil {
-					log.Error("[Pool] failed on updating failed job", "err", err, "jobType", job.GetType())
-				}
+		job, more := <-w.Channel()
+		if !more {
+			w.Stop()
+			break
+		}
+
+		if job == nil {
+			continue
+		}
+
+		log.Debug("processing job", "id", job.GetID(), "nextTry", job.GetNextTry(), "retryCount", job.GetRetryCount(), "type", job.GetType())
+		if err := w.ProcessJob(job); err != nil {
+			// update try and next retry time
+			if job.GetRetryCount()+1 > job.GetMaxTry() {
+				log.Info("[Pool][processJob] job reaches its maxTry", "jobTransaction", job.GetTransaction().GetHash().Hex())
+				p.FailedJobChan <- job
 				continue
 			}
-			log.Debug("processing job", "id", job.GetID(), "nextTry", job.GetNextTry(), "retryCount", job.GetRetryCount(), "type", job.GetType())
-			if err := w.ProcessJob(job); err != nil {
-				// update try and next retry time
-				if job.GetRetryCount()+1 > job.GetMaxTry() {
-					log.Info("[Pool][processJob] job reaches its maxTry", "jobTransaction", job.GetTransaction().GetHash().Hex())
-					p.FailedJobChan <- job
-					continue
-				}
-				job.IncreaseRetryCount()
-				job.UpdateNextTry(time.Now().Unix() + int64(job.GetRetryCount()*job.GetBackOff()))
-				// send to retry job chan
-				p.RetryJob(job)
-			}
-		case <-w.Context().Done():
-			w.Close()
-			return
+			job.IncreaseRetryCount()
+			job.UpdateNextTry(time.Now().Unix() + int64(job.GetRetryCount()*job.GetBackOff()))
+			// send to retry job chan
+			p.RetryJob(job)
 		}
 	}
 }
@@ -222,45 +216,47 @@ func (p *Pool) Start(closeFunc func()) {
 			if job == nil {
 				continue
 			}
-			if p.isClosed.Load().(bool) {
-				if err := job.Save(stores.STATUS_PENDING); err != nil {
-					log.Error("[Pool] failed on saving processing job", "err", err, "jobType", job.GetType())
-				}
-				continue
-			}
+
 			log.Debug("[Pool] jobChan received a job", "jobId", job.GetID(), "nextTry", job.GetNextTry(), "type", job.GetType())
 			workerCh := <-p.Queue
 			p.SendJobToWorker(workerCh, job)
 		case <-p.ctx.Done():
-			log.Info("closing pool...")
+			log.Info("Closing pool...")
 			p.isClosed.Store(true)
 
 			// call close function firstly
 			if closeFunc != nil {
 				closeFunc()
 			}
 
-			// wait for all worker finish their close
-			for _, worker := range p.Workers {
-				worker.Wait()
-			}
-
 			// close all available channels to prevent further data send to pool's channels
 			close(p.JobChan)
-			close(p.FailedJobChan)
-			close(p.RetryJobChan)
-			close(p.Queue)
 
-			log.Info("Saving unprocessed jobs.", "jobs", len(p.JobChan))
+			log.Info("Trying to process all pending jobs.", "jobs", len(p.JobChan))
 			for {
 				job, more := <-p.JobChan
 				if !more {
 					break
 				}
-				// update job
-				p.saveJob(job)
+
+				if job == nil {
+					continue
+				}
+
+				workerCh := <-p.Queue
+				p.SendJobToWorker(workerCh, job)
+			}
+
+			// wait for all worker finish their close
+			for _, worker := range p.Workers {
+				worker.Close()
+				worker.Wait()
 			}
 
+			close(p.FailedJobChan)
+			close(p.RetryJobChan)
+			close(p.Queue)
+
 			log.Info("Saving unprocessed retry jobs.", "jobs", len(p.RetryJobChan))
 			for {
 				job, more := <-p.RetryJobChan
@@ -303,9 +299,11 @@ func (p *Pool) Stats() Stats {
 func (p *Pool) PrepareRetryableJob(job JobHandler) {
 	// if pool is closed, try update job to db
 	if p.isClosed.Load().(bool) {
-		log.Info("pool closed, update retrying job to database")
+		log.Debug("pool closed, update retrying job to database")
 		p.saveJob(job)
+		return
 	}
+
 	dur := time.Until(time.Unix(job.GetNextTry(), 0))
 	if dur <= 0 {
 		return

worker.go

@@ -3,16 +3,16 @@ package bridge_core
 import (
 	"context"
 	"fmt"
+
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/go-stack/stack"
-	"sync/atomic"
 )
 
 type Worker interface {
 	Context() context.Context
 	Close()
 	ProcessJob(job JobHandler) error
-	IsClose() bool
+	Stop()
 	Channel() chan JobHandler
 	Wait()
 }
@@ -23,14 +23,15 @@ type BridgeWorker struct {
 	// workerChan is used to receive and process job
 	workerChan chan JobHandler
 	listeners  map[string]Listener
-	isClose    int32
+	stopCh     chan struct{}
 }
 
 func NewWorker(ctx context.Context, id int, size int, listeners map[string]Listener) *BridgeWorker {
 	return &BridgeWorker{
 		ctx:        ctx,
 		id:         id,
 		workerChan: make(chan JobHandler, size),
+		stopCh:     make(chan struct{}),
 		listeners:  listeners,
 	}
 }
@@ -55,23 +56,18 @@ func (w *BridgeWorker) Context() context.Context {
 	return w.ctx
 }
 
-func (w *BridgeWorker) IsClose() bool {
-	return atomic.LoadInt32(&w.isClose) != 0
-}
-
 func (w *BridgeWorker) Channel() chan JobHandler {
 	return w.workerChan
 }
 
 func (w *BridgeWorker) Close() {
-	atomic.StoreInt32(&w.isClose, 1)
 	close(w.workerChan)
 }
 
+func (w *BridgeWorker) Stop() {
+	close(w.stopCh)
+}
+
 func (w *BridgeWorker) Wait() {
-	for {
-		if atomic.LoadInt32(&w.isClose) == 1 {
-			return
-		}
-	}
+	<-w.stopCh
 }