Skip to content

Commit 587ba7e

Browse files
committed
fix test
1 parent e254569 commit 587ba7e

3 files changed

Lines changed: 61 additions & 67 deletions

File tree

pkg/download/engine/inject/stream/module.go

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package stream
22

33
import (
4+
"context"
45
_ "embed"
56
"errors"
67
"fmt"
@@ -254,8 +255,9 @@ type fetchRegistry struct {
254255
}
255256

256257
type fetchStream struct {
257-
body io.ReadCloser
258-
mu sync.Mutex
258+
body io.ReadCloser
259+
cancel context.CancelFunc
260+
closeOnce sync.Once
259261
}
260262

261263
type fetchRequest struct {
@@ -323,7 +325,9 @@ func (r *fetchRegistry) Open(runtime *goja.Runtime, proxyHandler func(r *http.Re
323325
if err != nil {
324326
return nil, err
325327
}
328+
ctx, cancel := context.WithCancel(context.Background())
326329
reqBuilder := client.R()
330+
reqBuilder.SetContext(ctx)
327331
reqBuilder.DisableAutoReadResponse()
328332
for _, header := range reqMeta.Headers {
329333
reqBuilder.SetHeader(header[0], header[1])
@@ -375,7 +379,7 @@ func (r *fetchRegistry) Open(runtime *goja.Runtime, proxyHandler func(r *http.Re
375379
bodyCloser = resp.Response.Body
376380
}
377381
r.mu.Lock()
378-
r.streams[id] = &fetchStream{body: bodyCloser}
382+
r.streams[id] = &fetchStream{body: bodyCloser, cancel: cancel}
379383
r.mu.Unlock()
380384
return meta, nil
381385
}
@@ -389,9 +393,7 @@ func (r *fetchRegistry) Read(id string, chunkSize int) ([]byte, bool, error) {
389393
chunkSize = 64 * 1024
390394
}
391395
buf := make([]byte, chunkSize)
392-
stream.mu.Lock()
393396
n, err := stream.body.Read(buf)
394-
stream.mu.Unlock()
395397
if n > 0 {
396398
return buf[:n], false, nil
397399
}
@@ -412,9 +414,12 @@ func (r *fetchRegistry) Close(id string) {
412414
delete(r.streams, id)
413415
r.mu.Unlock()
414416
if stream != nil {
415-
stream.mu.Lock()
416-
_ = stream.body.Close()
417-
stream.mu.Unlock()
417+
stream.closeOnce.Do(func() {
418+
if stream.cancel != nil {
419+
stream.cancel()
420+
}
421+
_ = stream.body.Close()
422+
})
418423
}
419424
}
420425

pkg/download/engine/inject/stream/stream.js

Lines changed: 27 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -580,48 +580,6 @@
580580
};
581581
}
582582

583-
function attachWritableObjectURL(url, value) {
584-
value._addObserver({
585-
write(chunk) {
586-
return Promise.resolve(writeWritableObjectURL(url, chunk));
587-
},
588-
close() {
589-
return Promise.resolve(closeWritableObjectURL(url));
590-
},
591-
abort(reason) {
592-
return Promise.resolve(abortWritableObjectURL(url, reason == null ? "" : String(reason)));
593-
}
594-
});
595-
value.__gopeedObjectURL = url;
596-
return url;
597-
}
598-
599-
function pipeReadableToWritable(readable, writable) {
600-
const writer = writable.getWriter();
601-
void (async () => {
602-
const reader = readable.getReader();
603-
try {
604-
while (true) {
605-
const { done, value } = await reader.read();
606-
if (done) {
607-
break;
608-
}
609-
await writer.write(value);
610-
}
611-
await writer.close();
612-
} catch (error) {
613-
try {
614-
await writer.abort(error);
615-
} catch (_) {
616-
}
617-
} finally {
618-
if (typeof reader.releaseLock === "function") {
619-
reader.releaseLock();
620-
}
621-
}
622-
})();
623-
}
624-
625583
function startReadableObjectURL(url, state) {
626584
activeReadableObjectURLs.set(url, state);
627585
setTimeout(() => {
@@ -661,6 +619,7 @@
661619
}
662620

663621
async function pumpReadableObjectURL(url, state) {
622+
let reader = null;
664623
try {
665624
let source;
666625
if (state.initialReadable) {
@@ -672,9 +631,10 @@
672631
}
673632
source = await state.openReadable(state.offset);
674633
}
675-
const reader = toReadableStreamReader(source, state.sourceLabel);
634+
reader = toReadableStreamReader(source, state.sourceLabel);
676635
if (activeReadableObjectURLs.get(url) !== state || state.cancelled) {
677636
releaseReader(reader, "stale gblob producer");
637+
reader = null;
678638
return;
679639
}
680640
state.reader = reader;
@@ -683,6 +643,7 @@
683643
const current = activeReadableObjectURLs.get(url);
684644
if (current !== state) {
685645
releaseReader(reader, "stale gblob producer");
646+
reader = null;
686647
return;
687648
}
688649
const { done, value } = await reader.read();
@@ -692,11 +653,13 @@
692653
await closeWritableObjectURL(url);
693654
}
694655
releaseReader(reader);
656+
reader = null;
695657
return;
696658
}
697659
const chunk = value instanceof Uint8Array ? value : new Uint8Array(value);
698660
if (activeReadableObjectURLs.get(url) !== state || state.cancelled) {
699661
releaseReader(reader, "stale gblob producer");
662+
reader = null;
700663
return;
701664
}
702665
await writeWritableObjectURL(url, chunk);
@@ -707,8 +670,16 @@
707670
}
708671
}
709672
releaseReader(reader, "gblob producer cancelled");
673+
reader = null;
710674
} catch (error) {
711675
if (isIgnorableGBlobObjectURLError(error)) {
676+
if (activeReadableObjectURLs.get(url) === state) {
677+
activeReadableObjectURLs.delete(url);
678+
}
679+
if (reader) {
680+
releaseReader(reader, "gblob source closed");
681+
reader = null;
682+
}
712683
return;
713684
}
714685
if (activeReadableObjectURLs.get(url) === state) {
@@ -718,6 +689,10 @@
718689
} catch (_) {
719690
}
720691
}
692+
if (reader) {
693+
releaseReader(reader, error);
694+
reader = null;
695+
}
721696
}
722697
}
723698

@@ -743,9 +718,15 @@
743718
return createBlobObjectURL(described.value._buffer, described.value.type || "");
744719
}
745720
if (described.kind === "readable") {
746-
const transform = new TransformStream();
747-
const url = attachWritableObjectURL(createWritableObjectURL(false), transform.writable);
748-
pipeReadableToWritable(described.initialReadable, transform.writable);
721+
const url = createWritableObjectURL(false);
722+
startReadableObjectURL(url, {
723+
initialReadable: described.initialReadable,
724+
openReadable: null,
725+
sourceLabel: described.sourceLabel,
726+
offset: 0,
727+
reader: null,
728+
cancelled: false,
729+
});
749730
return url;
750731
}
751732
if (described.kind === "opener") {

pkg/download/extension_test.go

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -397,15 +397,17 @@ func TestDownloader_Extension_GBlobHTTPStreamProxyReportsDownloadedBeforeComplet
397397
}
398398

399399
payload := strings.Repeat("gopeed-progress-", 64*1024)
400-
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
401-
if r.Method == http.MethodHead {
402-
w.Header().Set("Content-Length", strconv.Itoa(len(payload)))
403-
return
404-
}
400+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
401+
if r.Method == http.MethodHead {
402+
w.Header().Set("Connection", "close")
403+
w.Header().Set("Content-Length", strconv.Itoa(len(payload)))
404+
return
405+
}
405406

406-
w.Header().Set("Content-Type", "application/octet-stream")
407-
w.Header().Set("Content-Length", strconv.Itoa(len(payload)))
408-
flusher, _ := w.(http.Flusher)
407+
w.Header().Set("Connection", "close")
408+
w.Header().Set("Content-Type", "application/octet-stream")
409+
w.Header().Set("Content-Length", strconv.Itoa(len(payload)))
410+
flusher, _ := w.(http.Flusher)
409411
chunkSize := 4096
410412
for start := 0; start < len(payload); start += chunkSize {
411413
end := start + chunkSize
@@ -588,10 +590,12 @@ func TestDownloader_Extension_GBlobHTTPStreamDeleteWhileDownloading(t *testing.T
588590
if flusher != nil {
589591
flusher.Flush()
590592
}
591-
time.Sleep(20 * time.Millisecond)
592-
}
593-
}))
594-
defer server.Close()
593+
time.Sleep(20 * time.Millisecond)
594+
}
595+
}))
596+
server.Config.SetKeepAlivesEnabled(false)
597+
defer server.Close()
598+
defer server.CloseClientConnections()
595599

596600
rr, err := downloader.Resolve(&base.Request{
597601
URL: "https://example.com/http-stream?target=" + server.URL + "&name=delete.bin",
@@ -1945,7 +1949,11 @@ func waitForTaskTerminal(t *testing.T, downloader *Downloader, id string, timeou
19451949
t.Fatal(err)
19461950
}
19471951
case <-time.After(timeout):
1948-
t.Fatalf("timeout waiting for task %s", id)
1952+
task := downloader.GetTask(id)
1953+
if task == nil {
1954+
t.Fatalf("timeout waiting for task %s: task not found", id)
1955+
}
1956+
t.Fatalf("timeout waiting for task %s: status=%s downloaded=%d total=%d", id, task.Status, task.Progress.Downloaded, task.Meta.Res.Size)
19491957
}
19501958
}
19511959

0 commit comments

Comments
 (0)