parallel downloader take 2

lsdl2.go

// Copyright (C) 2022 Storj Labs, Inc.
// See LICENSE for copying information.

package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"

	"storj.io/common/sync2"
)

var url string
var concurrency int
var rangeSize int64

func init() {
	flag.IntVar(&concurrency, "concurrency", 4, "number of ranges to download at once")
	flag.Int64Var(&rangeSize, "range-size", 64*1000*1000, "bytes to download per range")
	flag.Parse()
	url = flag.Arg(0)
	if url == "" {
		programName, _ := os.Executable()
		fmt.Printf("Usage: %s [OPTIONS]\n", filepath.Base(programName))
		flag.PrintDefaults()
		os.Exit(1)
	}
}

func main() {
	ctx, cancel := context.WithCancel(context.Background())

	wait := make(chan struct{})
	close(wait)

	limiter := sync2.NewLimiter(concurrency)

	for j := 0; ; j++ {
		thisWait := wait
		done := make(chan struct{})
		i := j
		ok := limiter.Go(ctx, func() {
			defer close(done)

			status, data, err := downloadRange(int64(i))
			if err != nil || status > 206 {
				fmt.Fprintf(os.Stderr, "error on job %d: %+v\n", i, err)
				cancel()
				return
			}
			<-thisWait
			os.Stdout.Write(data)
		})
		if !ok {
			break
		}
		wait = done
	}

}

func downloadRange(rangeNum int64) (int, []byte, error) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return 500, []byte{}, err
	}
	startRange := rangeNum * rangeSize
	endRange := (rangeNum+1)*rangeSize - 1
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", startRange, endRange))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return 500, []byte{}, err
	}
	fmt.Fprintf(os.Stderr, "starting job %d [%d-%d] status %d\n", rangeNum, startRange, endRange, resp.StatusCode)

	defer resp.Body.Close()
	data, err := ioutil.ReadAll(resp.Body)
	return resp.StatusCode, data, err
}