ocdp-go/backend/internal/domain/service/quota_service.go

package service

import (
	"context"
	"github.com/ocdp/cluster-service/internal/domain/entity"
	"github.com/ocdp/cluster-service/internal/domain/repository"
)

// QuotaService 配额领域服务
type QuotaService struct {
	quotaRepo     repository.QuotaRepository
	instanceRepo  repository.InstanceRepository
	workspaceRepo repository.WorkspaceRepository
}

// NewQuotaService 创建配额服务
func NewQuotaService(
	quotaRepo repository.QuotaRepository,
	instanceRepo repository.InstanceRepository,
	workspaceRepo repository.WorkspaceRepository,
) *QuotaService {
	return &QuotaService{
		quotaRepo:     quotaRepo,
		instanceRepo:  instanceRepo,
		workspaceRepo: workspaceRepo,
	}
}

// CheckQuota 检查配额是否足够
func (s *QuotaService) CheckQuota(ctx context.Context, workspaceID string, cpu, gpu, gpuMemory float64) error {
	// 检查 CPU 配额
	if cpu > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceCPU)
		if err != nil {
			return err
		}
		if quota != nil && !quota.CanAllocate(cpu) {
			return entity.ErrQuotaExceeded
		}
	}

	// 检查 GPU 配额
	if gpu > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceGPU)
		if err != nil {
			return err
		}
		if quota != nil && !quota.CanAllocate(gpu) {
			return entity.ErrQuotaExceeded
		}
	}

	// 检查 GPU Memory 配额
	if gpuMemory > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceGPUMemory)
		if err != nil {
			return err
		}
		if quota != nil && !quota.CanAllocate(gpuMemory) {
			return entity.ErrQuotaExceeded
		}
	}

	return nil
}

// AllocateQuota 分配配额（部署实例成功后调用）
func (s *QuotaService) AllocateQuota(ctx context.Context, workspaceID string, cpu, gpu, gpuMemory float64) error {
	// 分配 CPU
	if cpu > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceCPU)
		if err != nil {
			return err
		}
		if quota != nil {
			quota.Allocate(cpu)
			if err := s.quotaRepo.Update(ctx, quota); err != nil {
				return err
			}
		}
	}

	// 分配 GPU
	if gpu > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceGPU)
		if err != nil {
			return err
		}
		if quota != nil {
			quota.Allocate(gpu)
			if err := s.quotaRepo.Update(ctx, quota); err != nil {
				return err
			}
		}
	}

	// 分配 GPU Memory
	if gpuMemory > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceGPUMemory)
		if err != nil {
			return err
		}
		if quota != nil {
			quota.Allocate(gpuMemory)
			if err := s.quotaRepo.Update(ctx, quota); err != nil {
				return err
			}
		}
	}

	return nil
}

// ReleaseQuota 释放配额（删除实例后调用）
func (s *QuotaService) ReleaseQuota(ctx context.Context, workspaceID string, cpu, gpu, gpuMemory float64) error {
	// 释放 CPU
	if cpu > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceCPU)
		if err != nil {
			return err
		}
		if quota != nil {
			quota.Release(cpu)
			if err := s.quotaRepo.Update(ctx, quota); err != nil {
				return err
			}
		}
	}

	// 释放 GPU
	if gpu > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceGPU)
		if err != nil {
			return err
		}
		if quota != nil {
			quota.Release(gpu)
			if err := s.quotaRepo.Update(ctx, quota); err != nil {
				return err
			}
		}
	}

	// 释放 GPU Memory
	if gpuMemory > 0 {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, entity.ResourceGPUMemory)
		if err != nil {
			return err
		}
		if quota != nil {
			quota.Release(gpuMemory)
			if err := s.quotaRepo.Update(ctx, quota); err != nil {
				return err
			}
		}
	}

	return nil
}

// GetQuotaUsage 获取配额使用情况
func (s *QuotaService) GetQuotaUsage(ctx context.Context, workspaceID string) (map[entity.ResourceType]*entity.WorkspaceQuota, error) {
	quotas, err := s.quotaRepo.GetByWorkspace(ctx, workspaceID)
	if err != nil {
		return nil, err
	}

	result := make(map[entity.ResourceType]*entity.WorkspaceQuota)
	for _, q := range quotas {
		result[q.ResourceType] = q
	}

	// 确保所有资源类型都有返回值
	for _, rt := range []entity.ResourceType{entity.ResourceCPU, entity.ResourceGPU, entity.ResourceGPUMemory} {
		if _, ok := result[rt]; !ok {
			result[rt] = &entity.WorkspaceQuota{
				WorkspaceID:  workspaceID,
				ResourceType: rt,
				HardLimit:    0,
				SoftLimit:    0,
				Used:         0,
			}
		}
	}

	return result, nil
}

// RecalculateQuota 重新计算配额使用量（从实例汇总）
func (s *QuotaService) RecalculateQuota(ctx context.Context, workspaceID string) error {
	// 获取 workspace 的所有实例
	instances, err := s.instanceRepo.GetByWorkspace(ctx, workspaceID)
	if err != nil {
		return err
	}

	// 汇总资源使用
	var totalCPU, totalGPU, totalGPUMemory float64
	for _, inst := range instances {
		totalCPU += inst.CPURequested
		totalGPU += inst.GPURequested
		// GPU Memory 需要解析字符串
		// 这里简化处理，实际需要解析 "16Gi" 这样的格式
	}

	// 更新配额
	resources := []entity.ResourceType{entity.ResourceCPU, entity.ResourceGPU, entity.ResourceGPUMemory}
	values := []float64{totalCPU, totalGPU, totalGPUMemory}

	for i, rt := range resources {
		quota, err := s.quotaRepo.GetByWorkspaceAndType(ctx, workspaceID, rt)
		if err != nil {
			return err
		}
		if quota != nil {
			quota.Used = values[i]
			if err := s.quotaRepo.Update(ctx, quota); err != nil {
				return err
			}
		}
	}

	return nil
}