geo/server/internal/worker/generate/compliance_review_worker.go

package generate

import (
	"context"
	"errors"
	"fmt"
	"os"
	"time"

	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"

	"github.com/geo-platform/tenant-api/internal/shared/config"
	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
	tenantcompliance "github.com/geo-platform/tenant-api/internal/tenant/compliance"
)

var errComplianceReviewConsumerClosed = errors.New("compliance review consumer channel closed")

type ComplianceReviewWorker struct {
	rabbitMQ          *rabbitmq.Client
	compliance        *tenantcompliance.Service
	logger            *zap.Logger
	configProvider    config.Provider
	retryInterval     time.Duration
	consumerPrefix    string
	workerConcurrency int
}

func NewComplianceReviewWorker(
	rabbitMQClient *rabbitmq.Client,
	complianceService *tenantcompliance.Service,
	logger *zap.Logger,
	cfg config.ComplianceConfig,
) *ComplianceReviewWorker {
	config.NormalizeComplianceConfig(&cfg)
	return &ComplianceReviewWorker{
		rabbitMQ:          rabbitMQClient,
		compliance:        complianceService,
		logger:            logger,
		retryInterval:     5 * time.Second,
		consumerPrefix:    fmt.Sprintf("compliance-review-%d", os.Getpid()),
		workerConcurrency: cfg.ReviewWorkerConcurrency,
	}
}

func (w *ComplianceReviewWorker) WithConfigProvider(provider config.Provider) *ComplianceReviewWorker {
	if w != nil {
		w.configProvider = provider
	}
	return w
}

func (w *ComplianceReviewWorker) Start(ctx context.Context) {
	if w == nil || w.rabbitMQ == nil || w.compliance == nil {
		return
	}
	concurrency := w.runtimeConfig().ReviewWorkerConcurrency
	if concurrency <= 0 {
		concurrency = 1
	}
	for i := 0; i < concurrency; i++ {
		consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1)
		go w.run(ctx, consumerName)
	}
}

func (w *ComplianceReviewWorker) run(ctx context.Context, consumerName string) {
	for {
		if err := w.consumeOnce(ctx, consumerName); err != nil && ctx.Err() == nil && w.logger != nil {
			if errors.Is(err, errComplianceReviewConsumerClosed) {
				w.logger.Info("compliance review consumer channel closed, retrying",
					zap.String("consumer", consumerName),
				)
			} else {
				w.logger.Warn("compliance review worker stopped unexpectedly",
					zap.Error(err),
					zap.String("consumer", consumerName),
				)
			}
		}

		select {
		case <-ctx.Done():
			return
		case <-time.After(w.retryInterval):
		}
	}
}

func (w *ComplianceReviewWorker) consumeOnce(ctx context.Context, consumerName string) error {
	deliveries, ch, err := w.rabbitMQ.ConsumeComplianceReviewTask(consumerName)
	if err != nil {
		return err
	}
	defer ch.Close()

	for {
		select {
		case <-ctx.Done():
			return nil
		case delivery, ok := <-deliveries:
			if !ok {
				return errComplianceReviewConsumerClosed
			}
			w.handleDelivery(ctx, delivery)
		}
	}
}

func (w *ComplianceReviewWorker) handleDelivery(parent context.Context, delivery amqp.Delivery) {
	msg, err := tenantcompliance.DecodeReviewJobMessage(delivery.Body)
	if err != nil {
		w.rejectDelivery(delivery, false, err, 0)
		return
	}

	ctx, cancel := context.WithTimeout(parent, w.runtimeConfig().ReviewWorkerTimeout)
	defer cancel()

	err = w.compliance.ProcessReviewJob(ctx, msg)
	if err != nil {
		w.rejectDelivery(delivery, false, err, msg.JobID)
		return
	}
	if err := delivery.Ack(false); err != nil && w.logger != nil {
		w.logger.Warn("compliance review ack failed",
			zap.Error(err),
			zap.Int64("job_id", msg.JobID),
			zap.Int64("check_record_id", msg.CheckRecordID),
		)
	}
}

func (w *ComplianceReviewWorker) runtimeConfig() config.ComplianceConfig {
	if w != nil && w.configProvider != nil {
		if cfg := w.configProvider.Current(); cfg != nil {
			out := cfg.Compliance
			config.NormalizeComplianceConfig(&out)
			return out
		}
	}
	out := config.ComplianceConfig{Enabled: true}
	config.NormalizeComplianceConfig(&out)
	return out
}

func (w *ComplianceReviewWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, jobID int64) {
	if nackErr := delivery.Nack(false, requeue); nackErr != nil && w.logger != nil {
		w.logger.Warn("compliance review nack failed",
			zap.Error(nackErr),
			zap.Int64("job_id", jobID),
		)
	}
	if w.logger != nil {
		w.logger.Warn("compliance review processing failed",
			zap.Error(err),
			zap.Bool("requeue", requeue),
			zap.Bool("redelivered", delivery.Redelivered),
			zap.Int64("job_id", jobID),
		)
	}
}