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 }