internal/ingestion/mediacache/store.go
Rama main · réplica de lectura del repositorio original.
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Package mediacache stores decoded parser media for the short-lived,
// run-scoped handoff between Parser, vision enhancement and Chunker. It uses
// an independent one-day TTL and media manifest; persistence success purges
// media through that manifest.
package mediacache
import (
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"strings"
"time"
"github.com/cespare/xxhash/v2"
"ragflow/internal/entity"
"ragflow/internal/ingestion/chunkcache"
"ragflow/internal/ingestion/component/globals"
"ragflow/internal/parser/media"
)
const (
keyNamespace = "mc"
TTL = 24 * time.Hour
)
// Store is a run-scoped media cache backed by a Backend. It dedupes writes by
// digest and bounds how many distinct images, and how many decoded bytes, one
// run may keep in the backend.
type Store struct {
backend Backend
scope string
maxImages int
maxBytes int64
}
// Backend provides the cache lifecycle operations and atomic quota admission.
// The quota, media value, and manifest update happen in one MediaPut operation
// so separate Store instances and processes cannot each admit a full run
// allowance.
type Backend interface {
chunkcache.Store
MediaPut(
ctx context.Context,
mediaKey, digestKey, stateKey, digest, value string,
manifestKey, manifestMember string,
maxImages int,
maxBytes int64,
ttl time.Duration,
) (accepted, inserted bool, err error)
}
// New returns a cache bound to one ingestion run. A nil backend is a valid
// value so callers can surface configured-cache failures only when media is
// actually encountered.
func New(backend Backend, scope string) *Store {
return &Store{
backend: backend,
scope: strings.TrimSpace(scope),
maxImages: media.MediaMaxImagesPerRun(),
maxBytes: media.MediaMaxBytesPerRun(),
}
}
// Put validates and atomically stores one decoded image with its manifest and
// quota state. Any write or lifecycle failure is returned.
func (s *Store) Put(ctx context.Context, format, mime string, data []byte) (entity.MediaRef, error) {
if s == nil || s.backend == nil {
return entity.MediaRef{}, errors.New("mediacache: cache is not configured")
}
if s.scope == "" {
return entity.MediaRef{}, errors.New("mediacache: run scope is empty")
}
if taskID := globals.TaskID(ctx); taskID != "" && taskID != s.scope {
return entity.MediaRef{}, fmt.Errorf("mediacache: run scope %q does not match task %q", s.scope, taskID)
}
if strings.TrimSpace(format) == "" {
return entity.MediaRef{}, errors.New("mediacache: media format is empty")
}
if strings.TrimSpace(mime) == "" {
return entity.MediaRef{}, errors.New("mediacache: media type is empty")
}
if len(data) == 0 || len(data) > media.ImageMaxBytes() {
return entity.MediaRef{}, fmt.Errorf("mediacache: image size %d is outside the 1..%d-byte limit", len(data), media.ImageMaxBytes())
}
sum := sha256.Sum256(data)
digest := hex.EncodeToString(sum[:])
ref := entity.MediaRef{
Format: format,
Kind: entity.MediaKindKV,
Scope: s.scope,
SHA256: digest,
MIME: strings.ToLower(strings.TrimSpace(mime)),
Size: int64(len(data)),
}
key := s.key(digest)
mk := ""
if globals.TaskID(ctx) != "" {
mk = manifestKey(s.scope)
}
digestKey, stateKey := quotaKeys(s.scope)
accepted, _, err := s.backend.MediaPut(
ctx,
key,
digestKey,
stateKey,
digest,
string(data),
mk,
key,
s.maxImages,
s.maxBytes,
TTL,
)
if err != nil {
return entity.MediaRef{}, fmt.Errorf("mediacache: write image %s: %w", mediaShortDigest(digest), err)
}
if !accepted {
return entity.MediaRef{}, fmt.Errorf(
"mediacache: run media limit exceeded (%d images / %d bytes): %w",
s.maxImages,
s.maxBytes,
media.ErrMediaLimitExceeded,
)
}
return ref, nil
}
// Get resolves a reference and validates its scope, digest and encoded size.
// A cache miss is media.ErrMediaNotFound; backend and corruption failures are
// distinct errors so callers cannot silently turn them into a skipped image.
func (s *Store) Get(ctx context.Context, ref entity.MediaRef) ([]byte, error) {
if err := ctx.Err(); err != nil {
return nil, err
}
if s == nil || s.backend == nil {
return nil, errors.New("mediacache: cache is not configured")
}
if s.scope == "" {
return nil, errors.New("mediacache: run scope is empty")
}
if taskID := globals.TaskID(ctx); taskID != "" && taskID != s.scope {
return nil, fmt.Errorf("mediacache: run scope %q does not match task %q", s.scope, taskID)
}
if !ref.Valid() || ref.Kind != entity.MediaKindKV {
return nil, errors.New("mediacache: invalid kvrocks media reference")
}
if ref.Scope != s.scope {
return nil, fmt.Errorf("mediacache: reference scope %q does not match run scope %q", ref.Scope, s.scope)
}
if ref.Size > int64(media.ImageMaxBytes()) {
return nil, fmt.Errorf("mediacache: image %s exceeds the %d-byte limit", mediaShortDigest(ref.SHA256), media.ImageMaxBytes())
}
value, err := s.backend.Get(ctx, s.key(ref.SHA256))
if err != nil {
return nil, fmt.Errorf("mediacache: read image %s: %w", mediaShortDigest(ref.SHA256), err)
}
if value == "" {
return nil, fmt.Errorf("%w: image %s", media.ErrMediaNotFound, mediaShortDigest(ref.SHA256))
}
data, err := media.ValidatePayload([]byte(value), ref)
if err != nil {
return nil, fmt.Errorf("mediacache: validate image %s: %w", mediaShortDigest(ref.SHA256), err)
}
return data, nil
}
func (s *Store) key(digest string) string {
return fmt.Sprintf("%s:media:%x:%s", keyNamespace, xxhash.Sum64String(s.scope), digest)
}
func mediaShortDigest(digest string) string {
if len(digest) <= 12 {
return digest
}
return digest[:12]
}
func manifestKey(scope string) string {
return fmt.Sprintf("%s:manifest:%x", keyNamespace, xxhash.Sum64String(scope))
}
func quotaKeys(scope string) (digestKey, stateKey string) {
hashedScope := xxhash.Sum64String(scope)
return fmt.Sprintf("%s:quota:%x:digests", keyNamespace, hashedScope),
fmt.Sprintf("%s:quota:%x:state", keyNamespace, hashedScope)
}
// PurgeTask removes media after durable persistence. Failed runs rely on TTL.
//
// Entries and quota keys are reclaimed in one batched DeleteMany; the manifest
// is deleted on its own (a separate Delete) only after that batch succeeds. This
// keeps the manifest intact if the entry batch fails or is only partially
// applied (a large manifest is split across several DELs, and a later batch may
// fail after an earlier one already removed its keys), so the task stays
// resumable instead of orphaning surviving entries under a deleted manifest.
func PurgeTask(ctx context.Context, backend chunkcache.Store, taskID string) error {
if backend == nil || taskID == "" {
return nil
}
mk := manifestKey(taskID)
keys, err := backend.SMembers(ctx, mk)
if err != nil {
return fmt.Errorf("mediacache: read manifest: %w", err)
}
digestKey, stateKey := quotaKeys(taskID)
batch := append([]string{}, keys...)
batch = append(batch, digestKey, stateKey)
if !backend.DeleteMany(ctx, batch...) {
return fmt.Errorf("mediacache: purge task %s", taskID)
}
if !backend.Delete(ctx, mk) {
return fmt.Errorf("mediacache: delete manifest %s", taskID)
}
return nil
}