diff --git a/main.go b/main.go index 4ff2da6..6e0a0f4 100644 --- a/main.go +++ b/main.go @@ -4,6 +4,7 @@ import ( "flag" "fmt" "log" + "net/http" "os" "strings" "time" @@ -24,11 +25,11 @@ func main() { noauth := flag.Bool("noauth", false, "Do not check incoming authorization header") flag.Parse() - log.Println("Service starting") + log.Println("[main]: Service starting") // load all upstreams config = readConfig(*configFile) - log.Println("Load upstreams number:", len(config.Upstreams)) + log.Println("[main]: Load upstreams number:", len(config.Upstreams)) // connect to database var db *gorm.DB @@ -40,7 +41,7 @@ func main() { SkipDefaultTransaction: true, }) if err != nil { - log.Fatalf("Error to connect sqlite database: %s", err) + log.Fatalf("[main]: Error to connect sqlite database: %s", err) } case "postgres": db, err = gorm.Open(postgres.Open(config.DBAddr), &gorm.Config{ @@ -48,14 +49,14 @@ func main() { SkipDefaultTransaction: true, }) if err != nil { - log.Fatalf("Error to connect postgres database: %s", err) + log.Fatalf("[main]: Error to connect postgres database: %s", err) } default: - log.Fatalf("Unsupported database type: '%s'", config.DBType) + log.Fatalf("[main]: Unsupported database type: '%s'", config.DBType) } db.AutoMigrate(&Record{}) - log.Println("Auto migrate database done") + log.Println("[main]: Auto migrate database done") if *listMode { fmt.Println("SK\tEndpoint") @@ -130,7 +131,7 @@ func main() { for index, upstream := range config.Upstreams { if upstream.Endpoint == "" || upstream.SK == "" { - c.AbortWithError(500, fmt.Errorf("invaild upstream '%s' '%s'", upstream.SK, upstream.Endpoint)) + c.AbortWithError(500, fmt.Errorf("[processRequest.begin]: invaild upstream '%s' '%s'", upstream.SK, upstream.Endpoint)) continue } @@ -142,30 +143,37 @@ func main() { err = processRequest(c, &upstream, &record, shouldResponse) if err != nil { - log.Println("Error from upstream", upstream.Endpoint, "should retry", err) + if err == http.ErrAbortHandler { + abortErr := "[processRequest.done]: AbortHandler, client's connection lost?, no upstream will try, stop here" + log.Println(abortErr) + record.Response += abortErr + record.Status = 500 + break + } + log.Println("[processRequest.done]: Error from upstream", upstream.Endpoint, "should retry", err) continue } break } - log.Println("Record result:", record.Status, record.Response) + log.Println("[final]: Record result:", record.Status, record.Response) record.ElapsedTime = time.Now().Sub(record.CreatedAt) // async record request go func() { // turncate request if too long if len(record.Body) > 1024*128 { - log.Println("Warning: Truncate request body") + log.Println("[async.record]: Warning: Truncate request body") record.Body = record.Body[:1024*128] } if db.Create(&record).Error != nil { - log.Println("Error to save record:", record) + log.Println("[async.record]: Error to save record:", record) } }() if record.Status != 200 { - errMessage := fmt.Sprintf("IP: %s request %s error %d with %s", record.IP, record.Model, record.Status, record.Response) + errMessage := fmt.Sprintf("[result.error]: IP: %s request %s error %d with %s", record.IP, record.Model, record.Status, record.Response) go sendFeishuMessage(errMessage) go sendMatrixMessage(errMessage) } diff --git a/process.go b/process.go index 69db7f8..3f6af48 100644 --- a/process.go +++ b/process.go @@ -31,6 +31,10 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s return err } + remote.Path = upstream.URL.Path + strings.TrimPrefix(c.Request.URL.Path, "/v1") + log.Println("[proxy.begin]:", remote) + log.Println("[proxy.begin]: shouldResposne:", shouldResponse) + haveResponse := false proxy := httputil.NewSingleHostReverseProxy(remote) @@ -48,7 +52,7 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s // read request body inBody, err = io.ReadAll(in.Body) if err != nil { - errCtx = errors.New("reverse proxy middleware failed to read request body " + err.Error()) + errCtx = errors.New("[proxy.rewrite]: reverse proxy middleware failed to read request body " + err.Error()) return } @@ -77,7 +81,7 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s go func() { time.Sleep(timeout) if !haveResponse { - log.Println("Timeout upstream", upstream.Endpoint) + log.Println("[proxy.timeout]: Timeout upstream", upstream.Endpoint, timeout) errCtx = errors.New("timeout") if shouldResponse { c.Header("Content-Type", "application/json") @@ -93,8 +97,6 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s out.URL.Scheme = remote.Scheme out.URL.Host = remote.Host - out.URL.Path = upstream.URL.Path + strings.TrimPrefix(in.URL.Path, "/v1") - out.Header = http.Header{} out.Header.Set("Host", remote.Host) if upstream.SK == "asis" { @@ -123,19 +125,20 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s } if !shouldResponse && r.StatusCode != 200 { - log.Println("upstream return not 200 and should not response", r.StatusCode) + log.Println("[proxy.modifyResponse]: upstream return not 200 and should not response", r.StatusCode) return errors.New("upstream return not 200 and should not response") } if r.StatusCode != 200 { body, err := io.ReadAll(r.Body) if err != nil { - record.Response += "[Error]: failed to read response from upstream " + err.Error() - return errors.New(record.Response) + errRet := errors.New("[proxy.modifyResponse]: failed to read response from upstream " + err.Error()) + return errRet } - record.Response += fmt.Sprintf("[Error]: openai-api-route upstream return '%s' with '%s'", r.Status, string(body)) + errRet := errors.New(fmt.Sprintf("[error]: openai-api-route upstream return '%s' with '%s'", r.Status, string(body))) + log.Println(errRet) record.Status = r.StatusCode - return fmt.Errorf(record.Response) + return errRet } // count success r.Body = io.NopCloser(io.TeeReader(r.Body, &buf)) @@ -145,7 +148,7 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) { haveResponse = true record.ResponseTime = time.Now().Sub(record.CreatedAt) - log.Println("Error", err, upstream.SK, upstream.Endpoint) + log.Println("[proxy.errorHandler]", err, upstream.SK, upstream.Endpoint) errCtx = err @@ -155,21 +158,26 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s c.AbortWithError(502, err) } - log.Println("response is", r.Response) + log.Println("[proxy.errorHandler]: response is", r.Response) if record.Status == 0 { record.Status = 502 } - if record.Response == "" { - record.Response += "[Error]: " + err.Error() - } + record.Response += "[proxy.ErrorHandler]: " + err.Error() if r.Response != nil { record.Status = r.Response.StatusCode } } - proxy.ServeHTTP(c.Writer, c.Request) + err = ServeHTTP(proxy, c.Writer, c.Request) + if err != nil { + log.Println("[proxy.serve]: error from ServeHTTP:", err) + // panic means client has abort the http connection + // since the connection is lost, we return + // and the reverse process should not try the next upsteam + return http.ErrAbortHandler + } // return context error if errCtx != nil { @@ -197,7 +205,7 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s err := json.Unmarshal([]byte(line), &chunk) if err != nil { - log.Println(err) + log.Println("[proxy.parseChunkError]:", err) continue } @@ -210,20 +218,20 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s var fetchResp FetchModeResponse err := json.Unmarshal(resp, &fetchResp) if err != nil { - log.Println("Error parsing fetch response:", err) + log.Println("[proxy.parseJSONError]: error parsing fetch response:", err) return nil } if !strings.HasPrefix(fetchResp.Model, "gpt-") { - log.Println("Not GPT model, skip recording response:", fetchResp.Model) + log.Println("[proxy.record]: Not GPT model, skip recording response:", fetchResp.Model) return nil } if len(fetchResp.Choices) == 0 { - log.Println("Error: fetch response choice length is 0") + log.Println("[proxy.record]: Error: fetch response choice length is 0") return nil } record.Response = fetchResp.Choices[0].Message.Content } else { - log.Println("Unknown content type", contentType) + log.Println("[proxy.record]: Unknown content type", contentType) } } diff --git a/recovery.go b/recovery.go new file mode 100644 index 0000000..ebbf687 --- /dev/null +++ b/recovery.go @@ -0,0 +1,24 @@ +package main + +import ( + "errors" + "log" + "net/http" + "net/http/httputil" + + "github.com/gin-gonic/gin" +) + +func ServeHTTP(proxy *httputil.ReverseProxy, w gin.ResponseWriter, r *http.Request) (errReturn error) { + + // recovery + defer func() { + if err := recover(); err != nil { + log.Println("[serve.panic]: ", err) + errReturn = errors.New("[serve.panic]: Panic recover in reverse proxy serve HTTP") + } + }() + + proxy.ServeHTTP(w, r) + return nil +}