85 lines
2.1 KiB
Go
85 lines
2.1 KiB
Go
// Filename: internal/domain/upstream/service.go
|
||
package upstream
|
||
|
||
import (
|
||
"errors"
|
||
"gemini-balancer/internal/models"
|
||
"math/rand"
|
||
"time"
|
||
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
type Service struct {
|
||
db *gorm.DB
|
||
}
|
||
|
||
func NewService(db *gorm.DB) *Service {
|
||
rand.New(rand.NewSource(time.Now().UnixNano()))
|
||
return &Service{db: db}
|
||
}
|
||
|
||
func (s *Service) SelectActiveWeighted(upstreams []*models.UpstreamEndpoint) (*models.UpstreamEndpoint, error) {
|
||
activeUpstreams := make([]*models.UpstreamEndpoint, 0)
|
||
totalWeight := 0
|
||
for _, u := range upstreams {
|
||
if u.Status == "active" {
|
||
activeUpstreams = append(activeUpstreams, u)
|
||
totalWeight += u.Weight
|
||
}
|
||
}
|
||
if len(activeUpstreams) == 0 {
|
||
return nil, errors.New("no active upstream endpoints available")
|
||
}
|
||
if totalWeight <= 0 || len(activeUpstreams) == 1 {
|
||
return activeUpstreams[0], nil
|
||
}
|
||
randomWeight := rand.Intn(totalWeight)
|
||
for _, u := range activeUpstreams {
|
||
randomWeight -= u.Weight
|
||
if randomWeight < 0 {
|
||
return u, nil
|
||
}
|
||
}
|
||
return activeUpstreams[len(activeUpstreams)-1], nil
|
||
}
|
||
|
||
// CRUD，供Handler调用
|
||
|
||
func (s *Service) Create(upstream *models.UpstreamEndpoint) error {
|
||
if upstream.Weight == 0 {
|
||
upstream.Weight = 100 // 默认权重
|
||
}
|
||
if upstream.Status == "" {
|
||
upstream.Status = "active" // 默认状态
|
||
}
|
||
return s.db.Create(upstream).Error
|
||
}
|
||
|
||
// List Service层只做数据库查询
|
||
func (s *Service) List() ([]models.UpstreamEndpoint, error) {
|
||
var upstreams []models.UpstreamEndpoint
|
||
err := s.db.Find(&upstreams).Error
|
||
return upstreams, err
|
||
}
|
||
|
||
// GetByID Service层只做数据库查询
|
||
func (s *Service) GetByID(id int) (*models.UpstreamEndpoint, error) {
|
||
var upstream models.UpstreamEndpoint
|
||
if err := s.db.First(&upstream, id).Error; err != nil {
|
||
return nil, err
|
||
}
|
||
return &upstream, nil
|
||
}
|
||
|
||
// Update Service层只做数据库更新
|
||
func (s *Service) Update(upstream *models.UpstreamEndpoint) error {
|
||
return s.db.Save(upstream).Error
|
||
}
|
||
|
||
// Delete Service层只做数据库删除
|
||
func (s *Service) Delete(id int) (int64, error) {
|
||
result := s.db.Delete(&models.UpstreamEndpoint{}, id)
|
||
return result.RowsAffected, result.Error
|
||
}
|