third-party-mirrors/ollama
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

only retry once on auth failure (#1175)

Browse Source
This commit is contained in:
Bruce MacDonald 2023-11-17 14:22:35 -05:00 committed by GitHub
parent 0813387414
commit 0b19e24d81
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 27 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
server/images.go
View File

@ -1146,43 +1146,42 @@ func GetSHA256Digest(r io.Reader) (string, int64) {
var errUnauthorized = fmt.Errorf("unauthorized") var errUnauthorized = fmt.Errorf("unauthorized")
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) { func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
lastErr := errMaxRetriesExceeded resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
for try := 0; try < maxRetries; try++ { if err != nil {
resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts) if !errors.Is(err, context.Canceled) {
log.Printf("request failed: %v", err)
}
return nil, err
}
switch {
case resp.StatusCode == http.StatusUnauthorized:
// Handle authentication error with one retry
auth := resp.Header.Get("www-authenticate")
authRedir := ParseAuthRedirectString(auth)
token, err := getAuthToken(ctx, authRedir)
if err != nil { if err != nil {
log.Printf("couldn't start upload: %v", err)
return nil, err return nil, err
} }
regOpts.Token = token
switch { if body != nil {
case resp.StatusCode == http.StatusUnauthorized: _, err = body.Seek(0, io.SeekStart)
auth := resp.Header.Get("www-authenticate")
authRedir := ParseAuthRedirectString(auth)
token, err := getAuthToken(ctx, authRedir)
if err != nil { if err != nil {
return nil, err return nil, err
} }
regOpts.Token = token
if body != nil {
body.Seek(0, io.SeekStart)
}
lastErr = errUnauthorized
case resp.StatusCode == http.StatusNotFound:
return nil, os.ErrNotExist
case resp.StatusCode >= http.StatusBadRequest:
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("%d: %s", resp.StatusCode, err)
}
return nil, fmt.Errorf("%d: %s", resp.StatusCode, body)
default:
return resp, nil
} }
return makeRequest(ctx, method, requestURL, headers, body, regOpts)
case resp.StatusCode == http.StatusNotFound:
return nil, os.ErrNotExist
case resp.StatusCode >= http.StatusBadRequest:
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("%d: %s", resp.StatusCode, err)
}
return nil, fmt.Errorf("%d: %s", resp.StatusCode, responseBody)
} }
return nil, lastErr return resp, nil
} }
func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) { func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {