fastget.go

package fastget

import (
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"time"

	"golang.org/x/net/context/ctxhttp"
	"golang.org/x/sync/errgroup"
)

// FastGetter Represents the information required to fastget a file url
type FastGetter struct {
	// URL of the file to be downloaded
	URL string
	// Number of concurrent worked should be used to download file
	Workers int
	// Path of output file to write download
	OutputFile string
	// Headers to be included to while making requests
	Headers map[string]string
	// OnBeforeStart to be called before even download start
	OnBeforeStart func(int64, int64)

	// OnStart to be called on started downloading a chunk / a part
	OnStart func(int, int64)
	// OnProgress to be called on change in progress of downloading a chunk / a part
	OnProgress func(int, int64)
	// OnFinish to be called on finished downloading a chunk / a part
	OnFinish func(int)
}

// chunkInfo holds the information about a chunk to be downloaded
type chunkInfo struct {
	ctx      context.Context
	client   *http.Client
	output   io.WriterAt
	url      string
	off, lim int64
	wid      int
	headers  map[string]string
}

// Result represents the result of fastget
type Result struct {
	URL         string
	Size        int64
	OutputFile  *os.File
	ElapsedTime time.Duration
}

// NewFastGetter creates and returns an instance of FastGetter
func NewFastGetter(url string) (*FastGetter, error) {
	fg := &FastGetter{
		URL:        url,
		Workers:    3,
		OutputFile: path.Base(url),
		Headers:    make(map[string]string),
	}
	return fg, nil
}

// Get ultrafast downloads the file
func (fg *FastGetter) Get() (*Result, error) {
	return fg.get()
}

func (fg *FastGetter) get() (*Result, error) {
	canFastGet, length, err := fg.checkEligibility()
	if err != nil {
		return nil, err
	}
	if !canFastGet {
		log.Println("WARN: FileURL doesn't support parellel download.")
		fg.Workers = 1
	}

	chunkLen := int64(length / int64(fg.Workers))

	if fg.OnBeforeStart != nil {
		fg.OnBeforeStart(length, chunkLen)
	}

	ctx := context.Background()
	client := http.DefaultClient

	output, err := os.OpenFile(fg.OutputFile, os.O_CREATE|os.O_RDWR, 0666)
	if err != nil {
		return nil, err
	}

	wg, ctx := errgroup.WithContext(ctx)
	startTime := time.Now()

	var start, end int64
	for wid := 1; wid <= fg.Workers; wid++ {

		if wid == fg.Workers {
			end = length // last part
		} else {
			end = start + chunkLen
		}

		wid := wid
		off := start
		lim := end

		wg.Go(func() error {
			return fg.getChunk(&chunkInfo{
				ctx: ctx, client: client, output: output, url: fg.URL,
				off: off, lim: lim, wid: wid, headers: fg.Headers,
			})
		})

		start = end
	}

	err = wg.Wait()
	if err != nil {
		return nil, err
	}
	elapsed := time.Since(startTime)

	result := &Result{
		URL:         fg.URL,
		Size:        length,
		OutputFile:  output,
		ElapsedTime: elapsed,
	}
	return result, nil
}

func (fg FastGetter) checkEligibility() (bool, int64, error) {
	req, err := http.NewRequest("HEAD", fg.URL, nil)
	if err != nil {
		return false, 0, err
	}
	for key, value := range fg.Headers {
		req.Header.Add(key, value)
	}
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return false, 0, err
	}
	defer res.Body.Close()
	acceptRanges := res.Header.Get("Accept-Ranges") == "bytes"
	length := res.ContentLength

	return acceptRanges, length, nil
}

// makeRequest performs the HTTP request and returns the response object
func (cInfo chunkInfo) makeRequest() (*http.Response, error) {
	req, err := http.NewRequest("GET", cInfo.url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", cInfo.off, cInfo.lim))
	// Add custom headers
	for key, value := range cInfo.headers {
		req.Header.Add(key, value)
	}

	resp, err := ctxhttp.Do(cInfo.ctx, cInfo.client, req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// getChunk fetches the part of file to download
func (fg FastGetter) getChunk(cInfo *chunkInfo) error {
	if fg.OnStart != nil {
		fg.OnStart(cInfo.wid, cInfo.lim-cInfo.off)
	}
	resp, err := cInfo.makeRequest()
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusPartialContent {
		return fmt.Errorf("server responded with %d status code, expected %d", resp.StatusCode, http.StatusPartialContent)
	}

	var written int64
	contentLen := resp.ContentLength

	buf := make([]byte, 1*1024*1024)
	for {
		nr, er := resp.Body.Read(buf)

		if nr > 0 {
			nw, err := cInfo.output.WriteAt(buf[0:nr], cInfo.off)
			if err != nil {
				return fmt.Errorf("error writing chunk. %s", err.Error())
			}
			if nr != nw {
				return fmt.Errorf("error writing chunk. written %d, but expected %d", nw, nr)
			}

			cInfo.off = int64(nw) + cInfo.off
			if nw > 0 {
				written += int64(nw)
			}
			if fg.OnProgress != nil {
				fg.OnProgress(cInfo.wid, written)
			}
		}

		if er != nil {
			if er.Error() == "EOF" {
				if contentLen == written {
					// Download successfully
				} else {
					return fmt.Errorf("error reading response. %s", er.Error())
				}
				break
			}
			return er
		}
	}
	if fg.OnFinish != nil {
		fg.OnFinish(cInfo.wid)
	}
	return nil
}