feat: 代理调试版本
This commit is contained in:
tbphp
@@ -2,171 +2,311 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gpt-load/internal/channel"
|
||||
"gpt-load/internal/config"
|
||||
app_errors "gpt-load/internal/errors"
|
||||
"gpt-load/internal/keypool"
|
||||
"gpt-load/internal/models"
|
||||
"gpt-load/internal/response"
|
||||
"io"
|
||||
"time"
|
||||
"gpt-load/internal/services"
|
||||
"gpt-load/internal/types"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// A list of errors that are considered normal during streaming when a client disconnects.
|
||||
var ignorableStreamErrors = []string{
|
||||
"context canceled",
|
||||
"connection reset by peer",
|
||||
"broken pipe",
|
||||
"use of closed network connection",
|
||||
}
|
||||
|
||||
// isIgnorableStreamError checks if the error is a common, non-critical error that can occur
|
||||
// when a client disconnects during a streaming response.
|
||||
func isIgnorableStreamError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
errStr := err.Error()
|
||||
for _, ignorableError := range ignorableStreamErrors {
|
||||
if strings.Contains(errStr, ignorableError) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ProxyServer represents the proxy server
|
||||
type ProxyServer struct {
|
||||
DB *gorm.DB
|
||||
channelFactory *channel.Factory
|
||||
keyProvider *keypool.KeyProvider
|
||||
requestLogChan chan models.RequestLog
|
||||
keyProvider *keypool.KeyProvider
|
||||
groupManager *services.GroupManager
|
||||
settingsManager *config.SystemSettingsManager
|
||||
channelFactory *channel.Factory
|
||||
}
|
||||
|
||||
// NewProxyServer creates a new proxy server
|
||||
func NewProxyServer(
|
||||
db *gorm.DB,
|
||||
channelFactory *channel.Factory,
|
||||
keyProvider *keypool.KeyProvider,
|
||||
requestLogChan chan models.RequestLog,
|
||||
groupManager *services.GroupManager,
|
||||
settingsManager *config.SystemSettingsManager,
|
||||
channelFactory *channel.Factory,
|
||||
) (*ProxyServer, error) {
|
||||
return &ProxyServer{
|
||||
DB: db,
|
||||
channelFactory: channelFactory,
|
||||
keyProvider: keyProvider,
|
||||
requestLogChan: requestLogChan,
|
||||
keyProvider: keyProvider,
|
||||
groupManager: groupManager,
|
||||
settingsManager: settingsManager,
|
||||
channelFactory: channelFactory,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// HandleProxy handles the main proxy logic
|
||||
// HandleProxy is the main entry point for proxy requests, refactored based on the stable .bak logic.
|
||||
func (ps *ProxyServer) HandleProxy(c *gin.Context) {
|
||||
startTime := time.Now()
|
||||
groupName := c.Param("group_name")
|
||||
|
||||
// 1. Find the group by name (without preloading keys)
|
||||
var group models.Group
|
||||
if err := ps.DB.Where("name = ?", groupName).First(&group).Error; err != nil {
|
||||
group, err := ps.groupManager.GetGroupByName(groupName)
|
||||
if err != nil {
|
||||
response.Error(c, app_errors.ParseDBError(err))
|
||||
return
|
||||
}
|
||||
|
||||
// 2. Select an available API key from the KeyPool
|
||||
apiKey, err := ps.keyProvider.SelectKey(group.ID)
|
||||
if err != nil {
|
||||
// Properly handle the case where no keys are available
|
||||
if apiErr, ok := err.(*app_errors.APIError); ok {
|
||||
response.Error(c, apiErr)
|
||||
} else {
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrInternalServer, err.Error()))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// 3. Get the appropriate channel handler from the factory
|
||||
channelHandler, err := ps.channelFactory.GetChannel(&group)
|
||||
channelHandler, err := ps.channelFactory.GetChannel(group)
|
||||
if err != nil {
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrInternalServer, fmt.Sprintf("Failed to get channel for group '%s': %v", groupName, err)))
|
||||
return
|
||||
}
|
||||
|
||||
// 4. Apply parameter overrides if they exist
|
||||
if len(group.ParamOverrides) > 0 {
|
||||
err := ps.applyParamOverrides(c, &group)
|
||||
if err != nil {
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrInternalServer, fmt.Sprintf("Failed to apply parameter overrides: %v", err)))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Forward the request using the channel handler
|
||||
err = channelHandler.Handle(c, apiKey, &group)
|
||||
|
||||
// 6. Update key status and log the request asynchronously
|
||||
isSuccess := err == nil
|
||||
ps.keyProvider.UpdateStatus(apiKey.ID, group.ID, isSuccess)
|
||||
|
||||
if !isSuccess {
|
||||
logrus.WithFields(logrus.Fields{
|
||||
"group": group.Name,
|
||||
"key_id": apiKey.ID,
|
||||
"error": err.Error(),
|
||||
}).Error("Channel handler failed")
|
||||
}
|
||||
go ps.logRequest(c, &group, apiKey, startTime)
|
||||
}
|
||||
|
||||
func (ps *ProxyServer) logRequest(c *gin.Context, group *models.Group, key *models.APIKey, startTime time.Time) {
|
||||
logEntry := models.RequestLog{
|
||||
ID: fmt.Sprintf("req_%d", time.Now().UnixNano()),
|
||||
Timestamp: startTime,
|
||||
GroupID: group.ID,
|
||||
KeyID: key.ID,
|
||||
SourceIP: c.ClientIP(),
|
||||
StatusCode: c.Writer.Status(),
|
||||
RequestPath: c.Request.URL.Path,
|
||||
RequestBodySnippet: "", // Can be implemented later if needed
|
||||
}
|
||||
|
||||
// Send to the logging channel without blocking
|
||||
select {
|
||||
case ps.requestLogChan <- logEntry:
|
||||
default:
|
||||
logrus.Warn("Request log channel is full. Dropping log entry.")
|
||||
}
|
||||
}
|
||||
|
||||
// Close cleans up resources
|
||||
func (ps *ProxyServer) Close() {
|
||||
close(ps.requestLogChan)
|
||||
}
|
||||
|
||||
func (ps *ProxyServer) applyParamOverrides(c *gin.Context, group *models.Group) error {
|
||||
// Read the original request body
|
||||
bodyBytes, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read request body: %w", err)
|
||||
logrus.Errorf("Failed to read request body: %v", err)
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "Failed to read request body"))
|
||||
return
|
||||
}
|
||||
c.Request.Body.Close()
|
||||
|
||||
// If body is empty, nothing to override, just restore the body
|
||||
if len(bodyBytes) == 0 {
|
||||
c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
|
||||
return nil
|
||||
// 4. Apply parameter overrides if any.
|
||||
finalBodyBytes, err := ps.applyParamOverrides(bodyBytes, group)
|
||||
if err != nil {
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrInternalServer, fmt.Sprintf("Failed to apply parameter overrides: %v", err)))
|
||||
return
|
||||
}
|
||||
|
||||
// Save the original Content-Type
|
||||
originalContentType := c.GetHeader("Content-Type")
|
||||
// 5. Determine if this is a streaming request.
|
||||
isStream := channelHandler.IsStreamRequest(c, bodyBytes)
|
||||
|
||||
// 6. Execute the request using the recursive retry logic.
|
||||
ps.executeRequestWithRetry(c, channelHandler, group, finalBodyBytes, isStream, startTime, 0, nil)
|
||||
}
|
||||
|
||||
// executeRequestWithRetry is the core recursive function for handling requests and retries.
|
||||
func (ps *ProxyServer) executeRequestWithRetry(
|
||||
c *gin.Context,
|
||||
channelHandler channel.ChannelProxy,
|
||||
group *models.Group,
|
||||
bodyBytes []byte,
|
||||
isStream bool,
|
||||
startTime time.Time,
|
||||
retryCount int,
|
||||
retryErrors []types.RetryError,
|
||||
) {
|
||||
cfg := group.EffectiveConfig
|
||||
if retryCount > cfg.MaxRetries {
|
||||
logrus.Errorf("Max retries exceeded for group %s after %d attempts.", group.Name, retryCount)
|
||||
if len(retryErrors) > 0 {
|
||||
lastError := retryErrors[len(retryErrors)-1]
|
||||
var errorJSON map[string]any
|
||||
if err := json.Unmarshal([]byte(lastError.ErrorMessage), &errorJSON); err == nil {
|
||||
c.JSON(lastError.StatusCode, errorJSON)
|
||||
} else {
|
||||
response.Error(c, app_errors.NewAPIErrorWithUpstream(lastError.StatusCode, "UPSTREAM_ERROR", lastError.ErrorMessage))
|
||||
}
|
||||
} else {
|
||||
response.Error(c, app_errors.ErrMaxRetriesExceeded)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
apiKey, err := ps.keyProvider.SelectKey(group.ID)
|
||||
if err != nil {
|
||||
logrus.Errorf("Failed to select a key for group %s on attempt %d: %v", group.Name, retryCount+1, err)
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrNoKeysAvailable, err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
upstreamURL, err := channelHandler.BuildUpstreamURL(c.Request.URL, group)
|
||||
if err != nil {
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrInternalServer, fmt.Sprintf("Failed to build upstream URL: %v", err)))
|
||||
return
|
||||
}
|
||||
|
||||
var ctx context.Context
|
||||
var cancel context.CancelFunc
|
||||
if isStream {
|
||||
ctx, cancel = context.WithCancel(c.Request.Context())
|
||||
} else {
|
||||
timeout := time.Duration(cfg.RequestTimeout) * time.Second
|
||||
ctx, cancel = context.WithTimeout(c.Request.Context(), timeout)
|
||||
}
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, c.Request.Method, upstreamURL, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
logrus.Errorf("Failed to create upstream request: %v", err)
|
||||
response.Error(c, app_errors.ErrInternalServer)
|
||||
return
|
||||
}
|
||||
req.ContentLength = int64(len(bodyBytes))
|
||||
|
||||
req.Header = c.Request.Header.Clone()
|
||||
channelHandler.ModifyRequest(req, apiKey, group)
|
||||
|
||||
client := channelHandler.GetHTTPClient()
|
||||
if isStream {
|
||||
client = channelHandler.GetStreamClient()
|
||||
req.Header.Set("X-Accel-Buffering", "no")
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
ps.keyProvider.UpdateStatus(apiKey.ID, group.ID, false)
|
||||
logrus.Warnf("Request failed (attempt %d/%d) for key %s: %v", retryCount+1, cfg.MaxRetries, apiKey.KeyValue[:8], err)
|
||||
|
||||
newRetryErrors := append(retryErrors, types.RetryError{
|
||||
StatusCode: 0,
|
||||
ErrorMessage: err.Error(),
|
||||
KeyID: fmt.Sprintf("%d", apiKey.ID),
|
||||
Attempt: retryCount + 1,
|
||||
})
|
||||
ps.executeRequestWithRetry(c, channelHandler, group, bodyBytes, isStream, startTime, retryCount+1, newRetryErrors)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
ps.keyProvider.UpdateStatus(apiKey.ID, group.ID, false)
|
||||
errorBody, readErr := io.ReadAll(resp.Body)
|
||||
if readErr != nil {
|
||||
logrus.Errorf("Failed to read error body: %v", readErr)
|
||||
// Even if reading fails, we should proceed with retry logic
|
||||
errorBody = []byte("Failed to read error body")
|
||||
}
|
||||
|
||||
// Check for gzip encoding and decompress if necessary.
|
||||
if resp.Header.Get("Content-Encoding") == "gzip" {
|
||||
reader, err := gzip.NewReader(bytes.NewReader(errorBody))
|
||||
if err == nil {
|
||||
decompressedBody, err := io.ReadAll(reader)
|
||||
if err == nil {
|
||||
errorBody = decompressedBody
|
||||
} else {
|
||||
logrus.Warnf("Failed to decompress gzip error body: %v", err)
|
||||
}
|
||||
reader.Close()
|
||||
} else {
|
||||
logrus.Warnf("Failed to create gzip reader for error body: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
logrus.Warnf("Request failed with status %d (attempt %d/%d) for key %s. Body: %s", resp.StatusCode, retryCount+1, cfg.MaxRetries, apiKey.KeyValue[:8], string(errorBody))
|
||||
|
||||
newRetryErrors := append(retryErrors, types.RetryError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ErrorMessage: string(errorBody),
|
||||
KeyID: fmt.Sprintf("%d", apiKey.ID),
|
||||
Attempt: retryCount + 1,
|
||||
})
|
||||
ps.executeRequestWithRetry(c, channelHandler, group, bodyBytes, isStream, startTime, retryCount+1, newRetryErrors)
|
||||
return
|
||||
}
|
||||
|
||||
ps.keyProvider.UpdateStatus(apiKey.ID, group.ID, true)
|
||||
logrus.Debugf("Request for group %s succeeded on attempt %d with key %s", group.Name, retryCount+1, apiKey.KeyValue[:8])
|
||||
|
||||
for key, values := range resp.Header {
|
||||
for _, value := range values {
|
||||
c.Header(key, value)
|
||||
}
|
||||
}
|
||||
c.Status(resp.StatusCode)
|
||||
|
||||
if isStream {
|
||||
ps.handleStreamingResponse(c, resp)
|
||||
} else {
|
||||
ps.handleNormalResponse(c, resp)
|
||||
}
|
||||
}
|
||||
|
||||
func (ps *ProxyServer) handleStreamingResponse(c *gin.Context, resp *http.Response) {
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
logrus.Error("Streaming unsupported by the writer")
|
||||
ps.handleNormalResponse(c, resp)
|
||||
return
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
for scanner.Scan() {
|
||||
if _, err := c.Writer.Write(scanner.Bytes()); err != nil {
|
||||
if !isIgnorableStreamError(err) {
|
||||
logrus.Errorf("Error writing to client: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if _, err := c.Writer.Write([]byte("\n")); err != nil {
|
||||
if !isIgnorableStreamError(err) {
|
||||
logrus.Errorf("Error writing newline to client: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil && !isIgnorableStreamError(err) {
|
||||
logrus.Errorf("Error reading from upstream: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (ps *ProxyServer) handleNormalResponse(c *gin.Context, resp *http.Response) {
|
||||
if _, err := io.Copy(c.Writer, resp.Body); err != nil {
|
||||
if !isIgnorableStreamError(err) {
|
||||
logrus.Errorf("Failed to copy response body to client: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (ps *ProxyServer) applyParamOverrides(bodyBytes []byte, group *models.Group) ([]byte, error) {
|
||||
if len(group.ParamOverrides) == 0 || len(bodyBytes) == 0 {
|
||||
return bodyBytes, nil
|
||||
}
|
||||
|
||||
// Unmarshal the body into a map
|
||||
var requestData map[string]any
|
||||
if err := json.Unmarshal(bodyBytes, &requestData); err != nil {
|
||||
// If not a valid JSON, just pass it through
|
||||
c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
|
||||
return nil
|
||||
logrus.Warnf("failed to unmarshal request body for param override, passing through: %v", err)
|
||||
return bodyBytes, nil
|
||||
}
|
||||
|
||||
// Merge the overrides into the request data
|
||||
for key, value := range group.ParamOverrides {
|
||||
requestData[key] = value
|
||||
}
|
||||
|
||||
// Marshal the new data back to JSON
|
||||
newBodyBytes, err := json.Marshal(requestData)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal new request body: %w", err)
|
||||
}
|
||||
|
||||
// Replace the request body with the new one
|
||||
c.Request.Body = io.NopCloser(bytes.NewBuffer(newBodyBytes))
|
||||
c.Request.ContentLength = int64(len(newBodyBytes))
|
||||
|
||||
// Restore the original Content-Type header
|
||||
if originalContentType != "" {
|
||||
c.Request.Header.Set("Content-Type", originalContentType)
|
||||
}
|
||||
|
||||
return nil
|
||||
return json.Marshal(requestData)
|
||||
}
|
||||
|
||||
func (ps *ProxyServer) Close() {
|
||||
// The HTTP clients are now managed by the channel factory and httpclient manager,
|
||||
// so the proxy server itself doesn't need to close them.
|
||||
// The httpclient manager will handle closing idle connections for all its clients.
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user