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"
@@ -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 {