feat: 优化队列实现

pkg/queue/queue.go

@@ -1,53 +1,45 @@
 package queue
 
 import (
+	"context"
 	"errors"
 	"time"
 )
 
 type Queue struct {
-	jobs       chan Jobs
-	isShutdown chan struct{}
-	done       chan struct{}
+	jobs chan JobItem
 }
 
-func New() *Queue {
+func New(bufferSize int) *Queue {
 	return &Queue{
-		jobs:       make(chan Jobs, 10),
-		isShutdown: make(chan struct{}),
-		done:       make(chan struct{}),
+		jobs: make(chan JobItem, bufferSize),
 	}
 }
 
 func (r *Queue) Push(job Job, args []any) error {
 	select {
-	case <-r.isShutdown:
-		return errors.New("queue is shutdown, cannot add new jobs")
-	default:
-		r.jobs <- Jobs{Job: job, Args: args}
+	case r.jobs <- JobItem{Job: job, Args: args}:
 		return nil
+	default:
+		return errors.New("job queue is full")
 	}
 }
 
-func (r *Queue) Bulk(jobs []Jobs) error {
+func (r *Queue) Bulk(jobs []JobItem) error {
 	for _, job := range jobs {
-		if job.Delay > 0 {
-			time.AfterFunc(time.Duration(job.Delay)*time.Second, func() {
-				select {
-				case <-r.isShutdown:
-					return
-				default:
-					r.jobs <- Jobs{Job: job.Job, Args: job.Args}
-				}
+		jobCopy := job
+		if jobCopy.Delay > 0 {
+			time.AfterFunc(time.Duration(jobCopy.Delay)*time.Second, func() {
+				r.jobs <- jobCopy
 			})
 			continue
 		}
 
 		select {
-		case <-r.isShutdown:
-			return errors.New("queue is shutdown, cannot add new jobs")
+		case r.jobs <- jobCopy:
+			return nil
 		default:
-			r.jobs <- job
+			return errors.New("job queue is full")
 		}
 	}
 
@@ -55,38 +47,41 @@ func (r *Queue) Bulk(jobs []Jobs) error {
 }
 
 func (r *Queue) Later(delay uint, job Job, args []any) error {
+	jobCopy := job
+	argsCopy := make([]any, len(args))
+	copy(argsCopy, args)
 	time.AfterFunc(time.Duration(delay)*time.Second, func() {
-		select {
-		case <-r.isShutdown:
-			return
-		default:
-			r.jobs <- Jobs{Job: job, Args: args}
-		}
+		r.jobs <- JobItem{Job: jobCopy, Args: argsCopy}
 	})
 
 	return nil
 }
 
-func (r *Queue) Run() {
+func (r *Queue) Run(ctx context.Context) {
 	go func() {
 		for {
 			select {
 			case job := <-r.jobs:
-				if err := job.Job.Handle(job.Args...); err != nil {
-					if errJob, ok := job.Job.(JobWithErrHandle); ok {
-						errJob.ErrHandle(err)
-					}
-				}
-			case <-r.isShutdown:
-				close(r.done)
+				processJob(job)
+			case <-ctx.Done():
 				return
 			}
 		}
 	}()
 }
 
-func (r *Queue) Shutdown() error {
-	close(r.isShutdown)
-	<-r.done
-	return nil
+func (r *Queue) Len() int {
+	return len(r.jobs)
+}
+
+func (r *Queue) IsFull() bool {
+	return len(r.jobs) == cap(r.jobs)
+}
+
+func processJob(job JobItem) {
+	if err := job.Job.Handle(job.Args...); err != nil {
+		if errJob, ok := job.Job.(JobWithErrHandle); ok {
+			errJob.ErrHandle(err)
+		}
+	}
 }