s3-balance/internal/balancer/balancer.go

package balancer

import (
	"fmt"
	"time"

	"github.com/DullJZ/s3-balance/internal/bucket"
	"github.com/DullJZ/s3-balance/internal/config"
	"github.com/DullJZ/s3-balance/internal/metrics"
)

// Balancer 负载均衡器
// 负责根据配置的策略选择合适的存储桶
type Balancer struct {
	manager  *bucket.Manager
	strategy Strategy
	config   *config.BalancerConfig
	metrics  *metrics.Metrics
}

// NewBalancer 创建新的负载均衡器
func NewBalancer(manager *bucket.Manager, cfg *config.BalancerConfig) (*Balancer, error) {
	var strategy Strategy

	// 根据配置创建对应的策略实例
	switch cfg.Strategy {
	case "round-robin":
		strategy = NewRoundRobinStrategy()
	case "least-space":
		strategy = NewLeastSpaceStrategy()
	case "weighted":
		strategy = NewWeightedStrategy()
	default:
		return nil, fmt.Errorf("unknown balancer strategy: %s", cfg.Strategy)
	}

	return &Balancer{
		manager:  manager,
		strategy: strategy,
		config:   cfg,
		metrics:  nil, // 将在main.go中设置
	}, nil
}

// SelectBucket 选择一个存储桶
// 首先过滤出有足够空间的存储桶，然后使用策略选择
func (b *Balancer) SelectBucket(key string, size int64) (*bucket.BucketInfo, error) {
	attempts := 1
	delay := time.Second

	if b.config != nil {
		if b.config.RetryAttempts > 0 {
			attempts = b.config.RetryAttempts
		}
		if b.config.RetryDelay > 0 {
			delay = b.config.RetryDelay
		}
	}

	var lastErr error
	for i := 0; i < attempts; i++ {
		selected, err := b.selectOnce(key, size)
		if err == nil {
			return selected, nil
		}
		lastErr = err
		if i < attempts-1 {
			time.Sleep(delay)
		}
	}

	return nil, lastErr
}

func (b *Balancer) selectOnce(key string, size int64) (*bucket.BucketInfo, error) {
	// 获取所有可用的存储桶
	buckets := b.manager.GetAvailableBuckets()
	if len(buckets) == 0 {
		return nil, fmt.Errorf("no available buckets")
	}

	// 过滤出有足够空间的存储桶
	var availableBuckets []*bucket.BucketInfo
	for _, bucket := range buckets {
		if bucket.GetAvailableSpace() >= size {
			availableBuckets = append(availableBuckets, bucket)
		}
	}

	if len(availableBuckets) == 0 {
		return nil, fmt.Errorf("no bucket has enough space for %d bytes", size)
	}

	// 使用策略选择存储桶
	selected, err := b.strategy.SelectBucket(availableBuckets, key, size)
	if err != nil {
		return nil, err
	}

	// 记录指标
	if b.metrics != nil && selected != nil {
		b.metrics.RecordBalancerDecision(b.strategy.Name(), selected.Config.Name)
	}

	return selected, nil
}

// GetStrategy 获取当前策略名称
func (b *Balancer) GetStrategy() string {
	return b.strategy.Name()
}

// SetStrategy 动态切换策略
// 允许在运行时更改负载均衡策略
func (b *Balancer) SetStrategy(strategyName string) error {
	var strategy Strategy

	switch strategyName {
	case "round-robin":
		strategy = NewRoundRobinStrategy()
	case "least-space":
		strategy = NewLeastSpaceStrategy()
	case "weighted":
		strategy = NewWeightedStrategy()
	default:
		return fmt.Errorf("unknown balancer strategy: %s", strategyName)
	}

	b.strategy = strategy
	return nil
}

// UpdateStrategy 更新负载均衡策略（热更新用）
func (b *Balancer) UpdateStrategy(strategyName string) error {
	return b.SetStrategy(strategyName)
}

// SetMetrics 设置指标服务
func (b *Balancer) SetMetrics(metrics *metrics.Metrics) {
	b.metrics = metrics
}

// GetAvailableBuckets 获取所有可用的存储桶
// 方便外部直接查询可用存储桶列表
func (b *Balancer) GetAvailableBuckets() []*bucket.BucketInfo {
	return b.manager.GetAvailableBuckets()
}

// GetBucketStats 获取存储桶统计信息
// 返回每个存储桶的使用情况
func (b *Balancer) GetBucketStats() map[string]BucketStats {
	stats := make(map[string]BucketStats)
	buckets := b.manager.GetAllBuckets()

	for _, bucket := range buckets {
		stats[bucket.Config.Name] = BucketStats{
			Name:           bucket.Config.Name,
			TotalSpace:     bucket.Config.MaxSizeBytes,
			UsedSpace:      bucket.GetUsedSize(),
			AvailableSpace: bucket.GetAvailableSpace(),
			IsAvailable:    bucket.IsAvailable(),
			Weight:         bucket.Config.Weight,
		}
	}

	return stats
}

// BucketStats 存储桶统计信息
type BucketStats struct {
	Name           string `json:"name"`
	TotalSpace     int64  `json:"total_space"`
	UsedSpace      int64  `json:"used_space"`
	AvailableSpace int64  `json:"available_space"`
	IsAvailable    bool   `json:"is_available"`
	Weight         int    `json:"weight"`
}