yovision/Sense/server/app/sense/media_shard/service.go

package media_shard

import (
	"context"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

var (
	ErrNotFound   = errors.New("媒体分片不存在")
	ErrNoCapacity = errors.New("没有可用的媒体分片容量")
	ErrInvalidID  = errors.New("媒体分片标识不符合要求")
)

const staleAfter = 30 * time.Second

type Probe func(context.Context, string) error

type Service struct {
	DB    *gorm.DB
	Probe Probe
	Now   func() time.Time
}

func NewService(db *gorm.DB, probe Probe) *Service {
	return &Service{DB: db, Probe: probe, Now: time.Now}
}

func (s *Service) SyncSpecs(ctx context.Context, specs []Spec) error {
	return s.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		for _, spec := range specs {
			var current Shard
			err := tx.First(&current, "id = ?", spec.ID).Error
			if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
				return err
			}
			if errors.Is(err, gorm.ErrRecordNotFound) {
				current = Shard{ID: spec.ID, Status: StatusUnknown, Detail: "等待健康探测", ConfigVersion: 1}
			} else if current.Name != spec.Name || current.Mode != spec.Mode || current.ControlAPI != spec.ControlAPI || current.Capacity != spec.Capacity {
				current.ConfigVersion++
			}
			current.Name, current.Mode, current.ControlAPI, current.Capacity = spec.Name, spec.Mode, spec.ControlAPI, spec.Capacity
			if err = tx.Save(&current).Error; err != nil {
				return err
			}
		}
		ids := make([]string, 0, len(specs))
		for _, spec := range specs {
			ids = append(ids, spec.ID)
		}
		return tx.Model(&Shard{}).Where("id NOT IN ?", ids).Updates(map[string]any{"status": StatusDisabled, "detail": "分片已从当前运行配置移除", "updated_at": s.now()}).Error
	})
}

func (s *Service) RefreshAll(ctx context.Context) error {
	var shards []Shard
	if err := s.DB.WithContext(ctx).Where("status <> ?", StatusDisabled).Find(&shards).Error; err != nil {
		return err
	}
	var result error
	for _, shard := range shards {
		status, detail := StatusRunning, "Control API 正常"
		if s.Probe == nil || s.Probe(ctx, shard.ControlAPI) != nil {
			status, detail = StatusFailed, "Control API 不可用"
		}
		now := s.now()
		if err := s.DB.WithContext(ctx).Model(&Shard{}).Where("id = ?", shard.ID).Updates(map[string]any{"status": status, "detail": detail, "last_probe_at": &now, "updated_at": now}).Error; err != nil {
			result = errors.Join(result, err)
		}
	}
	return result
}

func (s *Service) EnsureAssignment(ctx context.Context, routeID string) (Shard, error) {
	routeID = strings.TrimSpace(routeID)
	if routeID == "" {
		return Shard{}, ErrInvalidID
	}
	var selected Shard
	err := s.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var existing Assignment
		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&existing, "route_id = ?", routeID).Error; err == nil {
			return tx.First(&selected, "id = ?", existing.ShardID).Error
		} else if !errors.Is(err, gorm.ErrRecordNotFound) {
			return err
		}
		var candidates []Shard
		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("status IN ?", []string{StatusRunning, StatusUnknown}).Find(&candidates).Error; err != nil {
			return err
		}
		counts, err := assignmentCounts(tx)
		if err != nil {
			return err
		}
		eligible := candidates[:0]
		for _, item := range candidates {
			if counts[item.ID] < int64(item.Capacity) {
				eligible = append(eligible, item)
			}
		}
		if len(eligible) == 0 {
			return ErrNoCapacity
		}
		sort.SliceStable(eligible, func(i, j int) bool { return rendezvous(routeID, eligible[i].ID) > rendezvous(routeID, eligible[j].ID) })
		selected = eligible[0]
		now := s.now()
		assignment := Assignment{RouteID: routeID, ShardID: selected.ID, Algorithm: "rendezvous-v1", CreatedAt: now, UpdatedAt: now}
		if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&assignment).Error; err != nil {
			return err
		}
		if assignment.RouteID != "" { // reload also handles a concurrent winner on PostgreSQL.
			if err := tx.First(&existing, "route_id = ?", routeID).Error; err != nil {
				return err
			}
			return tx.First(&selected, "id = ?", existing.ShardID).Error
		}
		return nil
	})
	return selected, err
}

func (s *Service) ResolveRoute(ctx context.Context, routeID string) (Shard, error) {
	var shard Shard
	err := s.DB.WithContext(ctx).Table("sense_media_shards AS s").Select("s.*").Joins("JOIN sense_media_shard_assignments a ON a.shard_id = s.id").Where("a.route_id = ?", routeID).Scan(&shard).Error
	if err != nil {
		return Shard{}, err
	}
	if shard.ID == "" {
		return Shard{}, ErrNotFound
	}
	return shard, nil
}

func (s *Service) List(ctx context.Context) (PageResponse, error) {
	var shards []Shard
	if err := s.DB.WithContext(ctx).Order("name ASC, id ASC").Find(&shards).Error; err != nil {
		return PageResponse{}, err
	}
	counts, err := assignmentCounts(s.DB.WithContext(ctx))
	if err != nil {
		return PageResponse{}, err
	}
	result := PageResponse{List: make([]Response, 0, len(shards)), Count: int64(len(shards))}
	for _, shard := range shards {
		item := s.response(shard, int(counts[shard.ID]))
		result.List = append(result.List, item)
		result.Summary.Total++
		result.Summary.Configured += item.Capacity
		result.Summary.Assigned += item.AssignedPaths
		if item.Status == StatusRunning {
			result.Summary.Available++
		}
		if item.Status == StatusFailed {
			result.Summary.Failed++
			result.Summary.Impacted += item.AssignedPaths
		}
	}
	return result, nil
}

func (s *Service) Get(ctx context.Context, id string) (DetailResponse, error) {
	id = strings.TrimSpace(id)
	if id == "" || len(id) > 64 {
		return DetailResponse{}, ErrInvalidID
	}
	var shard Shard
	if err := s.DB.WithContext(ctx).First(&shard, "id = ?", id).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return DetailResponse{}, ErrNotFound
		}
		return DetailResponse{}, err
	}
	impact, err := s.impact(ctx, id)
	if err != nil {
		return DetailResponse{}, err
	}
	return DetailResponse{Response: s.response(shard, len(impact)), Impact: impact}, nil
}

func (s *Service) Preflight(ctx context.Context, id string) (PreflightResponse, error) {
	detail, err := s.Get(ctx, id)
	if err != nil {
		return PreflightResponse{}, err
	}
	result := PreflightResponse{SourceShardID: id, ImpactedPaths: len(detail.Impact), ExecutionAuthorized: false, RequiresIssue: true}
	result.Checks = append(result.Checks, Check{Name: "源分片状态", Passed: detail.Status == StatusRunning && !detail.Stale, Detail: detail.Detail})
	var candidates []Shard
	if err = s.DB.WithContext(ctx).Where("id <> ? AND status = ?", id, StatusRunning).Find(&candidates).Error; err != nil {
		return PreflightResponse{}, err
	}
	counts, err := assignmentCounts(s.DB.WithContext(ctx))
	if err != nil {
		return PreflightResponse{}, err
	}
	for _, item := range candidates {
		if item.Capacity-int(counts[item.ID]) >= len(detail.Impact) && (result.TargetShardID == "" || rendezvous(id, item.ID) > rendezvous(id, result.TargetShardID)) {
			result.TargetShardID, result.TargetShardName = item.ID, item.Name
		}
	}
	result.Checks = append(result.Checks, Check{Name: "目标容量", Passed: result.TargetShardID != "", Detail: map[bool]string{true: "存在可容纳全部受影响路径的健康目标分片", false: "没有可容纳全部受影响路径的健康目标分片"}[result.TargetShardID != ""]})
	result.Checks = append(result.Checks, Check{Name: "实施授权", Passed: false, Detail: "本工单只提供只读预检；实际迁移必须另建高风险工单并人工确认"})
	result.Ready = result.Checks[0].Passed && result.Checks[1].Passed
	return result, nil
}

func (s *Service) impact(ctx context.Context, id string) ([]Impact, error) {
	var rows []Impact
	err := s.DB.WithContext(ctx).Table("sense_media_shard_assignments AS a").Select("r.id AS route_id, r.device_id, COALESCE(d.name, '') AS device_name, COALESCE(d.location, '') AS location, r.profile_token, r.path, r.desired, r.actual").Joins("JOIN sense_media_routes r ON r.id = a.route_id").Joins("LEFT JOIN sense_devices d ON d.id = r.device_id").Where("a.shard_id = ?", id).Order("d.name ASC, r.profile_token ASC").Scan(&rows).Error
	return rows, err
}

func (s *Service) response(shard Shard, assigned int) Response {
	remaining := shard.Capacity - assigned
	if remaining < 0 {
		remaining = 0
	}
	stale := shard.LastProbeAt == nil || s.now().Sub(shard.LastProbeAt.UTC()) > staleAfter
	return Response{ID: shard.ID, Name: shard.Name, Mode: shard.Mode, Status: shard.Status, Detail: shard.Detail, Capacity: shard.Capacity, AssignedPaths: assigned, Remaining: remaining, LastProbeAt: shard.LastProbeAt, Stale: stale}
}

func assignmentCounts(db *gorm.DB) (map[string]int64, error) {
	type row struct {
		ShardID string
		Count   int64
	}
	var rows []row
	err := db.Model(&Assignment{}).Select("shard_id, COUNT(*) AS count").Group("shard_id").Scan(&rows).Error
	out := map[string]int64{}
	for _, r := range rows {
		out[r.ShardID] = r.Count
	}
	return out, err
}
func rendezvous(routeID, shardID string) uint64 {
	sum := sha256.Sum256([]byte(routeID + "\x00" + shardID))
	return binary.BigEndian.Uint64(sum[:8])
}
func (s *Service) now() time.Time {
	if s.Now != nil {
		return s.Now().UTC()
	}
	return time.Now().UTC()
}

func (s *Service) ControlAPI(ctx context.Context, id string) (string, error) {
	var shard Shard
	if err := s.DB.WithContext(ctx).First(&shard, "id = ?", id).Error; err != nil {
		return "", fmt.Errorf("get media shard control endpoint: %w", err)
	}
	return shard.ControlAPI, nil
}