xinmi/internal/accesscontrol/resource_version_knowledge_cleanup.go

package accesscontrol

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/Tencent/WeKnora/internal/knowledgeversion"
	"github.com/Tencent/WeKnora/internal/types"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// A copy can run for two hours. Wait longer before deleting abandoned targets,
// including targets left by a process whose heartbeat stopped before its work.
const knowledgeVersionCleanupDelay = 3 * time.Hour

func (s *Service) CleanupAbandonedKnowledgeVersions(ctx context.Context) error {
	cleaner, ok := s.versionKnowledge.(interface {
		CleanupKnowledgeVersion(context.Context, string) error
	})
	if !ok {
		return nil
	}
	now := time.Now()
	var job KnowledgeVersionJob
	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		result := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).
			Where("status IN ? AND updated_at < ?", []string{"completed", "failed"}, now.Add(-knowledgeVersionCleanupDelay)).
			Where(`(abandoned_ids IS NOT NULL AND CAST(abandoned_ids AS TEXT) <> '[]' AND CAST(abandoned_ids AS TEXT) <> 'null') OR (status='failed' AND EXISTS (SELECT 1 FROM knowledge_bases kb WHERE kb.id=workbench_knowledge_version_jobs.target_id AND kb.tenant_id=workbench_knowledge_version_jobs.tenant_id AND kb.is_temporary=true AND kb.version_state='copying' AND kb.deleted_at IS NULL))`).
			Order("updated_at").Limit(1).Find(&job)
		if result.Error != nil {
			return result.Error
		}
		if result.RowsAffected == 0 {
			return gorm.ErrRecordNotFound
		}
		return tx.Model(&job).UpdateColumn("updated_at", now).Error
	})
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil
	}
	if err != nil {
		return err
	}
	worker, cancel := context.WithTimeout(knowledgeversion.WithWorker(ctx), 5*time.Minute)
	defer cancel()
	worker = context.WithValue(worker, types.TenantIDContextKey, job.TenantID)
	worker = context.WithValue(worker, types.UserIDContextKey, job.OwnerID)
	worker = context.WithValue(worker, types.TenantRoleContextKey, types.TenantRoleAdmin)
	worker = context.WithValue(worker, types.SystemAdminContextKey, false)
	targets := append([]string{}, job.AbandonedIDs...)
	if job.Status == "failed" {
		targets = append(targets, job.TargetID)
	}
	seen := map[string]bool{}
	for _, target := range targets {
		if target == "" || seen[target] {
			continue
		}
		seen[target] = true
		var kb types.KnowledgeBase
		result := s.db.WithContext(worker).Where("tenant_id=? AND id=?", job.TenantID, target).Limit(1).Find(&kb)
		if result.Error != nil {
			return result.Error
		}
		if result.RowsAffected == 0 {
			if err := s.finishKnowledgeVersionCleanup(worker, job.ID, job.TenantID, target); err != nil {
				return err
			}
			continue
		}
		if !kb.IsTemporary || kb.VersionState != "copying" {
			return fmt.Errorf("refusing to clean non-staging knowledge version %s", target)
		}
		if err := cleaner.CleanupKnowledgeVersion(worker, target); err != nil {
			return err
		}
		if err := s.finishKnowledgeVersionCleanup(worker, job.ID, job.TenantID, target); err != nil {
			return err
		}
	}
	return nil
}

func (s *Service) finishKnowledgeVersionCleanup(ctx context.Context, jobID string, tenant uint64, target string) error {
	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var job KnowledgeVersionJob
		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=?", jobID, tenant).First(&job).Error; err != nil {
			return err
		}
		allowed := job.Status == "failed" && job.TargetID == target
		remaining := make(types.StringArray, 0, len(job.AbandonedIDs))
		for _, id := range job.AbandonedIDs {
			if id == target {
				allowed = true
			} else {
				remaining = append(remaining, id)
			}
		}
		if !allowed || job.TargetID == target && job.Status != "failed" {
			return fmt.Errorf("staging cleanup ownership changed")
		}
		var kb types.KnowledgeBase
		result := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id=? AND id=?", tenant, target).Limit(1).Find(&kb)
		if result.Error != nil {
			return result.Error
		}
		if result.RowsAffected == 1 {
			if !kb.IsTemporary || kb.VersionState != "copying" {
				return fmt.Errorf("staging cleanup target became published")
			}
			var docs []types.Knowledge
			if err := tx.Unscoped().Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Find(&docs).Error; err != nil {
				return err
			}
			var size int64
			for _, doc := range docs {
				size += doc.StorageSize
			}
			var chunks int64
			if err := tx.Model(&types.Chunk{}).Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Count(&chunks).Error; err != nil {
				return err
			}
			if chunks != 0 {
				return fmt.Errorf("staging external cleanup is incomplete")
			}
			if size > 0 {
				r := tx.Model(&types.Tenant{}).Where("id=? AND storage_used>=?", tenant, size).UpdateColumn("storage_used", gorm.Expr("storage_used - ?", size))
				if r.Error != nil {
					return r.Error
				}
				if r.RowsAffected != 1 {
					return fmt.Errorf("staging storage accounting mismatch")
				}
			}
			if err := knowledgeversion.DeleteGraph(tx, tenant, target); err != nil {
				return err
			}
			if err := tx.Unscoped().Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Delete(&types.Chunk{}).Error; err != nil {
				return err
			}
			if err := tx.Unscoped().Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Delete(&types.Knowledge{}).Error; err != nil {
				return err
			}
			if err := tx.Unscoped().Where("tenant_id=? AND id=?", tenant, target).Delete(&types.KnowledgeBase{}).Error; err != nil {
				return err
			}
		}
		return tx.Model(&job).Updates(map[string]any{"abandoned_ids": remaining}).Error
	})
}