refactor(app): 优化后台任务退出与外部 IMAP 认证

- 统一用 `WaitGroup` 管理后台 worker,关闭时等待任务结束后再释放数据库连接。
- 将外部 IMAP 的 OAuth 认证切换为 `XOAUTH2` 客户端实现,并补充对应格式测试。
- 调整发送队列的故障恢复逻辑,超时中断后恢复为排队状态以便重新投递。
This commit is contained in:
LanQin
2026-06-25 21:36:59 +08:00
parent 081cfad02d
commit efcdc691de
4 changed files with 66 additions and 9 deletions
+16 -5
View File
@@ -15,6 +15,7 @@ import (
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
"sync"
"time" "time"
"golang.org/x/crypto/bcrypt" "golang.org/x/crypto/bcrypt"
@@ -28,6 +29,7 @@ type App struct {
now func() time.Time now func() time.Time
policy *HTMLPolicy policy *HTMLPolicy
workerCancel context.CancelFunc workerCancel context.CancelFunc
workerWG sync.WaitGroup
maildirHealth *maildirSyncHealthTracker maildirHealth *maildirSyncHealthTracker
externalIMAP externalIMAPClientFactory externalIMAP externalIMAPClientFactory
} }
@@ -73,16 +75,24 @@ func New(cfg Config, logger *slog.Logger) (*App, error) {
} }
workerCtx, cancel := context.WithCancel(context.Background()) workerCtx, cancel := context.WithCancel(context.Background())
a.workerCancel = cancel a.workerCancel = cancel
go a.scheduledSendWorker(workerCtx) a.startWorker(func() { a.scheduledSendWorker(workerCtx) })
if strings.TrimSpace(a.cfg.MaildirRoot) != "" { if strings.TrimSpace(a.cfg.MaildirRoot) != "" {
go a.maildirWorker(workerCtx) a.startWorker(func() { a.maildirWorker(workerCtx) })
} }
go a.sendQueueWorker(workerCtx) a.startWorker(func() { a.sendQueueWorker(workerCtx) })
go a.externalIMAPWorker(workerCtx) a.startWorker(func() { a.externalIMAPWorker(workerCtx) })
go a.smtpEventsCleanupWorker(workerCtx) a.startWorker(func() { a.smtpEventsCleanupWorker(workerCtx) })
return a, nil return a, nil
} }
func (a *App) startWorker(fn func()) {
a.workerWG.Add(1)
go func() {
defer a.workerWG.Done()
fn()
}()
}
func (a *App) Close() error { func (a *App) Close() error {
if a == nil || a.db == nil { if a == nil || a.db == nil {
return nil return nil
@@ -90,6 +100,7 @@ func (a *App) Close() error {
if a.workerCancel != nil { if a.workerCancel != nil {
a.workerCancel() a.workerCancel()
} }
a.workerWG.Wait()
return a.db.Close() return a.db.Close()
} }
+24
View File
@@ -62,6 +62,13 @@ func newTestAppWithConfig(t *testing.T, cfg Config) *App {
return a return a
} }
func stopTestWorkers(a *App) {
if a != nil && a.workerCancel != nil {
a.workerCancel()
a.workerWG.Wait()
}
}
func defaultAdminUserAndMailbox(t *testing.T, a *App) (*User, *Mailbox) { func defaultAdminUserAndMailbox(t *testing.T, a *App) (*User, *Mailbox) {
t.Helper() t.Helper()
ctx := context.Background() ctx := context.Background()
@@ -581,6 +588,21 @@ func TestExternalIMAPOAuthEmailFromIDToken(t *testing.T) {
} }
} }
func TestExternalIMAPXOAUTH2ClientFormat(t *testing.T) {
client := newExternalIMAPXOAUTH2Client("user@example.com", "access-token")
mech, initialResponse, err := client.Start()
if err != nil {
t.Fatal(err)
}
if mech != "XOAUTH2" {
t.Fatalf("mechanism=%q, want XOAUTH2", mech)
}
want := "user=user@example.com\x01auth=Bearer access-token\x01\x01"
if string(initialResponse) != want {
t.Fatalf("initial response=%q, want %q", string(initialResponse), want)
}
}
func mustOAuthStateFromURL(t *testing.T, rawURL string) string { func mustOAuthStateFromURL(t *testing.T, rawURL string) string {
t.Helper() t.Helper()
u, err := url.Parse(rawURL) u, err := url.Parse(rawURL)
@@ -1588,6 +1610,7 @@ func TestMailSendRollsBackSentCopyWhenQueueInsertFails(t *testing.T) {
func TestSendQueueRecoversStaleSendingItems(t *testing.T) { func TestSendQueueRecoversStaleSendingItems(t *testing.T) {
a := newTestApp(t) a := newTestApp(t)
stopTestWorkers(a)
host, port, received := startCapturingSMTP(t, 1) host, port, received := startCapturingSMTP(t, 1)
a.cfg.SMTPHost = host a.cfg.SMTPHost = host
a.cfg.SMTPPort = port a.cfg.SMTPPort = port
@@ -1637,6 +1660,7 @@ func TestSendQueueRecoversStaleSendingItems(t *testing.T) {
func TestSendQueueStaleDeliveredMarkerDoesNotRedeliver(t *testing.T) { func TestSendQueueStaleDeliveredMarkerDoesNotRedeliver(t *testing.T) {
a := newTestApp(t) a := newTestApp(t)
stopTestWorkers(a)
host, port, received := startCapturingSMTP(t, 1) host, port, received := startCapturingSMTP(t, 1)
a.cfg.SMTPHost = host a.cfg.SMTPHost = host
a.cfg.SMTPPort = port a.cfg.SMTPPort = port
+18 -2
View File
@@ -24,7 +24,6 @@ import (
"github.com/emersion/go-imap/v2" "github.com/emersion/go-imap/v2"
"github.com/emersion/go-imap/v2/imapclient" "github.com/emersion/go-imap/v2/imapclient"
"github.com/emersion/go-sasl"
"github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5"
"golang.org/x/oauth2" "golang.org/x/oauth2"
) )
@@ -1498,7 +1497,7 @@ func (a *App) openExternalIMAPClient(ctx context.Context, account externalIMAPAc
c.Close() c.Close()
return nil, err return nil, err
} }
if err := c.Authenticate(sasl.NewOAuthBearerClient(&sasl.OAuthBearerOptions{Username: account.Username, Token: token, Host: account.Host, Port: account.Port})); err != nil { if err := c.Authenticate(newExternalIMAPXOAUTH2Client(account.Username, token)); err != nil {
c.Close() c.Close()
return nil, err return nil, err
} }
@@ -1516,6 +1515,23 @@ func (a *App) openExternalIMAPClient(ctx context.Context, account externalIMAPAc
return &goExternalIMAPClient{client: c}, nil return &goExternalIMAPClient{client: c}, nil
} }
type externalIMAPXOAUTH2Client struct {
username string
token string
}
func newExternalIMAPXOAUTH2Client(username, token string) externalIMAPXOAUTH2Client {
return externalIMAPXOAUTH2Client{username: username, token: token}
}
func (c externalIMAPXOAUTH2Client) Start() (string, []byte, error) {
return "XOAUTH2", []byte("user=" + c.username + "\x01auth=Bearer " + c.token + "\x01\x01"), nil
}
func (c externalIMAPXOAUTH2Client) Next(challenge []byte) ([]byte, error) {
return []byte{}, nil
}
func (a *App) externalIMAPOAuthAccessToken(ctx context.Context, account externalIMAPAccountRecord) (string, error) { func (a *App) externalIMAPOAuthAccessToken(ctx context.Context, account externalIMAPAccountRecord) (string, error) {
access, err := a.decryptExternalIMAPPassword(account.OAuthAccessTokenCiphertext) access, err := a.decryptExternalIMAPPassword(account.OAuthAccessTokenCiphertext)
if err != nil { if err != nil {
+8 -2
View File
@@ -126,6 +126,12 @@ func (a *App) sendQueueWorker(ctx context.Context) {
ticker := time.NewTicker(10 * time.Second) ticker := time.NewTicker(10 * time.Second)
defer ticker.Stop() defer ticker.Stop()
for { for {
select {
case <-ctx.Done():
a.log.Info("send queue worker stopped")
return
default:
}
if err := a.processDueSendQueue(ctx); err != nil { if err := a.processDueSendQueue(ctx); err != nil {
a.log.Warn("send queue worker failed", "error", err) a.log.Warn("send queue worker failed", "error", err)
} }
@@ -232,12 +238,12 @@ func (a *App) recoverStaleSendQueueItems(ctx context.Context) error {
} }
continue continue
} }
res, err := a.db.ExecContext(ctx, `UPDATE send_queue SET status=?,next_attempt_at=?,last_error=?,updated_at=? WHERE id=? AND status=?`, sendQueueStatusFailed, now, "send attempt interrupted", now, item.ID, sendQueueStatusSending) res, err := a.db.ExecContext(ctx, `UPDATE send_queue SET status=?,next_attempt_at=?,last_error=?,updated_at=? WHERE id=? AND status=?`, sendQueueStatusQueued, now, "send attempt interrupted", now, item.ID, sendQueueStatusSending)
if err != nil { if err != nil {
return err return err
} }
if n, _ := res.RowsAffected(); n > 0 { if n, _ := res.RowsAffected(); n > 0 {
a.recordSendAudit(ctx, sendAuditRetry, sendQueueStatusFailed, sendAuditInputFromQueue(item, "send attempt interrupted")) a.recordSendAudit(ctx, sendAuditRetry, sendQueueStatusQueued, sendAuditInputFromQueue(item, "send attempt interrupted"))
} }
} }
return nil return nil