github.com/laof/lite-speed-test@v0.0.0-20230930011949-1f39b7037845/download/download_range.go (about)

     1  package download
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  	"net/http"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/laof/lite-speed-test/common/pool"
    12  )
    13  
    14  var (
    15  	contentLength = 242743296
    16  )
    17  
    18  func DownloadRange(ctx context.Context, link string, part int, timeout time.Duration, handshakeTimeout time.Duration, resultChan chan<- int64, startChan chan<- time.Time) (int64, error) {
    19  	client, err := createClient(ctx, link)
    20  	if err != nil {
    21  		return 0, err
    22  	}
    23  
    24  	option := DownloadOption{
    25  		DownloadTimeout:  timeout,
    26  		HandshakeTimeout: handshakeTimeout,
    27  		URL:              downloadLink,
    28  		Ranges:           calcRange(int64(part), int64(contentLength), link),
    29  	}
    30  	return downloadRangeInternal(ctx, option, resultChan, startChan, client.Dial)
    31  }
    32  
    33  func downloadRangeInternal(ctx context.Context, option DownloadOption, resultChan chan<- int64, startOuterChan chan<- time.Time, dial func(network, addr string) (net.Conn, error)) (int64, error) {
    34  	var max int64 = 0
    35  	var wg sync.WaitGroup
    36  	totalCh := make(chan int64)
    37  	// remove
    38  	errorCh := make(chan error)
    39  	startCh := make(chan time.Time, 1)
    40  	for _, rng := range option.Ranges {
    41  		wg.Add(1)
    42  		go func(rng Range, totalChan chan<- int64, errorChan chan<- error, startChan chan<- time.Time) (int64, error) {
    43  			defer wg.Done()
    44  			var max int64 = 0
    45  			httpTransport := &http.Transport{}
    46  			httpClient := &http.Client{Transport: httpTransport, Timeout: option.HandshakeTimeout}
    47  			if dial != nil {
    48  				httpTransport.Dial = dial
    49  			}
    50  			req, err := http.NewRequest("GET", option.URL, nil)
    51  			if err != nil {
    52  				errorChan <- err
    53  				return max, err
    54  			}
    55  			// add range
    56  			ranges := rng.toHeader(int64(contentLength))
    57  			req.Header.Add("Range", ranges)
    58  			response, err := httpClient.Do(req)
    59  			if err != nil {
    60  				errorChan <- err
    61  				return max, err
    62  			}
    63  			defer response.Body.Close()
    64  			prev := time.Now()
    65  			startChan <- prev
    66  			var total int64
    67  			buf := pool.Get(20 * 1024)
    68  			pool.Put(buf)
    69  		Loop:
    70  			for {
    71  				if ctx.Err() != nil {
    72  					break
    73  				}
    74  				select {
    75  				case <-ctx.Done():
    76  					return max, err
    77  				default:
    78  					// buf := pool.Get(20 * 1024)
    79  					nr, er := response.Body.Read(buf)
    80  					total += int64(nr)
    81  					// pool.Put(buf)
    82  					now := time.Now()
    83  					if now.Sub(prev) >= 100*time.Millisecond || er != nil {
    84  						prev = now
    85  						if totalChan != nil {
    86  							totalChan <- total
    87  						}
    88  						if max < total {
    89  							max = total
    90  						}
    91  						total = 0
    92  					}
    93  					if er != nil {
    94  						if er != io.EOF {
    95  							errorChan <- err
    96  							err = er
    97  						}
    98  						break Loop
    99  					}
   100  				}
   101  
   102  			}
   103  			return max, nil
   104  
   105  		}(rng, totalCh, errorCh, startCh)
   106  	}
   107  	var sum int64 = 0
   108  	var errorResult error = nil
   109  
   110  	doneCh := make(chan bool, 1)
   111  	go func(doneChan chan<- bool) {
   112  		wg.Wait()
   113  		doneChan <- true
   114  	}(doneCh)
   115  	var prev time.Time
   116  	for {
   117  		if !prev.IsZero() {
   118  			now := time.Now()
   119  			if now.Sub(prev) >= time.Second {
   120  				prev = now
   121  				if resultChan != nil {
   122  					resultChan <- sum
   123  				}
   124  				if max < sum {
   125  					max = sum
   126  				}
   127  				sum = 0
   128  			}
   129  		}
   130  		select {
   131  		case total := <-totalCh:
   132  			if total < 0 {
   133  				return max, nil
   134  			}
   135  			sum += total
   136  		case err := <-errorCh:
   137  			if err != nil {
   138  				errorResult = err
   139  			}
   140  		case start := <-startCh:
   141  			// init only once
   142  			if prev.IsZero() {
   143  				prev = start
   144  				if startOuterChan != nil {
   145  					startOuterChan <- start
   146  				}
   147  			}
   148  		case <-doneCh:
   149  			return max, errorResult
   150  		case <-ctx.Done():
   151  			err := ctx.Err()
   152  			if max > 0 {
   153  				err = nil
   154  			}
   155  			return max, err
   156  		}
   157  	}
   158  }