diff --git a/internal/balancer/balancer.go b/internal/balancer/balancer.go index bf4a4b6..96e2146 100644 --- a/internal/balancer/balancer.go +++ b/internal/balancer/balancer.go @@ -2,6 +2,7 @@ package balancer import ( "fmt" + "time" "github.com/DullJZ/s3-balance/internal/bucket" "github.com/DullJZ/s3-balance/internal/config" @@ -20,7 +21,7 @@ type Balancer struct { // NewBalancer 创建新的负载均衡器 func NewBalancer(manager *bucket.Manager, cfg *config.BalancerConfig) (*Balancer, error) { var strategy Strategy - + // 根据配置创建对应的策略实例 switch cfg.Strategy { case "round-robin": @@ -46,6 +47,34 @@ func NewBalancer(manager *bucket.Manager, cfg *config.BalancerConfig) (*Balancer // 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 { @@ -126,7 +155,7 @@ func (b *Balancer) GetAvailableBuckets() []*bucket.BucketInfo { 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, @@ -137,7 +166,7 @@ func (b *Balancer) GetBucketStats() map[string]BucketStats { Weight: bucket.Config.Weight, } } - + return stats } diff --git a/internal/config/config.go b/internal/config/config.go index 3a0f9e8..f961d50 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -44,7 +44,7 @@ type BucketConfig struct { // BalancerConfig 负载均衡配置 type BalancerConfig struct { - Strategy string `yaml:"strategy"` // 负载均衡策略: "round-robin", "least-space", "weighted", "consistent-hash" + Strategy string `yaml:"strategy"` // 负载均衡策略: "round-robin", "least-space", "weighted", "consistent-hash" HealthCheckPeriod time.Duration `yaml:"health_check_period"` // 健康检查周期 UpdateStatsPeriod time.Duration `yaml:"update_stats_period"` // 统计更新周期 RetryAttempts int `yaml:"retry_attempts"` // 重试次数 @@ -60,22 +60,22 @@ type MetricsConfig struct { // S3APIConfig S3兼容API配置 type S3APIConfig struct { - AccessKey string `yaml:"access_key"` // S3访问密钥ID - SecretKey string `yaml:"secret_key"` // S3秘密访问密钥 - VirtualHost bool `yaml:"virtual_host"` // 是否使用虚拟主机模式 - ProxyMode bool `yaml:"proxy_mode"` // 是否使用代理模式(而非重定向) - AuthRequired bool `yaml:"auth_required"` // 是否需要认证 + AccessKey string `yaml:"access_key"` // S3访问密钥ID + SecretKey string `yaml:"secret_key"` // S3秘密访问密钥 + VirtualHost bool `yaml:"virtual_host"` // 是否使用虚拟主机模式 + ProxyMode bool `yaml:"proxy_mode"` // 是否使用代理模式(而非重定向) + AuthRequired bool `yaml:"auth_required"` // 是否需要认证 } // DatabaseConfig 数据库配置 type DatabaseConfig struct { - Type string `yaml:"type"` // 数据库类型: sqlite, mysql, postgres - DSN string `yaml:"dsn"` // 数据源名称 - MaxOpenConns int `yaml:"max_open_conns"` // 最大打开连接数 - MaxIdleConns int `yaml:"max_idle_conns"` // 最大空闲连接数 + Type string `yaml:"type"` // 数据库类型: sqlite, mysql, postgres + DSN string `yaml:"dsn"` // 数据源名称 + MaxOpenConns int `yaml:"max_open_conns"` // 最大打开连接数 + MaxIdleConns int `yaml:"max_idle_conns"` // 最大空闲连接数 ConnMaxLifetime int `yaml:"conn_max_lifetime"` // 连接最大生命周期(秒) - LogLevel string `yaml:"log_level"` // 日志级别: silent, error, warn, info - AutoMigrate bool `yaml:"auto_migrate"` // 是否自动迁移 + LogLevel string `yaml:"log_level"` // 日志级别: silent, error, warn, info + AutoMigrate bool `yaml:"auto_migrate"` // 是否自动迁移 } // Load 从文件加载配置 @@ -95,7 +95,7 @@ func Load(configPath string) (*Config, error) { // 解析容量大小 for i := range config.Buckets { if err := config.Buckets[i].ParseMaxSize(); err != nil { - return nil, fmt.Errorf("failed to parse max size for bucket %s: %w", + return nil, fmt.Errorf("failed to parse max size for bucket %s: %w", config.Buckets[i].Name, err) } }