// Package logic is the deliver worker and queue contracts. package logic import ( "context" "sync" "time" "go.uber.org/zap" "spark-server/internal/infrastructure/logger" "spark-server/internal/module/deliver/entity" "spark-server/internal/transport/tcp/protocol" "spark-server/internal/transport/tcp/session" ) // Handler fans a job out. Return error to leave the job in processing for RequeueStale. type Handler func(ctx context.Context, job *entity.Job) error type Worker struct { log *zap.Logger queue Queue done Deduper handle Handler n int timeout time.Duration stop chan struct{} wg sync.WaitGroup } type WorkerOptions struct { Workers int TakeTimeout time.Duration Handler Handler } func NewWorker(log *zap.Logger, q Queue, done Deduper, o WorkerOptions) *Worker { if log == nil { log = zap.NewNop() } if o.Workers < 1 { o.Workers = 1 } if o.TakeTimeout <= 0 { o.TakeTimeout = 5 * time.Second } if done == nil { done = NopDeduper{} } if o.Handler == nil { o.Handler = NopHandler } return &Worker{ log: log, queue: q, done: done, handle: o.Handler, n: o.Workers, timeout: o.TakeTimeout, stop: make(chan struct{}), } } func NopHandler(context.Context, *entity.Job) error { return nil } type NopDeduper struct{} func (NopDeduper) Mark(context.Context, int64) (bool, error) { return true, nil } func SessionFanout(sessions session.Manager) Handler { return func(ctx context.Context, job *entity.Job) error { if sessions == nil || job == nil || len(job.ToUserIDs) == 0 { return nil } cmd := job.Cmd if cmd == 0 { cmd = protocol.CmdMessage } pkt := &session.Packet{Command: cmd, Body: job.Payload} for _, uid := range job.ToUserIDs { _ = sessions.SendToUserExceptDevice(ctx, uid, job.ExcludeDevice, pkt) } return nil } } func (w *Worker) Start(ctx context.Context) { if w == nil || w.queue == nil { return } if err := w.queue.RequeueStale(ctx); err != nil { w.log.Error("deliver requeue", zap.Error(err)) } for i := 0; i < w.n; i++ { w.wg.Add(1) go w.loop() } } func (w *Worker) Stop() { if w == nil { return } select { case <-w.stop: default: close(w.stop) } w.wg.Wait() } func (w *Worker) loop() { defer w.wg.Done() for { select { case <-w.stop: return default: } ctx, cancel := context.WithTimeout(context.Background(), w.timeout+time.Second) job, err := w.queue.Take(ctx, w.timeout) cancel() if err != nil { select { case <-w.stop: return case <-time.After(200 * time.Millisecond): } continue } if job == nil { select { case <-w.stop: return default: continue } } w.process(job) } } func (w *Worker) process(job *entity.Job) { defer func() { if rec := recover(); rec != nil { w.log.Error("deliver panic", zap.Any("panic", rec), zap.Int64("msg_id", job.MsgID)) } }() ctx := logger.Attach(context.Background(), w.log) first, err := w.done.Mark(ctx, job.MsgID) if err != nil { w.log.Error("deliver done", zap.Error(err), zap.Int64("msg_id", job.MsgID)) return } if !first { _ = w.queue.Ack(ctx, job) return } if err := w.handle(ctx, job); err != nil { w.log.Error("deliver handle", zap.Error(err), zap.Int64("msg_id", job.MsgID)) return } if err := w.queue.Ack(ctx, job); err != nil { w.log.Error("deliver ack", zap.Error(err), zap.Int64("msg_id", job.MsgID)) } }