Skip to content

Commit d2576ac

Browse files
committed
Fix downloadChunked for larger files
closes #298
1 parent 14b47a0 commit d2576ac

File tree

2 files changed

+24
-13
lines changed

2 files changed

+24
-13
lines changed

client.go

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -338,7 +338,6 @@ func (c *Client) GetStreamContext(ctx context.Context, video *Video, format *For
338338
} else {
339339
// we have length information, let's download by chunks!
340340
c.downloadChunked(ctx, req, w, format)
341-
342341
}
343342

344343
return r, contentLength, nil
@@ -402,19 +401,21 @@ func (c *Client) downloadChunked(ctx context.Context, req *http.Request, w *io.P
402401
currentChunk := atomic.Uint32{}
403402
for i := 0; i < maxRoutines; i++ {
404403
go func() {
405-
i := int(currentChunk.Add(1)) - 1
406-
if i > len(chunks) {
407-
// no more chunks
408-
return
409-
}
404+
for {
405+
chunkIndex := int(currentChunk.Add(1)) - 1
406+
if chunkIndex >= len(chunks) {
407+
// no more chunks
408+
return
409+
}
410410

411-
chunk := &chunks[i]
412-
err := c.downloadChunk(req.Clone(cancelCtx), chunk)
413-
close(chunk.data)
411+
chunk := &chunks[chunkIndex]
412+
err := c.downloadChunk(req.Clone(cancelCtx), chunk)
413+
close(chunk.data)
414414

415-
if err != nil {
416-
abort(err)
417-
return
415+
if err != nil {
416+
abort(err)
417+
return
418+
}
418419
}
419420
}()
420421
}
@@ -424,6 +425,7 @@ func (c *Client) downloadChunked(ctx context.Context, req *http.Request, w *io.P
424425
for i := 0; i < len(chunks); i++ {
425426
select {
426427
case <-cancelCtx.Done():
428+
abort(context.Canceled)
427429
return
428430
case data := <-chunks[i].data:
429431
_, err := io.Copy(w, bytes.NewBuffer(data))

client_test.go

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,15 @@ func TestGetVideoWithManifestURL(t *testing.T) {
129129
func TestGetStream(t *testing.T) {
130130
assert, require := assert.New(t), require.New(t)
131131

132+
expectedSize := 988479
133+
134+
// Create testclient to enforce re-using of routines
135+
testClient := Client{
136+
Debug: true,
137+
MaxRoutines: 10,
138+
ChunkSize: int64(expectedSize) / 11,
139+
}
140+
132141
// Download should not last longer than a minute.
133142
// Otherwise we assume Youtube is throtteling us.
134143
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
@@ -141,7 +150,7 @@ func TestGetStream(t *testing.T) {
141150

142151
reader, size, err := testClient.GetStreamContext(ctx, video, &video.Formats[0])
143152
require.NoError(err)
144-
assert.EqualValues(988479, size)
153+
assert.EqualValues(expectedSize, size)
145154

146155
data, err := io.ReadAll(reader)
147156
require.NoError(err)

0 commit comments

Comments
 (0)