diff --git a/config/config.go b/config/config.go index fe8bb83..29fd106 100644 --- a/config/config.go +++ b/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 } diff --git a/core/requests.go b/core/requests.go index 859b6ce..f6551fb 100644 --- a/core/requests.go +++ b/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< 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< 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 } diff --git a/server/token_manager/token_manager.go b/server/token_manager/token_manager.go new file mode 100644 index 0000000..1f4e172 --- /dev/null +++ b/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 +}