Aether/backend/internal/documents/worker.go

package documents

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

	"github.com/rs/zerolog"

	"aether/internal/storage"
	"aether/pkg/bg"
)

// idleInterval ist die Pause, wenn nichts zu tun ist.
//
// Abfragen statt Benachrichtigen (LISTEN/NOTIFY) ist hier die richtige Wahl:
// Eine Abfrage auf einen Teilindex kostet nichts, und der Arbeiter muss ohnehin
// regelmäßig nachsehen – ein Auftrag kann auch dadurch entstehen, dass der
// Server beim letzten Mal mittendrin ausgefallen ist. Fünf Sekunden Verzug
// spielen bei einer Aufgabe, die selbst Sekunden bis Minuten dauert, keine
// Rolle.
const idleInterval = 5 * time.Second

// busyInterval: Wurde gerade etwas erledigt, sofort weitermachen. Eine kurze
// Pause bleibt trotzdem, damit ein Stapel von hundert Scans nicht die ganze
// Maschine belegt und die Cloud daneben zäh wird.
const busyInterval = 200 * time.Millisecond

// maxFileBytes: Größeres wird übersprungen. Der Arbeiter muss die Datei in eine
// temporäre Kopie schreiben, bevor die Werkzeuge sie lesen können; ein
// 4-GB-Video würde dabei die Platte füllen – und Text ist darin ohnehin keiner.
const maxFileBytes = 512 << 20 // 512 MiB

// Worker arbeitet die Warteschlange ab. Einer genügt: OCR lastet einen Kern
// voll aus, und mehrere parallel machen auf einem Heimserver die Oberfläche
// zäh, ohne insgesamt schneller fertig zu sein.
type Worker struct {
	store   *Store
	blobs   storage.Storage
	extract *Extractor
	log     zerolog.Logger
	// paused merkt sich, ob wegen Platzmangels gerade nichts läuft – nur damit
	// die Warnung einmal im Log steht und nicht alle fünf Sekunden.
	paused bool
}

func NewWorker(store *Store, blobs storage.Storage, ex *Extractor, log zerolog.Logger) *Worker {
	return &Worker{store: store, blobs: blobs, extract: ex, log: log.With().Str("component", "ocr").Logger()}
}

// Run läuft, bis der Kontext endet. Gedacht für `go worker.Run(ctx)` beim Start.
func (w *Worker) Run(ctx context.Context) {
	w.log.Info().
		Str("languages", w.extract.Languages).
		Int("max_pages", w.extract.MaxPages).
		Msg("OCR-Arbeiter gestartet")

	// Angefangene Dateien eines hart beendeten Laufs wegräumen (siehe
	// sweepTemp). Einmal beim Start genügt: Während des Betriebs räumt der
	// Arbeiter selbst auf.
	if n, freed := sweepTemp(w.extract.TmpDir); n > 0 {
		w.log.Info().Int("files", n).Int64("bytes", freed).
			Msg("Zwischendateien eines früheren Laufs entfernt")
	}

	for {
		wait := idleInterval
		// **Je Auftrag abgesichert, nicht je Arbeiter.** Fängt man eine Panik
		// um die ganze Schleife ab, überlebt zwar der Prozess – aber der
		// Arbeiter ist tot, und von da an wird still keine Datei mehr erkannt.
		// So kostet ein unverdauliches Dokument genau sich selbst.
		var etwasGetan bool
		bg.Run("Texterkennung", func() { etwasGetan = w.step(ctx) })
		if etwasGetan {
			wait = busyInterval
		}
		select {
		case <-ctx.Done():
			w.log.Info().Msg("OCR-Arbeiter beendet")
			return
		case <-time.After(wait):
		}
	}
}

// step erledigt höchstens einen Auftrag. Rückgabe: ob etwas zu tun war.
func (w *Worker) step(ctx context.Context) bool {
	// Vor dem Zugreifen nachsehen, nicht danach: Ein Auftrag, den wir schon
	// auf 'running' gesetzt haben, hätte einen Versuch verbraucht, obwohl es
	// gar nicht an ihm lag.
	if free, err := freeBytes(w.extract.TmpDir); err == nil && free < minFreeBytes {
		// Nur einmal je Zustandswechsel melden – sonst füllt die Warnung im
		// Fünf-Sekunden-Takt den Log, und das ausgerechnet, wenn der Platz
		// knapp ist.
		if !w.paused {
			w.paused = true
			w.log.Warn().
				Uint64("free_mb", free/(1<<20)).
				Uint64("needed_mb", uint64(minFreeBytes)/(1<<20)).
				Msg("Texterkennung pausiert: zu wenig Platz auf der Platte")
		}
		return false
	} else if w.paused {
		w.paused = false
		w.log.Info().Msg("Texterkennung fortgesetzt: wieder genug Platz")
	}

	job, err := w.store.Claim(ctx)
	if errors.Is(err, ErrNoJob) {
		return false
	}
	if err != nil {
		// Datenbank weg oder Kontext beendet – nicht in einer engen Schleife
		// weiterprobieren.
		if ctx.Err() == nil {
			w.log.Warn().Err(err).Msg("Auftrag holen fehlgeschlagen")
		}
		return false
	}

	start := time.Now()
	res, err := w.process(ctx, job)
	if err != nil {
		if ctx.Err() != nil {
			// Herunterfahren mitten im Lauf ist kein Fehler der Datei. Der
			// Auftrag steht auf 'running' und wird beim nächsten Start wieder
			// aufgegriffen – genau dafür nimmt Claim auch 'running' mit.
			return false
		}
		w.log.Warn().Err(err).
			Str("file", job.Name).Int("attempt", job.Attempts).
			Msg("Texterkennung fehlgeschlagen")
		if err := w.store.Fail(context.WithoutCancel(ctx), job, err); err != nil {
			w.log.Error().Err(err).Msg("Fehlschlag konnte nicht vermerkt werden")
		}
		return true
	}

	if err := w.store.Complete(ctx, job, res); err != nil {
		w.log.Error().Err(err).Str("file", job.Name).Msg("Text konnte nicht gespeichert werden")
		return true
	}
	w.log.Info().
		Str("file", job.Name).
		Str("source", res.Source).
		Int("chars", res.Chars).
		Int("pages", res.Pages).
		Dur("took", time.Since(start)).
		Msg("Text erkannt")
	return true
}

func (w *Worker) process(ctx context.Context, job *Job) (Result, error) {
	if job.Size > maxFileBytes {
		return Result{}, fmt.Errorf("Datei zu groß für die Texterkennung (%d Bytes)", job.Size)
	}
	if !Supported(job.MimeType, job.Name) {
		// Kann vorkommen, wenn der MIME-Typ nachträglich korrigiert wurde.
		// Kein Fehler – nur nichts zu tun.
		return Result{}, w.store.Skip(ctx, job.ID, "Dateityp trägt keinen Text")
	}

	// Die Werkzeuge sind externe Programme und brauchen einen Pfad; die
	// Ablageschicht gibt grundsätzlich nur einen Datenstrom heraus, weil sie
	// später auch S3 bedienen soll.
	//
	// Bei der lokalen Ablage gibt es den Pfad aber wirklich – dann wird
	// **nicht** kopiert. Das ist kein Feinschliff: Eine Kopie bedeutet, dass
	// neben jedem Dokument kurzzeitig noch einmal dieselbe Menge auf der
	// Platte liegt, und die Erkennung läuft ausgerechnet auf einer Maschine,
	// die auch alle Dateien beherbergt.
	path, cleanup, err := w.materialize(ctx, job)
	if err != nil {
		return Result{}, err
	}
	defer cleanup()

	return w.extract.Extract(ctx, path, job.MimeType, job.Name)
}

// materialize besorgt einen Pfad zur Datei – ohne Kopie, wo es geht.
//
// Der zurückgegebene cleanup räumt nur auf, wenn tatsächlich kopiert wurde. Er
// darf das Original niemals löschen; deshalb die klare Trennung der beiden
// Zweige statt eines gemeinsamen Pfads mit einem Merker.
func (w *Worker) materialize(ctx context.Context, job *Job) (string, func(), error) {
	if lp, ok := w.blobs.(storage.LocalPather); ok {
		path, err := lp.LocalPath(job.StorageKey)
		if err == nil {
			// Nichts zu räumen: Das ist die Datei selbst. Gelesen wird sie nur.
			return path, func() {}, nil
		}
		// Kein Abbruch – der Weg über die Kopie funktioniert weiterhin.
		w.log.Debug().Err(err).Str("key", job.StorageKey).
			Msg("kein lokaler Pfad, weiche auf Kopie aus")
	}

	rc, err := w.blobs.Get(ctx, job.StorageKey)
	if err != nil {
		return "", nil, fmt.Errorf("Datei lesen: %w", err)
	}
	defer rc.Close()

	// Die Endung mitnehmen: ocrmypdf und tesseract entscheiden auch danach,
	// was sie vor sich haben.
	pattern := "ocr-*" + extOf(job.Name)
	var tmp *os.File
	if tf, ok := w.blobs.(storage.TempFiler); ok {
		// Über die Ablageschicht, damit die Kopie auf derselben Platte landet
		// wie die Daten – nicht in einem knappen /tmp im Container.
		tmp, err = tf.CreateTemp(pattern)
	} else {
		tmp, err = os.CreateTemp("", pattern)
	}
	if err != nil {
		return "", nil, fmt.Errorf("Zwischendatei anlegen: %w", err)
	}
	name := tmp.Name()
	cleanup := func() { os.Remove(name) }

	if _, err := io.Copy(tmp, rc); err != nil {
		tmp.Close()
		cleanup()
		return "", nil, fmt.Errorf("Zwischendatei schreiben: %w", err)
	}
	if err := tmp.Close(); err != nil {
		cleanup()
		return "", nil, fmt.Errorf("Zwischendatei schließen: %w", err)
	}
	return name, cleanup, nil
}

func extOf(name string) string {
	for i := len(name) - 1; i >= 0 && len(name)-i <= 6; i-- {
		if name[i] == '.' {
			return name[i:]
		}
		if name[i] == '/' {
			break
		}
	}
	return ""
}