基于Go语言实现一个并发下载器

package probe

import (
	"fmt"
	"log"
	"net/http"
	"strconv"
)

type Probe struct {
	workers int
	url     string
}

func NewProbe(worker int, url string) *Probe {
	return &Probe{
		workers: worker,
		url:     url,
	}
}

func (p *Probe) GetFileSize() (int, error) {
	var size = -1

	client := &http.Client{}
	req, err := http.NewRequest("HEAD", p.url, nil)
	if err != nil {
		log.Fatal(err)
	}
	resp, err := client.Do(req)
	if err != nil {
		log.Fatal(err)
	}

	if header, ok := resp.Header["Content-Length"]; ok {
		fileSize, err := strconv.Atoi(header[0])
		if err != nil {
			log.Fatal("File size could not be determined : ", err)
		}

		size = fileSize / p.workers
	} else {
		log.Fatal("File size was not provided!")
		return size, fmt.Errorf("file size was not provided.")
	}
	return size, nil
}

type Downloader struct {
	result  chan Part
	size    int
	workers int
}

type Part struct {
	Data  []byte
	Index int
}

func (d *Downloader) Download(index int, url string) {
	client := &http.Client{}

	// calculate offset by multiplying
	// index with size
	start := index * d.size

	// Write data range in correct format
	// I'm reducing one from the end size to account for
	// the next chunk starting there
	dataRange := fmt.Sprintf("bytes=%d-%d", start, start+d.size-1)

	// if this is downloading the last chunk
	// rewrite the header. It's an easy way to specify
	// getting the rest of the file
	if index == d.workers-1 {
		dataRange = fmt.Sprintf("bytes=%d-", start)
	}
	log.Println(dataRange)

	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		// TODO: restart download
		return
	}
	req.Header.Add("Range", dataRange)
	resp, err := client.Do(req)
	if err != nil {
		// TODO: restart download
		return
	}

	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		// TODO: restart download
		return
	}

	d.result <- Part{Index: index, Data: body}
}

func (d *Downloader) Merge(filename string) error {
	log.Println("start to merge data")

	parts := make([][]byte, d.workers)
	counter := 0
	for part := range d.result {
		counter++
		parts[part.Index] = part.Data
		if counter == d.workers {
			break
		}
	}
	log.Println("sort data as original order")

	file := []byte{}
	for _, part := range parts {
		file = append(file, part...)
	}
	log.Println("write data into buffer array")
	err := ioutil.WriteFile(filename, file, 0777)
	return err
}

package main

import (
	"flag"
	"log"
	"time"

	"go-store/applications/downloader/download"
	"go-store/applications/downloader/probe"
)

var (
	// to test internet
	url = flag.String("url", "http://212.183.159.230/512MB.zip", "download url")
	// number of goroutines to spawn for download.
	workers = flag.Int("worker", 5, "concurrent downloader number")
	// filename for downloaded file
	filename = flag.String("file", "data.zip", "downloaded filename")
)

func main() {
	flag.Parse()
	start := time.Now()
	probe := probe.NewProbe(*workers, *url)
	size, err := probe.GetFileSize()
	if err != nil {
		panic(err)
	}

	results := make(chan download.Part, *workers)
	downloader := download.NewDownloader(results, size, *workers)
	for i := 0; i < *workers; i++ {
		go downloader.Download(i, *url)
	}

	err = downloader.Merge(*filename)
	end := time.Now()
	if err != nil {
		panic(err)
	}

	log.Println("cost time: ", end.Sub(start))
}