332 lines
12 KiB
Go
332 lines
12 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"os"
|
||
|
|
"strconv"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
apprepo "github.com/Tencent/WeKnora/internal/application/repository"
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/requestbudget"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"gorm.io/gorm"
|
||
|
|
"gorm.io/gorm/clause"
|
||
|
|
)
|
||
|
|
|
||
|
|
var ErrWikiRecoveryBudgetStopped = requestbudget.ErrStopped
|
||
|
|
|
||
|
|
// Check before configuration short-circuits so a disabled/deleted source or
|
||
|
|
// expired deadline still settles durably without waiting for a model call.
|
||
|
|
func (s *wikiIngestService) checkWikiRecoveryBudgets(ctx context.Context, payload WikiIngestPayload) error {
|
||
|
|
if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
var rows []types.WikiRecoveryBudget
|
||
|
|
if err := s.checkpointDB.WithContext(ctx).Where("tenant_id = ? AND knowledge_base_id = ? AND status IN ?", payload.TenantID, payload.KnowledgeBaseID, []string{"active", "stopped_budget"}).Find(&rows).Error; err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
|
||
|
|
for _, row := range rows {
|
||
|
|
if err := repo.Check(ctx, row); err != nil {
|
||
|
|
if archiveErr := repo.StopAndArchive(ctx, row, "guard_rejected", "preflight"); archiveErr != nil {
|
||
|
|
return errors.Join(err, archiveErr)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *wikiIngestService) rejectOrdinaryRecoveryOverlap(ctx context.Context, payload WikiIngestPayload, ops []WikiPendingOp) error {
|
||
|
|
if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
ids := make([]string, 0, len(ops))
|
||
|
|
for _, op := range ops {
|
||
|
|
ids = append(ids, op.KnowledgeID)
|
||
|
|
}
|
||
|
|
var count int64
|
||
|
|
if err := s.checkpointDB.WithContext(ctx).Model(&types.WikiRecoveryBudget{}).
|
||
|
|
Where("tenant_id = ? AND knowledge_base_id = ? AND knowledge_id IN ? AND status IN ?", payload.TenantID, payload.KnowledgeBaseID, ids, []string{"active", "stopped_budget"}).Count(&count).Error; err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if count != 0 {
|
||
|
|
return fmt.Errorf("%w: ordinary operation overlaps controlled recovery", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type wikiRecoveryGuard struct {
|
||
|
|
repo apprepo.WikiRecoveryBudgetRepository
|
||
|
|
row types.WikiRecoveryBudget
|
||
|
|
cancel context.CancelCauseFunc
|
||
|
|
mu sync.Mutex
|
||
|
|
stopped error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *wikiRecoveryGuard) Key() string {
|
||
|
|
return fmt.Sprintf("%d:%s", g.row.TenantID, g.row.Identity().Key())
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *wikiRecoveryGuard) fail(err error) error {
|
||
|
|
if err != nil {
|
||
|
|
g.mu.Lock()
|
||
|
|
if g.stopped == nil {
|
||
|
|
g.stopped = err
|
||
|
|
}
|
||
|
|
err = g.stopped
|
||
|
|
g.mu.Unlock()
|
||
|
|
g.cancel(err)
|
||
|
|
}
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *wikiRecoveryGuard) Check(ctx context.Context) error {
|
||
|
|
g.mu.Lock()
|
||
|
|
stopped := g.stopped
|
||
|
|
g.mu.Unlock()
|
||
|
|
if stopped != nil {
|
||
|
|
return stopped
|
||
|
|
}
|
||
|
|
if !time.Now().Before(g.row.DeadlineAt) {
|
||
|
|
cleanup, cancel := wikiIngestCleanupContext(ctx)
|
||
|
|
defer cancel()
|
||
|
|
return g.fail(g.repo.Check(cleanup, g.row))
|
||
|
|
}
|
||
|
|
if err := ctx.Err(); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return g.fail(g.repo.Check(ctx, g.row))
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *wikiRecoveryGuard) Reserve(ctx context.Context, model, phase string) error {
|
||
|
|
if err := g.Check(ctx); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return g.fail(g.repo.Reserve(ctx, g.row, model, phase))
|
||
|
|
}
|
||
|
|
|
||
|
|
func wikiRecoveryFromContext(ctx context.Context) *wikiRecoveryGuard {
|
||
|
|
guard, _ := requestbudget.FromContext(ctx).(*wikiRecoveryGuard)
|
||
|
|
return guard
|
||
|
|
}
|
||
|
|
|
||
|
|
func wikiRecoveryIdentityFromContext(ctx context.Context) types.WikiRecoveryIdentity {
|
||
|
|
if guard := wikiRecoveryFromContext(ctx); guard != nil {
|
||
|
|
return guard.row.Identity()
|
||
|
|
}
|
||
|
|
return types.WikiRecoveryIdentity{}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *wikiIngestService) bindWikiRecovery(ctx context.Context, payload WikiIngestPayload, identity types.WikiRecoveryIdentity, knowledgeID, modelID string) (context.Context, context.CancelFunc, error) {
|
||
|
|
if !identity.Present() {
|
||
|
|
return ctx, func() {}, nil
|
||
|
|
}
|
||
|
|
repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
|
||
|
|
row, err := repo.Load(ctx, payload.TenantID, identity)
|
||
|
|
if err != nil {
|
||
|
|
return ctx, func() {}, err
|
||
|
|
}
|
||
|
|
if row.KnowledgeBaseID != payload.KnowledgeBaseID || (knowledgeID != "" && row.KnowledgeID != knowledgeID) || row.ModelID != modelID {
|
||
|
|
return ctx, func() {}, fmt.Errorf("%w: task scope mismatch", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
// Deployment must explicitly attest its immutable release and image. A
|
||
|
|
// controlled op cannot execute in an unreviewed local or stale API build.
|
||
|
|
if row.ReviewedRelease == "" || row.ReviewedAPIImage == "" ||
|
||
|
|
os.Getenv("WEKNORA_REVIEWED_RELEASE") != row.ReviewedRelease || os.Getenv("WEKNORA_REVIEWED_API_IMAGE") != row.ReviewedAPIImage {
|
||
|
|
return ctx, func() {}, fmt.Errorf("%w: reviewed runtime mismatch", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
ctx, deadlineCancel := context.WithDeadlineCause(ctx, row.DeadlineAt, fmt.Errorf("%w: deadline", ErrWikiRecoveryBudgetStopped))
|
||
|
|
ctx, cancel := context.WithCancelCause(ctx)
|
||
|
|
guard := &wikiRecoveryGuard{repo: repo, row: *row, cancel: cancel}
|
||
|
|
ctx = requestbudget.WithGuard(ctx, guard)
|
||
|
|
if err := guard.Check(ctx); err != nil {
|
||
|
|
deadlineCancel()
|
||
|
|
cancel(err)
|
||
|
|
return ctx, func() {}, err
|
||
|
|
}
|
||
|
|
return ctx, func() { deadlineCancel(); cancel(nil) }, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *wikiIngestService) settleWikiRecoveryStop(ctx context.Context, identity types.WikiRecoveryIdentity, tenant uint64, err error) error {
|
||
|
|
if !identity.Present() || err == nil || errors.Is(err, errWikiCheckpointYield) {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
reason := wikiRecoveryFailureReason(ctx, err)
|
||
|
|
if !errors.Is(err, ErrWikiRecoveryBudgetStopped) {
|
||
|
|
err = fmt.Errorf("%w: recovery phase failed: %w", ErrWikiRecoveryBudgetStopped, err)
|
||
|
|
}
|
||
|
|
cleanup, cancel := wikiIngestCleanupContext(ctx)
|
||
|
|
defer cancel()
|
||
|
|
repo := apprepo.WikiRecoveryBudgetRepository{DB: s.checkpointDB}
|
||
|
|
row, loadErr := repo.Load(cleanup, tenant, identity)
|
||
|
|
if loadErr != nil {
|
||
|
|
return errors.Join(err, loadErr)
|
||
|
|
}
|
||
|
|
phase := row.LastPhase
|
||
|
|
if phase == "" {
|
||
|
|
phase = requestbudget.Phase(ctx)
|
||
|
|
}
|
||
|
|
if phase == "" {
|
||
|
|
phase = "recovery"
|
||
|
|
}
|
||
|
|
return errors.Join(err, repo.StopAndArchive(cleanup, *row, reason, phase))
|
||
|
|
}
|
||
|
|
|
||
|
|
// Persist only a bounded category, never provider response bodies or secrets.
|
||
|
|
// Phase failures still stop the same recovery; they are not budget exhaustion.
|
||
|
|
func wikiRecoveryFailureReason(ctx context.Context, err error) string {
|
||
|
|
if errors.Is(err, ErrWikiRecoveryBudgetStopped) || errors.Is(context.Cause(ctx), ErrWikiRecoveryBudgetStopped) {
|
||
|
|
return "guard_rejected"
|
||
|
|
}
|
||
|
|
var outputErr *wikiOutputError
|
||
|
|
if errors.As(err, &outputErr) {
|
||
|
|
if outputErr.incomplete {
|
||
|
|
return "model_output_incomplete"
|
||
|
|
}
|
||
|
|
return "model_output_invalid"
|
||
|
|
}
|
||
|
|
if errors.Is(err, context.DeadlineExceeded) {
|
||
|
|
return "phase_timeout"
|
||
|
|
}
|
||
|
|
if errors.Is(err, context.Canceled) {
|
||
|
|
return "phase_cancelled"
|
||
|
|
}
|
||
|
|
return "phase_failed"
|
||
|
|
}
|
||
|
|
|
||
|
|
func selectWikiFinalizeGroup(rows []*types.TaskPendingOp) ([]*types.TaskPendingOp, types.WikiRecoveryIdentity, error) {
|
||
|
|
var first types.WikiRecoveryIdentity
|
||
|
|
var selected []*types.TaskPendingOp
|
||
|
|
identities := make([]types.WikiRecoveryIdentity, len(rows))
|
||
|
|
for i, row := range rows {
|
||
|
|
var identity types.WikiRecoveryIdentity
|
||
|
|
var keyIdentity types.WikiRecoveryIdentity
|
||
|
|
if strings.HasPrefix(row.DedupKey, "recovery:") {
|
||
|
|
parts := strings.SplitN(row.DedupKey, ":", 6)
|
||
|
|
if len(parts) != 6 {
|
||
|
|
return nil, first, fmt.Errorf("%w: invalid finalize dedup identity", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
keyIdentity.DeadLetterID, _ = strconv.ParseInt(parts[1], 10, 64)
|
||
|
|
keyIdentity.Attempt, _ = strconv.Atoi(parts[2])
|
||
|
|
keyIdentity.PendingOpID, _ = strconv.ParseInt(parts[3], 10, 64)
|
||
|
|
keyIdentity.Version, _ = strconv.Atoi(parts[4])
|
||
|
|
}
|
||
|
|
if len(row.Payload) > 0 {
|
||
|
|
if err := json.Unmarshal(row.Payload, &identity); err != nil {
|
||
|
|
return nil, keyIdentity, fmt.Errorf("%w: invalid finalize payload: %v", ErrWikiRecoveryBudgetStopped, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if keyIdentity.Present() || identity.Present() {
|
||
|
|
if !keyIdentity.Valid() || keyIdentity != identity {
|
||
|
|
if keyIdentity.Valid() {
|
||
|
|
identity = keyIdentity
|
||
|
|
}
|
||
|
|
return nil, identity, fmt.Errorf("%w: finalize payload and dedup identity mismatch", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
var payload wikiFinalizeRow
|
||
|
|
if err := json.Unmarshal(row.Payload, &payload); err != nil {
|
||
|
|
return nil, identity, fmt.Errorf("%w: invalid finalize content: %v", ErrWikiRecoveryBudgetStopped, err)
|
||
|
|
}
|
||
|
|
if (row.Op != wikiFinalizeOpSlug && row.Op != wikiFinalizeOpChange) ||
|
||
|
|
(row.Op == wikiFinalizeOpSlug && payload.Slug == "") || (row.Op == wikiFinalizeOpChange && payload.Change == nil) {
|
||
|
|
return nil, identity, fmt.Errorf("%w: incomplete finalize content", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
identities[i] = identity
|
||
|
|
// Controlled rows go first so an ordinary index update cannot absorb
|
||
|
|
// this recovery's work and leave its own finalize lane unbudgeted.
|
||
|
|
if !first.Present() && identity.Present() {
|
||
|
|
first = identity
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for i, row := range rows {
|
||
|
|
if identities[i] == first {
|
||
|
|
selected = append(selected, row)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if first.Present() && !first.Valid() {
|
||
|
|
return nil, first, fmt.Errorf("%w: invalid finalize identity", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
return selected, first, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *wikiIngestService) rejectOrdinaryFinalizeRecoveryOverlap(ctx context.Context, payload WikiIngestPayload) error {
|
||
|
|
if s.checkpointDB == nil || !s.checkpointDB.Migrator().HasTable(&types.WikiRecoveryBudget{}) {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
var count int64
|
||
|
|
if err := s.checkpointDB.WithContext(ctx).Model(&types.WikiRecoveryBudget{}).
|
||
|
|
Where("tenant_id = ? AND knowledge_base_id = ? AND status = ?", payload.TenantID, payload.KnowledgeBaseID, "active").Count(&count).Error; err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if count > 0 {
|
||
|
|
return fmt.Errorf("%w: controlled recovery still owns index finalization", ErrWikiIngestConcurrent)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *wikiIngestService) enqueueControlledWikiFinalize(ctx context.Context, payload WikiIngestPayload, slugs []string, titles map[string]string, changes []wikiFinalizeChange) error {
|
||
|
|
guard := wikiRecoveryFromContext(ctx)
|
||
|
|
if guard == nil {
|
||
|
|
return fmt.Errorf("%w: missing recovery guard", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
if err := guard.Check(ctx); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
identity := guard.row.Identity()
|
||
|
|
var rows []types.TaskPendingOp
|
||
|
|
appendRow := func(op, key string, value wikiFinalizeRow) error {
|
||
|
|
value.WikiRecoveryIdentity = identity
|
||
|
|
body, err := json.Marshal(value)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
rows = append(rows, types.TaskPendingOp{TenantID: payload.TenantID, TaskType: wikiFinalizeTaskType, Scope: wikiTaskScope, ScopeID: payload.KnowledgeBaseID, Op: op, DedupKey: "recovery:" + identity.Key() + ":" + key, Payload: body, EnqueuedAt: time.Now().UTC()})
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
for _, slug := range slugs {
|
||
|
|
if err := appendRow(wikiFinalizeOpSlug, slug, wikiFinalizeRow{Slug: slug, Title: titles[slug]}); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for i := range changes {
|
||
|
|
if err := appendRow(wikiFinalizeOpChange, fmt.Sprintf("change:%d", i), wikiFinalizeRow{Change: &changes[i]}); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if len(changes) == 0 {
|
||
|
|
return fmt.Errorf("%w: recovery has no completed document to finalize", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
err := s.checkpointDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||
|
|
var budget types.WikiRecoveryBudget
|
||
|
|
if err := tx.Where("tenant_id = ? AND pending_op_id = ?", payload.TenantID, identity.PendingOpID).Clauses(clause.Locking{Strength: "UPDATE"}).Take(&budget).Error; err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if budget.Identity() != identity || budget.Status != "active" || !time.Now().Before(budget.DeadlineAt) {
|
||
|
|
return fmt.Errorf("%w: finalize handoff no longer active", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
var op types.TaskPendingOp
|
||
|
|
if err := tx.Where("tenant_id = ? AND id = ? AND task_type = ? AND scope_id = ?", payload.TenantID, identity.PendingOpID, wikiTaskType, payload.KnowledgeBaseID).Take(&op).Error; err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
var saved types.WikiRecoveryIdentity
|
||
|
|
if json.Unmarshal(op.Payload, &saved) != nil || saved != identity {
|
||
|
|
return fmt.Errorf("%w: handoff op mismatch", ErrWikiRecoveryBudgetStopped)
|
||
|
|
}
|
||
|
|
if err := tx.Create(&rows).Error; err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return tx.Where("tenant_id = ? AND id = ?", payload.TenantID, identity.PendingOpID).Delete(&types.TaskPendingOp{}).Error
|
||
|
|
})
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
s.scheduleFinalize(ctx, payload)
|
||
|
|
return nil
|
||
|
|
}
|