Implement token management and request retry system

Introduce a new TokenManager to handle token rotation, status
tracking, and automatic retries with exponential backoff for API
requests. This change aims to improve the reliability API
interactions by intelligently managing tokens and gracefully
handling request failures.

Update API request functions to utilise the new TokenManager. This
commit also adds configuration options for max retries, base timeout,
and max backoff time to fine-tune the retry behavior.

config/config.go

@@ -6,12 +6,11 @@ import (
 	"context"
 	"errors"
 	"log"
-	"math/rand"
 	"net/url"
 	"strings"
-	"sync/atomic"
 	"time"
 
+	"codeberg.org/vnpower/pixivfe/v2/server/token_manager"
 	"github.com/sethvargo/go-envconfig"
 )
 
@@ -27,19 +26,21 @@ type ServerConfig struct {
 	Port       string `env:"PIXIVFE_PORT"`
 	UnixSocket string `env:"PIXIVFE_UNIXSOCKET"`
 
-	Token          []string `env:"PIXIVFE_TOKEN,required"`                 // may be multiple tokens. delimiter is ','
-	LoadBalancing  string   `env:"PIXIVFE_TOKEN_LOAD_BALANCING,overwrite"` // 'round-robin' or 'random'
-	InDevelopment  bool     `env:"PIXIVFE_DEV"`
-	UserAgent      string   `env:"PIXIVFE_USERAGENT,overwrite"`
-	AcceptLanguage string   `env:"PIXIVFE_ACCEPTLANGUAGE,overwrite"`
-	RequestLimit   int      `env:"PIXIVFE_REQUESTLIMIT"` // if 0, request limit is disabled
+	Token              []string `env:"PIXIVFE_TOKEN,required"` // may be multiple tokens. delimiter is ','
+	TokenManager       *token_manager.TokenManager
+	TokenLoadBalancing string        `env:"PIXIVFE_TOKEN_LOAD_BALANCING,overwrite"`
+	MaxRetries         int           `env:"PIXIVFE_TOKEN_MAX_RETRIES,overwrite"`
+	BaseTimeout        time.Duration `env:"PIXIVFE_TOKEN_BASE_TIMEOUT,overwrite"`
+	MaxBackoffTime     time.Duration `env:"PIXIVFE_TOKEN_MAX_BACKOFF_TIME,overwrite"`
+	InDevelopment      bool          `env:"PIXIVFE_DEV"`
+	UserAgent          string        `env:"PIXIVFE_USERAGENT,overwrite"`
+	AcceptLanguage     string        `env:"PIXIVFE_ACCEPTLANGUAGE,overwrite"`
+	RequestLimit       int           `env:"PIXIVFE_REQUESTLIMIT"` // if 0, request limit is disabled
 
 	ProxyServer_staging string  `env:"PIXIVFE_IMAGEPROXY,overwrite"`
 	ProxyServer         url.URL // proxy server URL, may or may not contain authority part of the URL
 
 	ProxyCheckInterval time.Duration `env:"PIXIVFE_PROXY_CHECK_INTERVAL,overwrite"`
-
-	tokenIndex uint32 // Used for round-robin token selection
 }
 
 func (s *ServerConfig) LoadConfig() error {
@@ -53,7 +54,10 @@ func (s *ServerConfig) LoadConfig() error {
 	s.AcceptLanguage = "en-US,en;q=0.5"
 	s.ProxyServer_staging = BuiltinProxyUrl
 	s.ProxyCheckInterval = 8 * time.Hour
-	s.LoadBalancing = "round-robin"
+	s.TokenLoadBalancing = "round-robin"
+	s.MaxRetries = 5
+	s.BaseTimeout = 1 * time.Second
+	s.MaxBackoffTime = 32 * time.Second
 
 	// load config from from env vars
 	if err := envconfig.Process(context.Background(), s); err != nil {
@@ -86,25 +90,30 @@ func (s *ServerConfig) LoadConfig() error {
 	}
 	log.Printf("Proxy server set to: %s\n", s.ProxyServer.String())
 	log.Printf("Proxy check interval set to: %v\n", s.ProxyCheckInterval)
-	log.Printf("Token load balancing method: %s\n", s.LoadBalancing)
+
+	// Validate TokenLoadBalancing
+	switch s.TokenLoadBalancing {
+	case "round-robin", "random", "least-recently-used":
+		// Valid options
+	default:
+		log.Printf("Invalid PIXIVFE_TOKEN_LOAD_BALANCING value: %s. Defaulting to 'round-robin'.\n", s.TokenLoadBalancing)
+		s.TokenLoadBalancing = "round-robin"
+	}
+
+	// Initialize TokenManager
+	s.TokenManager = token_manager.NewTokenManager(s.Token, s.MaxRetries, s.BaseTimeout, s.MaxBackoffTime, s.TokenLoadBalancing)
+	log.Printf("Token manager initialized with %d tokens\n", len(s.Token))
+	log.Printf("Max retries: %d, Base timeout: %v, Max backoff time: %v\n", s.MaxRetries, s.BaseTimeout, s.MaxBackoffTime)
+	log.Printf("Token load balancing method: %s\n", s.TokenLoadBalancing)
 
 	return nil
 }
 
 func (s *ServerConfig) GetToken() string {
-	switch s.LoadBalancing {
-	case "random":
-		return s.getRandomToken()
-	default:
-		return s.getRoundRobinToken()
+	token := s.TokenManager.GetToken()
+	if token == nil {
+		log.Println("[WARNING] All tokens are timed out. Using the first available token.")
+		return s.Token[0]
 	}
-}
-
-func (s *ServerConfig) getRandomToken() string {
-	return s.Token[rand.Intn(len(s.Token))]
-}
-
-func (s *ServerConfig) getRoundRobinToken() string {
-	index := atomic.AddUint32(&s.tokenIndex, 1) % uint32(len(s.Token))
-	return s.Token[index]
+	return token.Value
 }

core/requests.go

@@ -12,6 +12,7 @@ import (
 	config "codeberg.org/vnpower/pixivfe/v2/config"
 	"codeberg.org/vnpower/pixivfe/v2/server/audit"
 	"codeberg.org/vnpower/pixivfe/v2/server/request_context"
+	"codeberg.org/vnpower/pixivfe/v2/server/token_manager"
 	"codeberg.org/vnpower/pixivfe/v2/server/utils"
 	"github.com/tidwall/gjson"
 )
@@ -21,24 +22,66 @@ type SimpleHTTPResponse struct {
 	Body       string
 }
 
-// send GET
-func API_GET(context context.Context, url string, token string) (SimpleHTTPResponse, error) {
-	start_time := time.Now()
-	res, resp, err := _API_GET(context, url, token)
-	end_time := time.Now()
-	audit.LogAPIRoundTrip(audit.APIRequestSpan{
-		RequestId: request_context.GetFromContext(context).RequestId,
-		Response:  resp, Error: err, Method: "GET", Url: url, Token: token, Body: res.Body, StartTime: start_time, EndTime: end_time})
-	if err != nil {
-		return SimpleHTTPResponse{}, fmt.Errorf("While GET %s: %w", url, err)
+func retryRequest(ctx context.Context, reqFunc func(context.Context, string) (SimpleHTTPResponse, *http.Response, error)) (SimpleHTTPResponse, error) {
+	var lastErr error
+	tokenManager := config.GlobalConfig.TokenManager
+
+	for i := 0; i < tokenManager.GetMaxRetries(); i++ {
+		token := tokenManager.GetToken()
+		if token == nil {
+			return SimpleHTTPResponse{}, errors.New("All tokens are timed out")
+		}
+
+		res, resp, err := reqFunc(ctx, token.Value)
+		if err == nil && res.StatusCode == http.StatusOK {
+			tokenManager.MarkTokenStatus(token, token_manager.Good)
+			return res, nil
+		}
+
+		lastErr = err
+		if err == nil {
+			lastErr = fmt.Errorf("HTTP status code: %d", res.StatusCode)
+		}
+
+		tokenManager.MarkTokenStatus(token, token_manager.TimedOut)
+
+		backoffDuration := tokenManager.GetBaseTimeout() * time.Duration(1<<uint(i))
+		if backoffDuration > tokenManager.GetMaxBackoffTime() {
+			backoffDuration = tokenManager.GetMaxBackoffTime()
+		}
+
+		select {
+		case <-ctx.Done():
+			return SimpleHTTPResponse{}, ctx.Err()
+		case <-time.After(backoffDuration):
+		}
+
+		audit.LogAPIRoundTrip(audit.APIRequestSpan{
+			RequestId: request_context.GetFromContext(ctx).RequestId,
+			Response:  resp,
+			Error:     err,
+			Method:    "GET",
+			Token:     token.Value,
+			Body:      res.Body,
+			StartTime: time.Now(),
+			EndTime:   time.Now().Add(backoffDuration),
+		})
 	}
-	return res, nil
+
+	return SimpleHTTPResponse{}, fmt.Errorf("Max retries reached. Last error: %v", lastErr)
 }
 
-func _API_GET(context context.Context, url string, token string) (SimpleHTTPResponse, *http.Response, error) {
+// send GET
+func API_GET(ctx context.Context, url string, _ string) (SimpleHTTPResponse, error) {
+	return retryRequest(ctx, func(ctx context.Context, token string) (SimpleHTTPResponse, *http.Response, error) {
+		return _API_GET(ctx, url, token)
+	})
+}
+
+func _API_GET(ctx context.Context, url string, token string) (SimpleHTTPResponse, *http.Response, error) {
 	var res SimpleHTTPResponse
 
-	req, err := http.NewRequestWithContext(context, "GET", url, nil)
+	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
 	if err != nil {
 		return res, nil, err
 	}
@@ -46,17 +89,10 @@ func _API_GET(context context.Context, url string, token string) (SimpleHTTPResp
 	req.Header.Add("User-Agent", config.GlobalConfig.UserAgent)
 	req.Header.Add("Accept-Language", config.GlobalConfig.AcceptLanguage)
 
-	if token == "" {
-		req.AddCookie(&http.Cookie{
-			Name:  "PHPSESSID",
-			Value: config.GlobalConfig.GetToken(),
-		})
-	} else {
-		req.AddCookie(&http.Cookie{
-			Name:  "PHPSESSID",
-			Value: token,
-		})
-	}
+	req.AddCookie(&http.Cookie{
+		Name:  "PHPSESSID",
+		Value: token,
+	})
 
 	// Make the request
 	resp, err := utils.HttpClient.Do(req)
@@ -77,8 +113,8 @@ func _API_GET(context context.Context, url string, token string) (SimpleHTTPResp
 	return res, resp, nil
 }
 
-func API_GET_UnwrapJson(context context.Context, url string, token string) (string, error) {
-	resp, err := API_GET(context, url, token)
+func API_GET_UnwrapJson(ctx context.Context, url string, _ string) (string, error) {
+	resp, err := API_GET(ctx, url, "")
 	if err != nil {
 		return "", err
 	}
@@ -101,23 +137,63 @@ func API_GET_UnwrapJson(context context.Context, url string, token string) (stri
 }
 
 // send POST
-func API_POST(context context.Context, url, payload, token, csrf string, isJSON bool) error {
-	start_time := time.Now()
-	resp, err := _API_POST(context, url, payload, token, csrf, isJSON)
-	end_time := time.Now()
-	audit.LogAPIRoundTrip(audit.APIRequestSpan{
-		RequestId: request_context.GetFromContext(context).RequestId,
-		Response:  resp, Error: err, Method: "POST", Url: url, Token: token, Body: "", StartTime: start_time, EndTime: end_time})
-	if err != nil {
-		return fmt.Errorf("While POST %s: %w", url, err)
+func API_POST(ctx context.Context, url, payload, _, csrf string, isJSON bool) error {
+	tokenManager := config.GlobalConfig.TokenManager
+
+	var lastErr error
+	for i := 0; i < tokenManager.GetMaxRetries(); i++ {
+		token := tokenManager.GetToken()
+		if token == nil {
+			return errors.New("All tokens are timed out")
+		}
+
+		start_time := time.Now()
+		resp, err := _API_POST(ctx, url, payload, token.Value, csrf, isJSON)
+		end_time := time.Now()
+
+		audit.LogAPIRoundTrip(audit.APIRequestSpan{
+			RequestId: request_context.GetFromContext(ctx).RequestId,
+			Response:  resp,
+			Error:     err,
+			Method:    "POST",
+			Url:       url,
+			Token:     token.Value,
+			Body:      "",
+			StartTime: start_time,
+			EndTime:   end_time,
+		})
+
+		if err == nil && resp.StatusCode == http.StatusOK {
+			tokenManager.MarkTokenStatus(token, token_manager.Good)
+			return nil
+		}
+
+		lastErr = err
+		if err == nil {
+			lastErr = fmt.Errorf("HTTP status code: %d", resp.StatusCode)
+		}
+
+		tokenManager.MarkTokenStatus(token, token_manager.TimedOut)
+
+		backoffDuration := tokenManager.GetBaseTimeout() * time.Duration(1<<uint(i))
+		if backoffDuration > tokenManager.GetMaxBackoffTime() {
+			backoffDuration = tokenManager.GetMaxBackoffTime()
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(backoffDuration):
+		}
 	}
-	return err
+
+	return fmt.Errorf("Max retries reached. Last error: %v", lastErr)
 }
 
-func _API_POST(context context.Context, url, payload, token, csrf string, isJSON bool) (*http.Response, error) {
+func _API_POST(ctx context.Context, url, payload, token, csrf string, isJSON bool) (*http.Response, error) {
 	requestBody := []byte(payload)
 
-	req, err := http.NewRequestWithContext(context, "POST", url, bytes.NewBuffer(requestBody))
+	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(requestBody))
 	if err != nil {
 		return nil, err
 	}

server/token_manager/token_manager.go

@@ -0,0 +1,169 @@
+package token_manager
+
+import (
+	"math"
+	"math/rand"
+	"sort"
+	"sync"
+	"time"
+)
+
+type TokenStatus int
+
+const (
+	Good TokenStatus = iota
+	TimedOut
+)
+
+type Token struct {
+	Value               string
+	Status              TokenStatus
+	TimeoutUntil        time.Time
+	FailureCount        int
+	LastUsed            time.Time
+	BaseTimeoutDuration time.Duration
+}
+
+type TokenManager struct {
+	tokens              []*Token
+	mu                  sync.Mutex
+	maxRetries          int
+	baseTimeout         time.Duration
+	maxBackoffTime      time.Duration
+	loadBalancingMethod string
+	currentIndex        int
+}
+
+func NewTokenManager(tokenValues []string, maxRetries int, baseTimeout, maxBackoffTime time.Duration, loadBalancingMethod string) *TokenManager {
+	tokens := make([]*Token, len(tokenValues))
+	for i, value := range tokenValues {
+		tokens[i] = &Token{
+			Value:               value,
+			Status:              Good,
+			BaseTimeoutDuration: time.Second,
+		}
+	}
+	return &TokenManager{
+		tokens:              tokens,
+		maxRetries:          maxRetries,
+		baseTimeout:         baseTimeout,
+		maxBackoffTime:      maxBackoffTime,
+		loadBalancingMethod: loadBalancingMethod,
+		currentIndex:        0,
+	}
+}
+
+func (tm *TokenManager) GetToken() *Token {
+	tm.mu.Lock()
+	defer tm.mu.Unlock()
+
+	now := time.Now()
+	healthyTokens := tm.getHealthyTokens()
+
+	if len(healthyTokens) == 0 {
+		return tm.getFallbackToken(now)
+	}
+
+	var selectedToken *Token
+	switch tm.loadBalancingMethod {
+	case "round-robin":
+		selectedToken = tm.roundRobinSelection(healthyTokens)
+	case "random":
+		selectedToken = tm.randomSelection(healthyTokens)
+	case "least-recently-used":
+		selectedToken = tm.leastRecentlyUsedSelection(healthyTokens)
+	default:
+		selectedToken = tm.roundRobinSelection(healthyTokens)
+	}
+
+	selectedToken.LastUsed = now
+	return selectedToken
+}
+
+func (tm *TokenManager) getHealthyTokens() []*Token {
+	healthyTokens := make([]*Token, 0)
+	for _, token := range tm.tokens {
+		if token.Status == Good {
+			healthyTokens = append(healthyTokens, token)
+		}
+	}
+	return healthyTokens
+}
+
+func (tm *TokenManager) getFallbackToken(now time.Time) *Token {
+	var bestToken *Token
+	for _, token := range tm.tokens {
+		if token.Status == TimedOut && (bestToken == nil || token.TimeoutUntil.Before(bestToken.TimeoutUntil)) {
+			bestToken = token
+		}
+	}
+
+	if bestToken != nil && now.After(bestToken.TimeoutUntil) {
+		bestToken.Status = Good
+		bestToken.LastUsed = now
+		return bestToken
+	}
+
+	return bestToken
+}
+
+func (tm *TokenManager) roundRobinSelection(healthyTokens []*Token) *Token {
+	if tm.currentIndex >= len(healthyTokens) {
+		tm.currentIndex = 0
+	}
+	selectedToken := healthyTokens[tm.currentIndex]
+	tm.currentIndex++
+	return selectedToken
+}
+
+func (tm *TokenManager) randomSelection(healthyTokens []*Token) *Token {
+	return healthyTokens[rand.Intn(len(healthyTokens))]
+}
+
+func (tm *TokenManager) leastRecentlyUsedSelection(healthyTokens []*Token) *Token {
+	sort.Slice(healthyTokens, func(i, j int) bool {
+		return healthyTokens[i].LastUsed.Before(healthyTokens[j].LastUsed)
+	})
+	return healthyTokens[0]
+}
+
+func (tm *TokenManager) MarkTokenStatus(token *Token, status TokenStatus) {
+	tm.mu.Lock()
+	defer tm.mu.Unlock()
+
+	token.Status = status
+	if status == TimedOut {
+		token.FailureCount++
+		timeoutDuration := time.Duration(math.Min(
+			float64(token.BaseTimeoutDuration)*math.Pow(2, float64(token.FailureCount-1)),
+			float64(tm.maxBackoffTime),
+		))
+		token.TimeoutUntil = time.Now().Add(timeoutDuration)
+	} else {
+		token.FailureCount = 0
+		token.BaseTimeoutDuration = time.Second
+	}
+}
+
+func (tm *TokenManager) ResetAllTokens() {
+	tm.mu.Lock()
+	defer tm.mu.Unlock()
+
+	for _, token := range tm.tokens {
+		token.Status = Good
+		token.FailureCount = 0
+		token.BaseTimeoutDuration = time.Second
+	}
+}
+
+func (tm *TokenManager) GetMaxRetries() int {
+	return tm.maxRetries
+}
+
+func (tm *TokenManager) GetBaseTimeout() time.Duration {
+	return tm.baseTimeout
+}
+
+func (tm *TokenManager) GetMaxBackoffTime() time.Duration {
+	return tm.maxBackoffTime
+}