gochat/backend/internal/autoassignment/assignment_job.go

package autoassignment

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/worker"
	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
)

const (
	TaskTypeAssignmentJob = "auto_assignment:assign_inbox"
	assignmentInFlightTTL = 5 * time.Minute
)

var errAssignmentJobUnavailable = errors.New("auto-assignment job requires worker and redis")

type assignmentJobPayload struct {
	InboxID   uint   `json:"inbox_id"`
	AccountID uint   `json:"account_id"`
	Token     string `json:"token"`
}

// AssignmentJob durably processes one bounded inbox backlog batch.
type AssignmentJob struct {
	service *AssignmentService
	worker  *worker.WorkerPool
	redis   *redis.Client
}

func NewAssignmentJob(service *AssignmentService, workerPool *worker.WorkerPool, rdb *redis.Client) *AssignmentJob {
	job := &AssignmentJob{service: service, worker: workerPool, redis: rdb}
	if workerPool != nil {
		workerPool.Register(TaskTypeAssignmentJob, job.perform)
		workerPool.RegisterFailureHandler(TaskTypeAssignmentJob, job.afterFailure)
	}
	return job
}

func (j *AssignmentJob) EnqueueForInbox(ctx context.Context, inboxID, accountID uint) (bool, error) {
	if j.worker == nil || j.redis == nil {
		return false, errAssignmentJobUnavailable
	}

	payload := assignmentJobPayload{InboxID: inboxID, AccountID: accountID, Token: uuid.NewString()}
	claimed, err := j.redis.SetNX(ctx, j.lockKey(inboxID), payload.Token, assignmentInFlightTTL).Result()
	if err != nil || !claimed {
		return false, err
	}

	if _, err := j.worker.Enqueue(ctx, TaskTypeAssignmentJob, payload, worker.WithMaxAttempts(3)); err != nil {
		_ = j.release(ctx, inboxID, payload.Token)
		return false, err
	}
	return true, nil
}

func (j *AssignmentJob) perform(ctx context.Context, backgroundJob *model.BackgroundJob) (err error) {
	var payload assignmentJobPayload
	if err := json.Unmarshal(backgroundJob.Payload, &payload); err != nil {
		return worker.Permanent(fmt.Errorf("unmarshal auto-assignment job: %w", err))
	}
	keepToken := false
	defer func() {
		if keepToken || err != nil {
			return
		}
		_ = j.release(context.WithoutCancel(ctx), payload.InboxID, payload.Token)
	}()

	assigned, err := j.service.AssignUnassignedConversations(ctx, payload.InboxID, payload.AccountID)
	if err != nil || len(assigned) < assignmentBatchLimit {
		return err
	}
	_, err = j.worker.Enqueue(ctx, TaskTypeAssignmentJob, payload, worker.WithMaxAttempts(3))
	keepToken = err == nil
	return err
}

func (j *AssignmentJob) afterFailure(ctx context.Context, backgroundJob *model.BackgroundJob, retryAfter time.Duration) error {
	var payload assignmentJobPayload
	if err := json.Unmarshal(backgroundJob.Payload, &payload); err != nil {
		return fmt.Errorf("unmarshal auto-assignment failure payload: %w", err)
	}
	if backgroundJob.Status == model.BackgroundJobStatusDead {
		return j.release(ctx, payload.InboxID, payload.Token)
	}
	if retryAfter < 0 {
		retryAfter = 0
	}
	return j.renew(ctx, payload.InboxID, payload.Token, assignmentInFlightTTL+retryAfter)
}

func (j *AssignmentJob) renew(ctx context.Context, inboxID uint, token string, ttl time.Duration) error {
	if j.redis == nil || token == "" {
		return nil
	}
	return j.redis.Eval(ctx,
		`if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("PEXPIRE", KEYS[1], ARGV[2]) else return 0 end`,
		[]string{j.lockKey(inboxID)}, token, ttl.Milliseconds(),
	).Err()
}

func (j *AssignmentJob) release(ctx context.Context, inboxID uint, token string) error {
	if j.redis == nil || token == "" {
		return nil
	}
	return j.redis.Eval(ctx,
		`if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("DEL", KEYS[1]) else return 0 end`,
		[]string{j.lockKey(inboxID)}, token,
	).Err()
}

func (j *AssignmentJob) lockKey(inboxID uint) string {
	return fmt.Sprintf("gochat:auto_assignment:in_flight:%d", inboxID)
}