nexarch/ai/internal/queue/worker.go

package queue

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

	"gitea.perlbach24.de/scripte/nexarch/ai/internal/limits"
	"gitea.perlbach24.de/scripte/nexarch/ai/internal/provider"
)

// UsageRecorder ist die schmale Schnittstelle zu GOV-04 (ai/internal/quota):
// zaehlt einen tatsaechlich ausgefuehrten Inferenz-Aufruf. Als Interface
// hier statt Import, damit queue (ENG-05) nicht von quota abhaengt —
// quota.Store erfuellt diese Schnittstelle strukturell.
type UsageRecorder interface {
	RecordCall(ctx context.Context, tenantSlug string) error
}

// Worker holt Jobs aus der Queue und fuehrt sie ueber den Provider aus.
// Ein Ausfuehrungsfehler blockiert den Anfrage-Pfad nicht (AK1) — der
// aufrufende Code hat laengst eine Job-ID zurueckbekommen und pollt oder
// wird benachrichtigt, dieser Worker laeuft unabhaengig im Hintergrund.
type Worker struct {
	queue    *Queue
	provider provider.Provider
	id       string
	limits   limits.Limits
	usage    UsageRecorder
}

func NewWorker(q *Queue, p provider.Provider, workerID string) *Worker {
	return &Worker{queue: q, provider: p, id: workerID, limits: limits.FromEnv()}
}

// SetLimits ueberschreibt die maximale Verarbeitungszeit je Anfrage (ENG-07),
// vor allem fuer Tests mit kleinen, deterministischen Werten.
func (w *Worker) SetLimits(l limits.Limits) { w.limits = l }

// SetUsageRecorder verknuepft den Worker mit GOV-04s Verbrauchszaehlung; ohne
// Aufruf bleibt die Ausfuehrung unveraendert ungezaehlt.
func (w *Worker) SetUsageRecorder(u UsageRecorder) { w.usage = u }

// ProcessOne holt einen Job und fuehrt ihn aus. Liefert ErrNoJobAvailable,
// wenn die Queue leer ist (kein Fehlerzustand).
func (w *Worker) ProcessOne(ctx context.Context) error {
	job, err := w.queue.Dequeue(ctx, w.id)
	if err != nil {
		return err
	}

	execCtx, cancel := w.limits.WithTimeout(ctx)
	defer cancel()
	result, execErr := w.execute(execCtx, job)
	if execErr != nil {
		if errors.Is(execErr, context.DeadlineExceeded) {
			execErr = fmt.Errorf("%w: job %s", limits.ErrProcessingTimeout, job.ID)
		}
		if failErr := w.queue.Fail(ctx, job.ID, execErr); failErr != nil {
			log.Printf("queue: fehlschlag fuer job %s konnte nicht erfasst werden: %v", job.ID, failErr)
		}
		return execErr
	}
	if completeErr := w.queue.Complete(ctx, job.ID, result); completeErr != nil {
		return fmt.Errorf("queue: job %s abschliessen: %w", job.ID, completeErr)
	}
	// Zaehlung erst NACH erfolgreichem Abschluss (GOV-04 Pruefung 3): der
	// Zaehlerstand bildet tatsaechlich ausgefuehrte Aufrufe ab, nicht
	// blosse Versuche — ein Fehlschlag der Zaehlung faellt bereits
	// erfolgreiche Inferenz nicht rueckwirkend zurueck, sondern wird nur
	// geloggt.
	if w.usage != nil && job.TenantSlug != "" {
		if usageErr := w.usage.RecordCall(ctx, job.TenantSlug); usageErr != nil {
			log.Printf("queue: verbrauchszaehlung fuer job %s (tenant %s) fehlgeschlagen: %v", job.ID, job.TenantSlug, usageErr)
		}
	}
	return nil
}

func (w *Worker) execute(ctx context.Context, job *Job) (any, error) {
	switch job.JobType {
	case KindChat:
		var req provider.ChatRequest
		if err := json.Unmarshal(job.Payload, &req); err != nil {
			return nil, fmt.Errorf("chat-payload dekodieren: %w", err)
		}
		return w.provider.Chat(ctx, req)
	case KindCompletion:
		var req provider.CompletionRequest
		if err := json.Unmarshal(job.Payload, &req); err != nil {
			return nil, fmt.Errorf("completion-payload dekodieren: %w", err)
		}
		return w.provider.Complete(ctx, req)
	case KindEmbedding:
		var req provider.EmbeddingRequest
		if err := json.Unmarshal(job.Payload, &req); err != nil {
			return nil, fmt.Errorf("embedding-payload dekodieren: %w", err)
		}
		return w.provider.Embed(ctx, req)
	default:
		return nil, fmt.Errorf("queue: unbekannter job-typ %q", job.JobType)
	}
}

// Run pollt die Queue in einem Intervall, bis ctx beendet wird. Damit laesst
// sich eine Lastspitze (Pruefung 1) mit mehreren nebenlaeufigen Run-Aufrufen
// (mehrere Goroutinen/Worker-IDs) ohne blockierten Anfrage-Pfad abarbeiten.
func (w *Worker) Run(ctx context.Context, pollInterval time.Duration) {
	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			for {
				err := w.ProcessOne(ctx)
				if errors.Is(err, ErrNoJobAvailable) {
					break
				}
				if err != nil {
					log.Printf("queue worker %s: %v", w.id, err)
				}
				if ctx.Err() != nil {
					return
				}
			}
		}
	}
}