IPSExtractor/services/services-count-facturas-medicas-serving/internal/worker/runner.go

package worker

import (
	"context"
	"log"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/nats-io/nats.go/jetstream"
	"github.com/pys/pdf-count/internal/config"
	"github.com/pys/pdf-count/internal/domain"
	"github.com/pys/pdf-count/internal/model"
	"github.com/pys/pdf-count/internal/ocr"
	"github.com/pys/pdf-count/internal/pipeline"
	"github.com/pys/pdf-count/internal/queue"
	"github.com/pys/pdf-count/internal/storage"
	"github.com/pys/pdf-count/internal/store"
)

type Runner struct {
	cfg       config.Config
	predictor *model.Predictor
	ocrEngine *ocr.Engine
	store     *store.Store
	storage   *storage.Client
	consumer  jetstream.Consumer
	ncClose   func()
}

func NewRunner(cfg config.Config) (*Runner, error) {
	predictor, err := model.NewPredictor(cfg.ModelPath, cfg.InferencePath)
	if err != nil {
		return nil, err
	}
	poolSize := cfg.OCRPoolSize
	if poolSize <= 0 {
		poolSize = cfg.WorkerConcurrency
	}
	ocrEngine, err := ocr.NewEngine(cfg.OCRLang, poolSize)
	if err != nil {
		predictor.Close()
		return nil, err
	}
	st, err := store.New(cfg.RedisURL, cfg.JobResultTTL)
	if err != nil {
		ocrEngine.Close()
		predictor.Close()
		return nil, err
	}
	stg, err := storage.New(cfg.MinIOEndpoint, cfg.MinIOAccessKey, cfg.MinIOSecretKey, cfg.MinIOBucket, cfg.MinIOUseSSL)
	if err != nil {
		ocrEngine.Close()
		st.Close()
		predictor.Close()
		return nil, err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	if err := stg.EnsureBucket(ctx); err != nil {
		ocrEngine.Close()
		st.Close()
		predictor.Close()
		return nil, err
	}

	_, cons, nc, err := queue.NewConsumer(cfg.NATSURL, cfg.StreamName, cfg.SubjectPrefix, "pdf-workers")
	if err != nil {
		ocrEngine.Close()
		st.Close()
		predictor.Close()
		return nil, err
	}

	return &Runner{
		cfg:       cfg,
		predictor: predictor,
		ocrEngine: ocrEngine,
		store:     st,
		storage:   stg,
		consumer:  cons,
		ncClose:   func() { nc.Close() },
	}, nil
}

func (r *Runner) Close() {
	r.ncClose()
	_ = r.ocrEngine.Close()
	_ = r.store.Close()
	_ = r.predictor.Close()
}

func (r *Runner) processSlots() int {
	c := r.cfg.WorkerConcurrency
	pool := r.cfg.OCRPoolSize
	if pool <= 0 {
		pool = r.cfg.WorkerConcurrency
	}
	if pool <= 0 {
		pool = 4
	}
	if c > pool {
		c = pool
	}
	if c <= 0 {
		c = 4
	}
	return c
}

func (r *Runner) Run(ctx context.Context) error {
	slots := r.processSlots()
	sem := make(chan struct{}, slots)
	var wg sync.WaitGroup

	go func() {
		<-ctx.Done()
		wg.Wait()
	}()

	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}

		msgs, err := r.consumer.Fetch(slots, jetstream.FetchMaxWait(r.cfg.NATSFetchMaxWait))
		if err != nil {
			continue
		}
		for msg := range msgs.Messages() {
			sem <- struct{}{}
			wg.Add(1)
			m := msg
			go func() {
				defer func() {
					<-sem
					wg.Done()
				}()
				r.handleMessage(ctx, m)
			}()
		}
		_ = msgs.Error()
	}
}

func (r *Runner) handleMessage(ctx context.Context, msg jetstream.Msg) {
	processCtx, cancel := context.WithTimeout(ctx, r.cfg.PDFProcessTimeout)
	defer cancel()

	qmsg, err := queue.DecodeMessage(msg.Data())
	if err != nil {
		_ = msg.Term()
		return
	}

	cancelled, err := r.store.IsCancelled(processCtx, qmsg.JobID)
	if err != nil || cancelled {
		_ = msg.Ack()
		return
	}

	data, err := r.storage.Get(processCtx, qmsg.ObjectKey)
	if err != nil {
		_ = r.store.SaveResult(processCtx, qmsg.JobID, domain.PDFResult{
			Filename: qmsg.Filename,
			Status:   domain.StatusError,
			Error:    err.Error(),
		})
		_ = r.store.TryCompleteJob(processCtx, qmsg.JobID)
		_ = msg.Ack()
		return
	}

	proc := pipeline.NewProcessor(r.predictor, r.ocrEngine, pipeline.Options{
		DPI:               r.cfg.OCRDPI,
		Lang:              r.cfg.OCRLang,
		DBSCANEps:         r.cfg.DBSCANEps,
		DBSCANMinSamples:  r.cfg.DBSCANMinSamples,
		IncludeRejected:   qmsg.IncludeRejected,
		LogPipelineTiming: r.cfg.LogPipelineTiming,
	})
	result := proc.ProcessPDF(qmsg.Filename, data)
	data = nil

	if err := r.store.SaveResult(processCtx, qmsg.JobID, result); err != nil {
		_ = msg.NakWithDelay(5 * time.Second)
		return
	}
	_ = r.store.TryCompleteJob(processCtx, qmsg.JobID)
	_ = msg.Ack()
}

func RunMain() {
	cfg := config.Load()
	runner, err := NewRunner(cfg)
	if err != nil {
		log.Fatalf("worker init: %v", err)
	}
	defer runner.Close()

	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer cancel()

	log.Printf("worker started slots=%d", runner.processSlots())
	if err := runner.Run(ctx); err != nil {
		log.Fatalf("worker: %v", err)
	}
	log.Println("worker stopped")
}