email_queue_service.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. package service
  2. import (
  3. "context"
  4. "fmt"
  5. "sync"
  6. "time"
  7. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  8. )
  9. // Task type constants
  10. const (
  11. TaskTypeVerifyCode = "verify_code"
  12. TaskTypePasswordReset = "password_reset"
  13. )
  14. // EmailTask 邮件发送任务
  15. type EmailTask struct {
  16. Email string
  17. SiteName string
  18. TaskType string // "verify_code" or "password_reset"
  19. ResetURL string // Only used for password_reset task type
  20. Locale string // Optional Accept-Language locale hint
  21. }
  22. // EmailQueueService 异步邮件队列服务
  23. type EmailQueueService struct {
  24. emailService *EmailService
  25. taskChan chan EmailTask
  26. wg sync.WaitGroup
  27. stopChan chan struct{}
  28. workers int
  29. }
  30. // NewEmailQueueService 创建邮件队列服务
  31. func NewEmailQueueService(emailService *EmailService, workers int) *EmailQueueService {
  32. if workers <= 0 {
  33. workers = 3 // 默认3个工作协程
  34. }
  35. service := &EmailQueueService{
  36. emailService: emailService,
  37. taskChan: make(chan EmailTask, 100), // 缓冲100个任务
  38. stopChan: make(chan struct{}),
  39. workers: workers,
  40. }
  41. // 启动工作协程
  42. service.start()
  43. return service
  44. }
  45. // start 启动工作协程
  46. func (s *EmailQueueService) start() {
  47. for i := 0; i < s.workers; i++ {
  48. s.wg.Add(1)
  49. go s.worker(i)
  50. }
  51. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Started %d workers", s.workers)
  52. }
  53. // worker 工作协程
  54. func (s *EmailQueueService) worker(id int) {
  55. defer s.wg.Done()
  56. for {
  57. select {
  58. case task := <-s.taskChan:
  59. s.processTask(id, task)
  60. case <-s.stopChan:
  61. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Worker %d stopping", id)
  62. return
  63. }
  64. }
  65. }
  66. // processTask 处理任务
  67. func (s *EmailQueueService) processTask(workerID int, task EmailTask) {
  68. ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
  69. defer cancel()
  70. switch task.TaskType {
  71. case TaskTypeVerifyCode:
  72. if err := s.emailService.SendVerifyCode(ctx, task.Email, task.SiteName, task.Locale); err != nil {
  73. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Worker %d failed to send verify code to %s: %v", workerID, task.Email, err)
  74. } else {
  75. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Worker %d sent verify code to %s", workerID, task.Email)
  76. }
  77. case TaskTypePasswordReset:
  78. if err := s.emailService.SendPasswordResetEmailWithCooldown(ctx, task.Email, task.SiteName, task.ResetURL, task.Locale); err != nil {
  79. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Worker %d failed to send password reset to %s: %v", workerID, task.Email, err)
  80. } else {
  81. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Worker %d sent password reset to %s", workerID, task.Email)
  82. }
  83. default:
  84. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Worker %d unknown task type: %s", workerID, task.TaskType)
  85. }
  86. }
  87. // EnqueueVerifyCode 将验证码发送任务加入队列
  88. func (s *EmailQueueService) EnqueueVerifyCode(email, siteName string, locale ...string) error {
  89. task := EmailTask{
  90. Email: email,
  91. SiteName: siteName,
  92. TaskType: TaskTypeVerifyCode,
  93. Locale: firstEmailLocale(locale),
  94. }
  95. select {
  96. case s.taskChan <- task:
  97. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Enqueued verify code task for %s", email)
  98. return nil
  99. default:
  100. return fmt.Errorf("email queue is full")
  101. }
  102. }
  103. // EnqueuePasswordReset 将密码重置邮件任务加入队列
  104. func (s *EmailQueueService) EnqueuePasswordReset(email, siteName, resetURL string, locale ...string) error {
  105. task := EmailTask{
  106. Email: email,
  107. SiteName: siteName,
  108. TaskType: TaskTypePasswordReset,
  109. ResetURL: resetURL,
  110. Locale: firstEmailLocale(locale),
  111. }
  112. select {
  113. case s.taskChan <- task:
  114. logger.LegacyPrintf("service.email_queue", "[EmailQueue] Enqueued password reset task for %s", email)
  115. return nil
  116. default:
  117. return fmt.Errorf("email queue is full")
  118. }
  119. }
  120. // Stop 停止队列服务
  121. func (s *EmailQueueService) Stop() {
  122. close(s.stopChan)
  123. s.wg.Wait()
  124. logger.LegacyPrintf("service.email_queue", "%s", "[EmailQueue] All workers stopped")
  125. }