spark-server/internal/module/deliver/logic/worker.go

// 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))
	}
}