291 lines
7.9 KiB
Go
291 lines
7.9 KiB
Go
package playback
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"reflect"
|
|
"strconv"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"litepan/internal/domain"
|
|
)
|
|
|
|
type firstWriteBuffer struct {
|
|
b []byte
|
|
once sync.Once
|
|
first chan struct{}
|
|
}
|
|
|
|
func (w *firstWriteBuffer) Write(p []byte) (int, error) {
|
|
w.b = append(w.b, p...)
|
|
if len(p) > 0 {
|
|
w.once.Do(func() { close(w.first) })
|
|
}
|
|
return len(p), nil
|
|
}
|
|
|
|
func (w *firstWriteBuffer) Bytes() []byte { return w.b }
|
|
|
|
func TestProxyStreamUsesDriverConcurrencyAndPreservesOrder(t *testing.T) {
|
|
const partSize = 64 * 1024
|
|
data := make([]byte, 7*partSize+321)
|
|
for i := range data {
|
|
data[i] = byte(i % 251)
|
|
}
|
|
const streamStart = int64(17)
|
|
streamEnd := int64(len(data) - 29)
|
|
|
|
var mu sync.Mutex
|
|
active := 0
|
|
peak := 0
|
|
ranges := make(map[string]int)
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
start, end, err := parseSingleRange(r.Header.Get("Range"), int64(len(data)))
|
|
if err != nil {
|
|
http.Error(w, err.Error(), http.StatusRequestedRangeNotSatisfiable)
|
|
return
|
|
}
|
|
mu.Lock()
|
|
active++
|
|
if active > peak {
|
|
peak = active
|
|
}
|
|
ranges[r.Header.Get("Range")]++
|
|
mu.Unlock()
|
|
|
|
index := int((start - streamStart) / partSize)
|
|
delays := []time.Duration{45 * time.Millisecond, 5 * time.Millisecond, 20 * time.Millisecond}
|
|
time.Sleep(delays[index%len(delays)])
|
|
span := data[start : end+1]
|
|
w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, len(data)))
|
|
w.Header().Set("Content-Length", strconv.Itoa(len(span)))
|
|
w.WriteHeader(http.StatusPartialContent)
|
|
_, _ = w.Write(span)
|
|
|
|
mu.Lock()
|
|
active--
|
|
mu.Unlock()
|
|
}))
|
|
t.Cleanup(server.Close)
|
|
|
|
svc := &Service{clientHTTP1: server.Client(), clientH2: server.Client()}
|
|
lh := &linkHolder{
|
|
svc: svc,
|
|
link: domain.DownloadInfo{URL: server.URL, ChunkSize: partSize, Concurrency: 3},
|
|
accountID: 7,
|
|
fileID: "file",
|
|
ua: "test",
|
|
}
|
|
var got bytes.Buffer
|
|
if err := svc.streamUpstreamBody(context.Background(), &got, lh, streamStart, streamEnd, partSize); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !bytes.Equal(got.Bytes(), data[streamStart:streamEnd+1]) {
|
|
t.Fatal("并行代理返回内容顺序错误")
|
|
}
|
|
|
|
wantRanges := make(map[string]int)
|
|
for start := streamStart; start <= streamEnd; start += partSize {
|
|
end := start + partSize - 1
|
|
if end > streamEnd {
|
|
end = streamEnd
|
|
}
|
|
wantRanges[fmt.Sprintf("bytes=%d-%d", start, end)] = 1
|
|
}
|
|
mu.Lock()
|
|
gotPeak := peak
|
|
gotRanges := ranges
|
|
mu.Unlock()
|
|
if gotPeak != 3 {
|
|
t.Fatalf("代理峰值并发=%d，期望 3", gotPeak)
|
|
}
|
|
if !reflect.DeepEqual(gotRanges, wantRanges) {
|
|
t.Fatalf("代理 Range=%v，期望 %v", gotRanges, wantRanges)
|
|
}
|
|
}
|
|
|
|
func TestProxyStreamWritesFirstBytesBeforePrefetchCompletes(t *testing.T) {
|
|
const partSize = 64 * 1024
|
|
data := make([]byte, 3*partSize)
|
|
for i := range data {
|
|
data[i] = byte(i % 239)
|
|
}
|
|
release := make(chan struct{})
|
|
var releaseOnce sync.Once
|
|
firstHandler := make(chan struct{})
|
|
prefixFlushed := make(chan struct{})
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
start, end, err := parseSingleRange(r.Header.Get("Range"), int64(len(data)))
|
|
if err != nil {
|
|
http.Error(w, err.Error(), http.StatusRequestedRangeNotSatisfiable)
|
|
return
|
|
}
|
|
span := data[start : end+1]
|
|
w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, len(data)))
|
|
w.Header().Set("Content-Length", strconv.Itoa(len(span)))
|
|
w.WriteHeader(http.StatusPartialContent)
|
|
if start == 0 {
|
|
close(firstHandler)
|
|
_, _ = w.Write(span[:1024])
|
|
w.(http.Flusher).Flush()
|
|
close(prefixFlushed)
|
|
<-release
|
|
_, _ = w.Write(span[1024:])
|
|
return
|
|
}
|
|
<-release
|
|
_, _ = w.Write(span)
|
|
}))
|
|
t.Cleanup(func() {
|
|
releaseOnce.Do(func() { close(release) })
|
|
server.Close()
|
|
})
|
|
|
|
svc := &Service{clientHTTP1: server.Client(), clientH2: server.Client()}
|
|
lh := &linkHolder{
|
|
svc: svc,
|
|
link: domain.DownloadInfo{URL: server.URL, ChunkSize: partSize, Concurrency: 3},
|
|
accountID: 10,
|
|
fileID: "file",
|
|
ua: "test",
|
|
}
|
|
got := &firstWriteBuffer{first: make(chan struct{})}
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
done <- svc.streamUpstreamBody(context.Background(), got, lh, 0, int64(len(data)-1), partSize)
|
|
}()
|
|
select {
|
|
case <-firstHandler:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("首片上游请求没有启动")
|
|
}
|
|
select {
|
|
case <-prefixFlushed:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("测试上游没有刷出首片前缀")
|
|
}
|
|
select {
|
|
case <-got.first:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("代理等待完整预取后才写出首字节")
|
|
}
|
|
releaseOnce.Do(func() { close(release) })
|
|
select {
|
|
case err := <-done:
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("释放上游后代理没有完成")
|
|
}
|
|
if !bytes.Equal(got.Bytes(), data) {
|
|
t.Fatal("首片流式写出后内容顺序错误")
|
|
}
|
|
}
|
|
|
|
func TestProxyStreamCancellationStopsOutstandingRanges(t *testing.T) {
|
|
const partSize = 64 * 1024
|
|
started := make(chan struct{}, 3)
|
|
canceled := make(chan struct{}, 3)
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Range", "bytes 0-65535/196608")
|
|
w.Header().Set("Content-Length", strconv.Itoa(partSize))
|
|
w.WriteHeader(http.StatusPartialContent)
|
|
w.(http.Flusher).Flush()
|
|
started <- struct{}{}
|
|
<-r.Context().Done()
|
|
canceled <- struct{}{}
|
|
}))
|
|
t.Cleanup(server.Close)
|
|
|
|
svc := &Service{clientHTTP1: server.Client(), clientH2: server.Client()}
|
|
lh := &linkHolder{
|
|
svc: svc,
|
|
link: domain.DownloadInfo{URL: server.URL, ChunkSize: partSize, Concurrency: 3},
|
|
accountID: 8,
|
|
fileID: "file",
|
|
ua: "test",
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
done <- svc.streamUpstreamBody(ctx, &bytes.Buffer{}, lh, 0, 3*partSize-1, partSize)
|
|
}()
|
|
for range 3 {
|
|
select {
|
|
case <-started:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("并行代理请求未全部启动")
|
|
}
|
|
}
|
|
cancel()
|
|
select {
|
|
case err := <-done:
|
|
if err == nil {
|
|
t.Fatal("取消后代理返回 nil 错误")
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("取消后代理没有及时退出")
|
|
}
|
|
for range 3 {
|
|
select {
|
|
case <-canceled:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("取消后仍有上游 Range 未退出")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestProxyStreamRejectsIncompleteRange(t *testing.T) {
|
|
const size = 64 * 1024
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
w.WriteHeader(http.StatusPartialContent)
|
|
_, _ = w.Write(make([]byte, size/2))
|
|
}))
|
|
t.Cleanup(server.Close)
|
|
|
|
svc := &Service{clientHTTP1: server.Client(), clientH2: server.Client()}
|
|
lh := &linkHolder{
|
|
svc: svc,
|
|
link: domain.DownloadInfo{URL: server.URL, ChunkSize: size, Concurrency: 1},
|
|
accountID: 9,
|
|
fileID: "file",
|
|
ua: "test",
|
|
}
|
|
if err := svc.streamUpstreamBody(context.Background(), &bytes.Buffer{}, lh, 0, size-1, size); err == nil {
|
|
t.Fatal("上游 Range 缺失一半数据时未返回错误")
|
|
}
|
|
}
|
|
|
|
func TestProxyStreamDoesNotWriteBytesPastRange(t *testing.T) {
|
|
const size = 64 * 1024
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
w.Header().Set("Content-Length", strconv.Itoa(size+1))
|
|
w.WriteHeader(http.StatusPartialContent)
|
|
_, _ = w.Write(make([]byte, size+1))
|
|
}))
|
|
t.Cleanup(server.Close)
|
|
|
|
svc := &Service{clientHTTP1: server.Client(), clientH2: server.Client()}
|
|
lh := &linkHolder{
|
|
svc: svc,
|
|
link: domain.DownloadInfo{URL: server.URL, ChunkSize: size, Concurrency: 1},
|
|
accountID: 11,
|
|
fileID: "file",
|
|
ua: "test",
|
|
}
|
|
var got bytes.Buffer
|
|
if err := svc.streamUpstreamBody(context.Background(), &got, lh, 0, size-1, size); err == nil {
|
|
t.Fatal("上游返回超出 Range 的数据时未报错")
|
|
}
|
|
if got.Len() != size {
|
|
t.Fatalf("代理向客户端写出 %d 字节，期望严格限制为 %d", got.Len(), size)
|
|
}
|
|
}
|