Changeset 0.12.0 (#63)
This commit is contained in:
294
server/extraction/queue.go
Normal file
294
server/extraction/queue.go
Normal file
@@ -0,0 +1,294 @@
|
||||
package extraction
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ── Status Constants ───────────────────────
|
||||
|
||||
const (
|
||||
StatusPending = "pending"
|
||||
StatusProcessing = "processing"
|
||||
StatusComplete = "complete"
|
||||
StatusFailed = "failed"
|
||||
StatusSkipped = "skipped" // not extractable (images, etc.)
|
||||
)
|
||||
|
||||
// ── Queue Item ─────────────────────────────
|
||||
|
||||
// QueueItem is the status.json written to the processing directory.
|
||||
// The sidecar extractor watches this directory for pending items.
|
||||
type QueueItem struct {
|
||||
AttachmentID string `json:"attachment_id"`
|
||||
StorageKey string `json:"storage_key"`
|
||||
ContentType string `json:"content_type"`
|
||||
Filename string `json:"filename"`
|
||||
Status string `json:"status"`
|
||||
Error string `json:"error,omitempty"`
|
||||
OutputKey string `json:"output_key,omitempty"` // path to extracted text
|
||||
QueuedAt string `json:"queued_at"`
|
||||
StartedAt string `json:"started_at,omitempty"`
|
||||
CompletedAt string `json:"completed_at,omitempty"`
|
||||
}
|
||||
|
||||
// ── Queue Manager ──────────────────────────
|
||||
|
||||
// Queue manages the filesystem-based extraction queue.
|
||||
// Processing state lives in {storagePath}/processing/{attachment_id}/status.json.
|
||||
// The sidecar extractor watches for pending items and updates status.
|
||||
type Queue struct {
|
||||
storagePath string
|
||||
concurrency int
|
||||
sem chan struct{} // semaphore for concurrent extraction limit
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// NewQueue creates a new extraction queue manager.
|
||||
// storagePath is the PVC mount (e.g., /data/storage).
|
||||
// concurrency limits parallel extractions (default 3).
|
||||
func NewQueue(storagePath string, concurrency int) (*Queue, error) {
|
||||
if concurrency <= 0 {
|
||||
concurrency = 3
|
||||
}
|
||||
|
||||
procDir := filepath.Join(storagePath, "processing")
|
||||
if err := os.MkdirAll(procDir, 0750); err != nil {
|
||||
return nil, fmt.Errorf("extraction: create processing dir: %w", err)
|
||||
}
|
||||
|
||||
return &Queue{
|
||||
storagePath: storagePath,
|
||||
concurrency: concurrency,
|
||||
sem: make(chan struct{}, concurrency),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Enqueue adds an attachment to the extraction queue.
|
||||
// Creates {storagePath}/processing/{id}/status.json with status "pending".
|
||||
func (q *Queue) Enqueue(attachmentID, storageKey, contentType, filename string) error {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
|
||||
itemDir := filepath.Join(q.storagePath, "processing", attachmentID)
|
||||
if err := os.MkdirAll(itemDir, 0750); err != nil {
|
||||
return fmt.Errorf("extraction: create item dir: %w", err)
|
||||
}
|
||||
|
||||
item := QueueItem{
|
||||
AttachmentID: attachmentID,
|
||||
StorageKey: storageKey,
|
||||
ContentType: contentType,
|
||||
Filename: filename,
|
||||
Status: StatusPending,
|
||||
QueuedAt: time.Now().UTC().Format(time.RFC3339),
|
||||
}
|
||||
|
||||
return q.writeStatus(attachmentID, &item)
|
||||
}
|
||||
|
||||
// GetStatus reads the current extraction status for an attachment.
|
||||
func (q *Queue) GetStatus(attachmentID string) (*QueueItem, error) {
|
||||
statusPath := filepath.Join(q.storagePath, "processing", attachmentID, "status.json")
|
||||
data, err := os.ReadFile(statusPath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil // not queued
|
||||
}
|
||||
return nil, fmt.Errorf("extraction: read status: %w", err)
|
||||
}
|
||||
|
||||
var item QueueItem
|
||||
if err := json.Unmarshal(data, &item); err != nil {
|
||||
return nil, fmt.Errorf("extraction: parse status: %w", err)
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
// MarkComplete updates status to complete and records the output key.
|
||||
func (q *Queue) MarkComplete(attachmentID, outputKey string) error {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
|
||||
item, err := q.GetStatus(attachmentID)
|
||||
if err != nil || item == nil {
|
||||
return fmt.Errorf("extraction: item not found: %s", attachmentID)
|
||||
}
|
||||
|
||||
item.Status = StatusComplete
|
||||
item.OutputKey = outputKey
|
||||
item.CompletedAt = time.Now().UTC().Format(time.RFC3339)
|
||||
return q.writeStatus(attachmentID, item)
|
||||
}
|
||||
|
||||
// MarkFailed updates status to failed with an error message.
|
||||
func (q *Queue) MarkFailed(attachmentID, errMsg string) error {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
|
||||
item, err := q.GetStatus(attachmentID)
|
||||
if err != nil || item == nil {
|
||||
return fmt.Errorf("extraction: item not found: %s", attachmentID)
|
||||
}
|
||||
|
||||
item.Status = StatusFailed
|
||||
item.Error = errMsg
|
||||
item.CompletedAt = time.Now().UTC().Format(time.RFC3339)
|
||||
return q.writeStatus(attachmentID, item)
|
||||
}
|
||||
|
||||
// Cleanup removes the processing directory for a completed/failed item.
|
||||
func (q *Queue) Cleanup(attachmentID string) error {
|
||||
itemDir := filepath.Join(q.storagePath, "processing", attachmentID)
|
||||
return os.RemoveAll(itemDir)
|
||||
}
|
||||
|
||||
// RecoverStale scans for items stuck in "processing" state (crash recovery).
|
||||
// Items older than maxAge are reset to "pending" for re-processing.
|
||||
func (q *Queue) RecoverStale(maxAge time.Duration) (int, error) {
|
||||
procDir := filepath.Join(q.storagePath, "processing")
|
||||
entries, err := os.ReadDir(procDir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return 0, nil
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
|
||||
recovered := 0
|
||||
cutoff := time.Now().Add(-maxAge)
|
||||
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
item, err := q.GetStatus(entry.Name())
|
||||
if err != nil || item == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
if item.Status != StatusProcessing {
|
||||
continue
|
||||
}
|
||||
|
||||
// Check if started_at is older than maxAge
|
||||
if item.StartedAt != "" {
|
||||
startedAt, err := time.Parse(time.RFC3339, item.StartedAt)
|
||||
if err == nil && startedAt.Before(cutoff) {
|
||||
q.mu.Lock()
|
||||
item.Status = StatusPending
|
||||
item.StartedAt = ""
|
||||
item.Error = "recovered from stale processing state"
|
||||
q.writeStatus(entry.Name(), item)
|
||||
q.mu.Unlock()
|
||||
recovered++
|
||||
log.Printf("extraction: recovered stale item %s", entry.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return recovered, nil
|
||||
}
|
||||
|
||||
// ListPending returns all items with status "pending".
|
||||
func (q *Queue) ListPending() ([]QueueItem, error) {
|
||||
return q.listByStatus(StatusPending)
|
||||
}
|
||||
|
||||
// ListAll returns all items in the processing directory.
|
||||
func (q *Queue) ListAll() ([]QueueItem, error) {
|
||||
procDir := filepath.Join(q.storagePath, "processing")
|
||||
entries, err := os.ReadDir(procDir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var items []QueueItem
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
item, err := q.GetStatus(entry.Name())
|
||||
if err != nil || item == nil {
|
||||
continue
|
||||
}
|
||||
items = append(items, *item)
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// ── Internal Helpers ───────────────────────
|
||||
|
||||
func (q *Queue) writeStatus(attachmentID string, item *QueueItem) error {
|
||||
statusPath := filepath.Join(q.storagePath, "processing", attachmentID, "status.json")
|
||||
data, err := json.MarshalIndent(item, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Atomic write: temp + rename
|
||||
tmpPath := statusPath + ".tmp"
|
||||
if err := os.WriteFile(tmpPath, data, 0640); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmpPath, statusPath)
|
||||
}
|
||||
|
||||
func (q *Queue) listByStatus(status string) ([]QueueItem, error) {
|
||||
all, err := q.ListAll()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var filtered []QueueItem
|
||||
for _, item := range all {
|
||||
if item.Status == status {
|
||||
filtered = append(filtered, item)
|
||||
}
|
||||
}
|
||||
return filtered, nil
|
||||
}
|
||||
|
||||
// ── Extractable Check ──────────────────────
|
||||
|
||||
// IsExtractable returns true if the content type supports text extraction.
|
||||
// Images and plain text do not need extraction (handled inline by upload handler).
|
||||
var extractableTypes = map[string]bool{
|
||||
"application/pdf": true,
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document": true,
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": true,
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation": true,
|
||||
"application/msword": true,
|
||||
"application/vnd.ms-excel": true,
|
||||
"application/vnd.oasis.opendocument.text": true,
|
||||
"application/vnd.oasis.opendocument.spreadsheet": true,
|
||||
"application/vnd.oasis.opendocument.presentation": true,
|
||||
"application/rtf": true,
|
||||
}
|
||||
|
||||
// IsExtractable returns true if the given content type requires sidecar extraction.
|
||||
func IsExtractable(contentType string) bool {
|
||||
return extractableTypes[contentType]
|
||||
}
|
||||
|
||||
// ── Null Queue ─────────────────────────────
|
||||
|
||||
// NullQueue is a no-op implementation for when extraction is disabled.
|
||||
type NullQueue struct{}
|
||||
|
||||
func (NullQueue) Enqueue(_, _, _, _ string) error { return nil }
|
||||
func (NullQueue) GetStatus(_ string) (*QueueItem, error) { return nil, nil }
|
||||
func (NullQueue) MarkComplete(_, _ string) error { return nil }
|
||||
func (NullQueue) MarkFailed(_, _ string) error { return nil }
|
||||
func (NullQueue) Cleanup(_ string) error { return nil }
|
||||
func (NullQueue) RecoverStale(_ time.Duration) (int, error) { return 0, nil }
|
||||
func (NullQueue) ListPending() ([]QueueItem, error) { return nil, nil }
|
||||
func (NullQueue) ListAll() ([]QueueItem, error) { return nil, nil }
|
||||
Reference in New Issue
Block a user