fix: keep proxy slot through response body
This commit is contained in:
@@ -10,6 +10,7 @@ import (
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
@@ -50,6 +51,7 @@ func (s *Service) ProxyHTTP(w http.ResponseWriter, r *http.Request, targetIP, pr
|
||||
Transport: retryTransport{
|
||||
base: &http.Transport{
|
||||
DisableKeepAlives: true,
|
||||
DisableCompression: true,
|
||||
ResponseHeaderTimeout: 5 * time.Second,
|
||||
},
|
||||
limiter: s.proxyLimiter(targetIP),
|
||||
@@ -139,10 +141,20 @@ func (t retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := base.RoundTrip(nextReq)
|
||||
releaseProxySlot(t.limiter)
|
||||
if err == nil {
|
||||
if resp == nil || resp.Body == nil {
|
||||
releaseProxySlot(t.limiter)
|
||||
return resp, nil
|
||||
}
|
||||
resp.Body = &releaseOnCloseReadCloser{
|
||||
ReadCloser: resp.Body,
|
||||
release: func() {
|
||||
releaseProxySlot(t.limiter)
|
||||
},
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
releaseProxySlot(t.limiter)
|
||||
lastErr = err
|
||||
if !shouldRetryProxyRequest(req, err) || attempt == attempts-1 {
|
||||
return nil, err
|
||||
@@ -152,6 +164,22 @@ func (t retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
type releaseOnCloseReadCloser struct {
|
||||
io.ReadCloser
|
||||
once sync.Once
|
||||
release func()
|
||||
}
|
||||
|
||||
func (r *releaseOnCloseReadCloser) Close() error {
|
||||
err := r.ReadCloser.Close()
|
||||
r.once.Do(func() {
|
||||
if r.release != nil {
|
||||
r.release()
|
||||
}
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func acquireProxySlot(ctx context.Context, limiter chan struct{}, timeout time.Duration) error {
|
||||
if limiter == nil {
|
||||
return nil
|
||||
|
||||
Reference in New Issue
Block a user