Select Retry

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
 }
 

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
+	AccessKey    string `yaml:"access_key"`    // S3访问密钥ID
-	SecretKey      string `yaml:"secret_key"`       // S3秘密访问密钥
+	SecretKey    string `yaml:"secret_key"`    // S3秘密访问密钥
-	VirtualHost    bool   `yaml:"virtual_host"`     // 是否使用虚拟主机模式
+	VirtualHost  bool   `yaml:"virtual_host"`  // 是否使用虚拟主机模式
-	ProxyMode      bool   `yaml:"proxy_mode"`       // 是否使用代理模式（而非重定向）
+	ProxyMode    bool   `yaml:"proxy_mode"`    // 是否使用代理模式（而非重定向）
-	AuthRequired   bool   `yaml:"auth_required"`    // 是否需要认证
+	AuthRequired bool   `yaml:"auth_required"` // 是否需要认证
 }
 
 // DatabaseConfig 数据库配置
 type DatabaseConfig struct {
-	Type            string `yaml:"type"`             // 数据库类型: sqlite, mysql, postgres
+	Type            string `yaml:"type"`              // 数据库类型: sqlite, mysql, postgres
-	DSN             string `yaml:"dsn"`              // 数据源名称
+	DSN             string `yaml:"dsn"`               // 数据源名称
-	MaxOpenConns    int    `yaml:"max_open_conns"`   // 最大打开连接数
+	MaxOpenConns    int    `yaml:"max_open_conns"`    // 最大打开连接数
-	MaxIdleConns    int    `yaml:"max_idle_conns"`   // 最大空闲连接数
+	MaxIdleConns    int    `yaml:"max_idle_conns"`    // 最大空闲连接数
 	ConnMaxLifetime int    `yaml:"conn_max_lifetime"` // 连接最大生命周期（秒）
-	LogLevel        string `yaml:"log_level"`        // 日志级别: silent, error, warn, info
+	LogLevel        string `yaml:"log_level"`         // 日志级别: silent, error, warn, info
-	AutoMigrate     bool   `yaml:"auto_migrate"`     // 是否自动迁移
+	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)
 		}
 	}