Skip to content

Commit c5d7adf

Browse files
zhaoxianhuagithubgxll
authored andcommitted
[fix][dingospeed] Fixed download interruption, integrity validation failure
1 parent 0043b2a commit c5d7adf

12 files changed

Lines changed: 239 additions & 73 deletions

File tree

docker/Dockerfile-alpine

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,11 @@ ENV TZ=Asia/Shanghai
1919
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
2020

2121
# 安装 huggingface-hub(指定版本)
22-
RUN pip3 install --no-cache-dir --break-system-packages huggingface-hub==1.2.3
22+
RUN pip3 install --no-cache-dir --break-system-packages huggingface-hub==1.12.0
2323
RUN pip3 install --no-cache-dir --break-system-packages hf_transfer==0.1.9
24-
RUN pip3 install --no-cache-dir --break-system-packages hf-xet==1.2.0
24+
RUN pip3 install --no-cache-dir --break-system-packages hf-xet==1.4.3
2525
# 验证安装
2626
RUN hf version
2727
CMD ["/bin/sh"]
2828

29-
# docker build -t huggingface-hub:alpine-1.2.3 -f Dockerfile-alpine .
29+
# docker build -t huggingface-hub:alpine-1.12.0 -f Dockerfile-alpine .

internal/dao/downloader_dao.go

Lines changed: 30 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -43,54 +43,52 @@ func NewDownloaderDao(schedulerDao *SchedulerDao) *DownloaderDao {
4343
}
4444

4545
// 整个文件
46-
func (d *DownloaderDao) FileDownload(chanErr chan error, startPos, endPos int64, isInnerRequest bool, taskParam *downloader.TaskParam) {
47-
var (
48-
wg sync.WaitGroup
49-
)
50-
defer close(chanErr)
51-
defer close(taskParam.ResponseChan)
46+
func (d *DownloaderDao) FileDownload(startPos, endPos int64, isInnerRequest bool, taskParam *downloader.TaskParam) error {
5247
dingCacheManager := downloader.GetInstance()
5348
dingFile, err := dingCacheManager.GetDingFile(taskParam.BlobsFile, taskParam.FileSize)
5449
if err != nil {
5550
zap.S().Errorf("GetDingFile err.%v", err)
56-
chanErr <- myerr.NewAppendCode(http.StatusInternalServerError, "Get DingFile err")
57-
return
51+
return myerr.NewAppendCode(http.StatusInternalServerError, "Get DingFile err")
5852
}
59-
defer func() {
60-
dingCacheManager.ReleasedDingFile(taskParam.BlobsFile)
61-
}()
6253
taskParam.DingFile = dingFile
6354
tasks, err := d.constructTask(startPos, endPos, isInnerRequest, taskParam)
6455
if err != nil {
65-
chanErr <- err
66-
return
56+
return err
6757
}
68-
wg.Add(1)
6958
go func() {
59+
defer close(taskParam.ResponseChan)
7060
defer func() {
71-
wg.Done()
61+
dingCacheManager.ReleasedDingFile(taskParam.BlobsFile)
7262
}()
73-
for i := 0; i < len(tasks); i++ {
74-
if taskParam.Context.Err() != nil {
75-
break
76-
}
77-
task := tasks[i]
78-
if i == 0 {
79-
task.GetResponseChan() <- []byte{} // 先建立长连接
80-
}
81-
task.OutResult()
82-
}
83-
}()
84-
if len(tasks) > 0 {
63+
var wg sync.WaitGroup
8564
wg.Add(1)
8665
go func() {
8766
defer func() {
8867
wg.Done()
8968
}()
90-
doTask(taskParam.Context, tasks)
69+
for i := 0; i < len(tasks); i++ {
70+
if taskParam.Context.Err() != nil {
71+
break
72+
}
73+
task := tasks[i]
74+
if i == 0 {
75+
task.GetResponseChan() <- []byte{} // 先建立长连接
76+
}
77+
task.OutResult()
78+
}
9179
}()
92-
}
93-
wg.Wait() // 等待协程池所有远程下载任务执行完毕
80+
if len(tasks) > 0 {
81+
wg.Add(1)
82+
go func() {
83+
defer func() {
84+
wg.Done()
85+
}()
86+
doTask(taskParam.Context, tasks)
87+
}()
88+
}
89+
wg.Wait() // 等待协程池所有远程下载任务执行完毕
90+
}()
91+
return nil
9492
}
9593

9694
func (d *DownloaderDao) constructTask(startPos, endPos int64, isInnerRequest bool, taskParam *downloader.TaskParam) ([]common.DownloadTask, error) {
@@ -107,7 +105,7 @@ func (d *DownloaderDao) constructTask(startPos, endPos int64, isInnerRequest boo
107105
// 分析下载类型是否全部存在,若文件不完整,返回当前已缓存的最大偏移量
108106
fileComplete, curPos = analysisFilePosition(taskParam.DingFile, startPos, endPos)
109107
if !fileComplete && !config.SysConfig.Online() { // 文件不完整,且当前节点为离线
110-
return nil, myerr.NewAppendCode(http.StatusNotFound, "model file is not exist")
108+
return nil, myerr.NewAppendCode(http.StatusNotFound, "Entry not found")
111109
}
112110
// isInnerRequest为true,即内部请求,是已经被调度过后,设置为内部域名的请求,这种请求将不会再次参与调度,直接做下载即可。
113111
if !isInnerRequest && config.SysConfig.IsCluster() && !fileComplete {
@@ -237,7 +235,7 @@ func getContiguousRanges(startPos, endPos int64, taskParam *downloader.TaskParam
237235
if startPos == 0 && endPos == 0 {
238236
return
239237
}
240-
if startPos < 0 || endPos <= startPos || endPos > dingFile.GetFileSize() {
238+
if startPos < 0 || endPos <= startPos || (endPos-1) > dingFile.GetFileSize() {
241239
zap.S().Errorf("Invalid pos path=%s, startPos=%d, endPos=%d", dingFile.GetPath(), startPos, endPos)
242240
return
243241
}

internal/dao/file_dao.go

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -227,18 +227,25 @@ func (f *FileDao) FileGetGenerator(c echo.Context, repoType, orgRepo, commit, fi
227227

228228
func constructRespHeader(c echo.Context, pathInfo *common.PathsInfo, commit, fileName string) (map[string]string, string, int64, int64) {
229229
var startPos, endPos int64
230+
var clientHasRange bool
230231
if pathInfo.Size > 0 { // There exists a file of size 0
231232
var headRange = c.Request().Header.Get("Range")
232-
if headRange == "" {
233-
headRange = fmt.Sprintf("bytes=%d-%d", 0, pathInfo.Size-1)
233+
defaultEndPos := pathInfo.Size - 1
234+
if headRange != "" {
235+
clientHasRange = true
236+
} else {
237+
headRange = fmt.Sprintf("bytes=%d-%d", 0, defaultEndPos)
234238
}
235-
startPos, endPos = parseRangeParams(headRange, pathInfo.Size)
239+
startPos, endPos = parseRangeParams(headRange, defaultEndPos)
236240
endPos = endPos + 1
237241
} else if pathInfo.Size == 0 {
238242
zap.S().Warnf("file %s size: %d", fileName, pathInfo.Size)
239243
}
240244
respHeaders := map[string]string{}
241245
respHeaders[consts.HUGGINGFACE_HEADER_CONTENT_LENGTH] = util.Itoa(endPos - startPos)
246+
if clientHasRange && pathInfo.Size > 0 {
247+
respHeaders["Content-Range"] = fmt.Sprintf("bytes %d-%d/%d", startPos, endPos-1, pathInfo.Size)
248+
}
242249
if commit != "" {
243250
respHeaders[strings.ToLower(consts.HUGGINGFACE_HEADER_X_REPO_COMMIT)] = commit
244251
}
@@ -455,10 +462,11 @@ func (f *FileDao) FileChunkGet(c echo.Context, taskParam *downloader.TaskParam,
455462
taskParam.Context = ctx
456463
taskParam.ResponseChan = responseChan
457464
taskParam.Cancel = cancel
458-
fileErrCh := make(chan error, 1)
459465
fileName := fmt.Sprintf("%s/%s", taskParam.OrgRepo, taskParam.FileName)
460-
go f.downloaderDao.FileDownload(fileErrCh, startPos, endPos, isInnerRequest, taskParam)
461-
if err := util.ResponseStream(ctx, c, fileName, respHeaders, responseChan, fileErrCh); err != nil {
466+
if err := f.downloaderDao.FileDownload(startPos, endPos, isInnerRequest, taskParam); err != nil {
467+
return util.MultipleErrorProxyError(err, c)
468+
}
469+
if err := util.ResponseStream(c, fileName, respHeaders, responseChan); err != nil {
462470
zap.S().Errorf("FileChunkGet stream err.%v", err)
463471
return util.ErrorProxyTimeout(c)
464472
}
@@ -526,7 +534,7 @@ func (f *FileDao) GetFileOffset(dataType string, org string, repo string, etag s
526534
return curPos
527535
}
528536

529-
func parseRangeParams(fileRange string, fileSize int64) (int64, int64) {
537+
func parseRangeParams(fileRange string, defaultEndPos int64) (int64, int64) {
530538
if strings.Contains(fileRange, "/") {
531539
split := strings.SplitN(fileRange, "/", 2)
532540
fileRange = split[0]
@@ -548,7 +556,7 @@ func parseRangeParams(fileRange string, fileSize int64) (int64, int64) {
548556
if s2 != 0 {
549557
endPos = s2
550558
} else {
551-
endPos = fileSize - 1
559+
endPos = defaultEndPos
552560
}
553561
return startPos, endPos
554562
}

internal/dao/meta_dao.go

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -134,9 +134,13 @@ func (m *MetaDao) GetMetadata(repoType, orgRepo, revision, method, authorization
134134
}
135135
}
136136
} else {
137-
if cacheContent, err = m.fileDao.ReadCacheRequest(apiMetaPath); err != nil {
138-
zap.S().Errorf("ReadCacheRequest err.%v", err)
139-
return nil, err
137+
if util.FileExists(apiMetaPath) {
138+
if cacheContent, err = m.fileDao.ReadCacheRequest(apiMetaPath); err != nil {
139+
zap.S().Errorf("ReadCacheRequest err.%v", err)
140+
return nil, err
141+
}
142+
} else {
143+
return nil, myerr.NewAppendCode(http.StatusNotFound, fmt.Sprintf("%s not exist", orgRepo))
140144
}
141145
}
142146
return cacheContent, nil

internal/downloader/cache_task.go

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ func (c *CacheFileTask) DoTask() {
8383
func (c *CacheFileTask) OutResult() {
8484
startBlock := c.RangeStartPos / c.DingFile.GetBlockSize()
8585
endBlock := (c.RangeEndPos - 1) / c.DingFile.GetBlockSize()
86+
blockNumber := c.DingFile.getBlockNumber()
8687
curPos := c.RangeStartPos
8788
for curBlock := startBlock; curBlock <= endBlock; curBlock++ {
8889
if c.Context.Err() != nil {
@@ -104,8 +105,9 @@ func (c *CacheFileTask) OutResult() {
104105
zap.S().Errorf("ReadBlock err file:%s, %v", c.FileName, err)
105106
continue
106107
}
107-
sPos := max(c.RangeStartPos, blockStartPos) - blockStartPos
108-
ePos := min(c.RangeEndPos, blockEndPos) - blockStartPos
108+
maxStart, minEnd := max(c.RangeStartPos, blockStartPos), min(c.RangeEndPos, blockEndPos)
109+
sPos := maxStart - blockStartPos
110+
ePos := minEnd - blockStartPos
109111
rawLen := int64(len(rawBlock))
110112
if rawLen == 0 || sPos > rawLen {
111113
zap.S().Errorf("read rawBlock err,%s, rawLen:%d, sPos:%d,ePos:%d, %v", c.FileName, rawLen, sPos, ePos, err)
@@ -118,6 +120,7 @@ func (c *CacheFileTask) OutResult() {
118120
chunk := rawBlock[sPos:ePos]
119121
select {
120122
case c.ResponseChan <- chunk:
123+
zap.S().Debugf("%s/%s, taskNo:%d, block:%d(%d)write done, range:%d-%d.", c.OrgRepo, c.FileName, c.TaskNo, curBlock, blockNumber, maxStart, minEnd)
121124
case <-c.Context.Done():
122125
return
123126
}

internal/downloader/remote_task.go

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,7 @@ func (r *RemoteFileTask) OutResult() {
206206
select {
207207
case chunk, ok := <-r.Queue:
208208
if !ok {
209-
zap.S().Debugf("end remote outResult close. taskNo:%d, %s/%s", r.TaskNo, r.OrgRepo, r.FileName)
209+
zap.S().Debugf("close remote outResult. taskNo:%d, %s/%s", r.TaskNo, r.OrgRepo, r.FileName)
210210
return
211211
}
212212
select {
@@ -216,7 +216,7 @@ func (r *RemoteFileTask) OutResult() {
216216
return
217217
}
218218
case <-r.Context.Done():
219-
zap.S().Debugf("end remote outResult fileName:%s/%s,err:%v", r.OrgRepo, r.FileName, r.Context.Err())
219+
zap.S().Debugf("close remote outResult fileName:%s/%s,err:%v", r.OrgRepo, r.FileName, r.Context.Err())
220220
return
221221
}
222222
}
@@ -239,7 +239,9 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC
239239
if r.Authorization != "" {
240240
headers["authorization"] = r.Authorization
241241
}
242-
headers["range"] = fmt.Sprintf("bytes=%d-%d", startPos, endPos-1)
242+
if startPos > 0 || endPos < r.DingFile.GetFileSize() {
243+
headers["range"] = fmt.Sprintf("bytes=%d-%d", startPos, endPos-1)
244+
}
243245
for i := 0; i < attempts; {
244246
if _, err = util.RetryRequest(func() (*common.Response, error) {
245247
err = util.GetStream(r.Domain, r.Uri, headers, func(resp *http.Response) error {
@@ -276,6 +278,12 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC
276278
}
277279
if err != nil {
278280
if err == io.EOF {
281+
if int64(chunkByteLen) < (endPos - startPos) {
282+
// 数据不完整,将EOF视为读取错误以触发重试/断点续传
283+
zap.S().Errorf("file:%s/%s, taskNo:%d, premature EOF: expected %d bytes, got %d", r.OrgRepo, r.FileName, r.TaskNo, endPos-startPos, chunkByteLen)
284+
headers["range"] = fmt.Sprintf("bytes=%d-%d", startPos+int64(chunkByteLen), endPos-1)
285+
return fmt.Errorf("premature EOF: expected %d bytes, got %d", endPos-startPos, chunkByteLen)
286+
}
279287
return nil
280288
}
281289
zap.S().Errorf("file:%s/%s, taskNo:%d, statusCode:%d, chunkByteLen:%d, %v", r.OrgRepo, r.FileName, r.TaskNo, resp.StatusCode, chunkByteLen, err)

internal/handler/meta_handler.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
package handler
1616

1717
import (
18-
"context"
1918
"net/http"
2019
"strings"
2120

@@ -69,7 +68,7 @@ func (handler *MetaHandler) GetMetadataHandler(c echo.Context) error {
6968
var bodyStreamChan = make(chan []byte, consts.RespChanSize)
7069
bodyStreamChan <- cacheContent.OriginContent
7170
close(bodyStreamChan)
72-
err = util.ResponseStream(context.Background(), c, orgRepo, cacheContent.Headers, bodyStreamChan, nil)
71+
err = util.ResponseStream(c, orgRepo, cacheContent.Headers, bodyStreamChan)
7372
if err != nil {
7473
return err
7574
}

internal/service/meta_service.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
package service
1616

1717
import (
18-
"context"
1918
"fmt"
2019
"io"
2120
"sort"
@@ -105,7 +104,7 @@ func (m *MetaService) RepoRefs(c echo.Context, repoType, org, repo string) error
105104
var bodyStreamChan = make(chan []byte, consts.RespChanSize)
106105
bodyStreamChan <- cacheContent.OriginContent
107106
close(bodyStreamChan)
108-
return util.ResponseStream(context.Background(), c, orgRepo, cacheContent.Headers, bodyStreamChan, nil)
107+
return util.ResponseStream(c, orgRepo, cacheContent.Headers, bodyStreamChan)
109108
}
110109

111110
func (m *MetaService) ForwardToNewSite(c echo.Context) error {

internal/service/task/preheat_cache_task.go

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,9 +140,7 @@ func (p *PreheatCacheTask) startPreheat(hfUri, orgRepo, fileName, commit, etag,
140140
return p.result(ctx, responseChan)
141141
})
142142
eg.Go(func() error {
143-
fileErrCh := make(chan error, 1)
144-
p.DownloaderDao.FileDownload(fileErrCh, offset, fileSize, false, taskParam)
145-
if err := <-fileErrCh; err != nil {
143+
if err := p.DownloaderDao.FileDownload(offset, fileSize, false, taskParam); err != nil {
146144
return err
147145
}
148146
return nil

pkg/util/http_util.go

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ package util
1616

1717
import (
1818
"bytes"
19-
"context"
2019
"fmt"
2120
"io"
2221
"net"
@@ -29,7 +28,6 @@ import (
2928
"dingospeed/pkg/common"
3029
"dingospeed/pkg/config"
3130
"dingospeed/pkg/consts"
32-
myerr "dingospeed/pkg/error"
3331
"dingospeed/pkg/prom"
3432

3533
"github.com/avast/retry-go"
@@ -306,26 +304,25 @@ func doPost(client *http.Client, targetURL string, contentType string, data []by
306304
}, nil
307305
}
308306

309-
func ResponseStream(ctx context.Context, c echo.Context, fileName string, headers map[string]string, content <-chan []byte, fileErrCh chan error) error {
307+
func ResponseStream(c echo.Context, fileName string, headers map[string]string, content <-chan []byte) error {
310308
c.Response().Header().Set("Content-Type", "text/event-stream")
311309
c.Response().Header().Set("Cache-Control", "no-cache")
312310
c.Response().Header().Set("Connection", "keep-alive")
313311
for k, v := range headers {
312+
// 流式响应不应预设Content-Length,因为数据是流式传输的,无法保证字节数与预估值一致。
313+
// 若实际传输字节数与Content-Length不符,客户端会报错(如curl: (18) transfer closed with N bytes remaining to read)。
314314
c.Response().Header().Set(k, v)
315315
}
316-
if fileErrCh != nil {
317-
if err := <-fileErrCh; err != nil {
318-
if e, ok := err.(myerr.Error); ok {
319-
return c.String(e.StatusCode(), err.Error())
320-
}
321-
return ErrorProxyError(c)
322-
}
316+
// 根据 headers 中是否包含 Content-Range 来决定状态码
317+
statusCode := http.StatusOK
318+
if c.Response().Header().Get("Content-Range") != "" {
319+
statusCode = http.StatusPartialContent
323320
}
321+
c.Response().WriteHeader(statusCode)
324322
flusher, ok := c.Response().Writer.(http.Flusher)
325323
if !ok {
326324
return c.String(http.StatusInternalServerError, "Streaming unsupported!")
327325
}
328-
c.Response().WriteHeader(http.StatusOK)
329326
for {
330327
select {
331328
case b, ok := <-content:
@@ -346,8 +343,6 @@ func ResponseStream(ctx context.Context, c echo.Context, fileName string, header
346343
}
347344
}
348345
flusher.Flush()
349-
case <-ctx.Done():
350-
return c.String(http.StatusInternalServerError, ctx.Err().Error())
351346
}
352347
}
353348
}

0 commit comments

Comments
 (0)