xinmi/internal/application/service/extract_graph_budget.go

package service

import (
	"context"
	"fmt"
	"time"

	"github.com/Tencent/WeKnora/internal/models/requestbudget"
	"github.com/Tencent/WeKnora/internal/types"
)

// Reserving the whole execution in the existing dispatch ledger fences both
// concurrent delivery and worker restart. A crash never restores its three sends.
func (s *ChunkExtractService) claimGraphExecution(ctx context.Context, p types.ExtractChunkPayload, taskID string) error {
	e := p.ExpectedSource
	if s.executionDB == nil || e == nil || taskID == "" || p.ModelID == "" || p.TenantID == 0 ||
		p.ChunkID == "" || e.KnowledgeBaseID == "" || e.KnowledgeID == "" || e.ContentSHA256 == "" {
		return fmt.Errorf("%w: missing graph execution receipt", requestbudget.ErrStopped)
	}
	result := s.executionDB.WithContext(ctx).Exec(`UPDATE native_graph_backfill_failures
		SET state='execution_reserved',updated_at=?
		WHERE task_id=? AND tenant_id=? AND knowledge_base_id=? AND chunk_id=? AND knowledge_id=? AND content_hash=?
		AND state='dispatch_requested' AND manual_retries=0 AND reason=''`,
		time.Now().UTC(), taskID, p.TenantID, e.KnowledgeBaseID, p.ChunkID, e.KnowledgeID, e.ContentSHA256)
	if result.Error != nil || result.RowsAffected != 1 {
		return fmt.Errorf("%w: graph execution receipt unavailable or already spent", requestbudget.ErrStopped)
	}
	return nil
}