diff --git a/process.go b/process.go index 694f935..65767e7 100644 --- a/process.go +++ b/process.go @@ -11,8 +11,10 @@ import ( "net/http/httputil" "net/url" "strings" + "time" "github.com/gin-gonic/gin" + "golang.org/x/net/context" ) func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, shouldResponse bool) error { @@ -29,11 +31,24 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s c.AbortWithError(500, errors.New("can't parse reverse proxy remote URL")) return err } + + // set timeout, default is 5 second + timeout := 5 * time.Second + haveResponse := false + if upstream.Timeout > 0 { + // convert upstream.Timeout(second) to nanosecond + timeout = time.Duration(upstream.Timeout) * time.Second + } + proxy := httputil.NewSingleHostReverseProxy(remote) proxy.Director = nil var inBody []byte proxy.Rewrite = func(proxyRequest *httputil.ProxyRequest) { in := proxyRequest.In + + ctx, cancel := context.WithCancel(context.Background()) + proxyRequest.Out = proxyRequest.Out.WithContext(ctx) + out := proxyRequest.Out // read request body @@ -46,6 +61,16 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s // record chat message from user record.Body = string(inBody) + // timeout out request + go func() { + time.Sleep(timeout) + if !haveResponse { + log.Println("Timeout", upstream.Endpoint) + errCtx = errors.New("timeout") + cancel() + } + }() + out.Body = io.NopCloser(bytes.NewReader(inBody)) out.Host = remote.Host @@ -64,6 +89,8 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s var buf bytes.Buffer var contentType string proxy.ModifyResponse = func(r *http.Response) error { + haveResponse = true + log.Println("haveResponse set to true") record.Status = r.StatusCode if !shouldResponse && r.StatusCode != 200 { log.Println("upstream return not 200 and should not response", r.StatusCode) diff --git a/structure.go b/structure.go index 08509ca..c47c116 100644 --- a/structure.go +++ b/structure.go @@ -9,4 +9,5 @@ type OPENAI_UPSTREAM struct { gorm.Model SK string `gorm:"index:idx_sk_endpoint,unique"` // key Endpoint string `gorm:"index:idx_sk_endpoint,unique"` // endpoint + Timeout int64 // timeout in seconds }