richkago/richkago.go

package richkago

import (
	"errors"
	"fmt"
	"golang.org/x/sync/errgroup"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"
)

const (
	version        = "Alpha/0.0.1"
	userAgent      = "Richkago" + version
	coroutineLimit = 10
	sliceThreshold = 10 * 1024 * 1024 // 10 MiB
	timeout        = 3 * time.Second
	retryTimes     = 5
	chunkSize      = 102400
)

// readWithTimeout read data from source with timeout limit
func readWithTimeout(reader io.Reader, buffer []byte, controller *Controller, file *os.File, downloaded *int64, chunkID string) error {
	// Make a error chan
	ch := make(chan error, 1)

	go func() {
		// Read data
		n, err := reader.Read(buffer)
		if n > 0 {
			// Write to file
			_, err = file.Write(buffer[:n])
			if err != nil {
				ch <- err
				return
			}
			// Calc amount of downloaded data
			*downloaded += int64(n)
		}
		ch <- err
	}()

	// Error and timeout handler
	select {
	case err := <-ch:
		// Update amount of downloaded data
		if controller != nil {
			controller.UpdateProgress(*downloaded, chunkID)
		}

		return err
	case <-time.After(timeout):
		return errors.New("timeout while reading data")
	}
}

// downloadRange download a file with many slices
func downloadRange(client *http.Client, url string, start, end int64, destination string, controller *Controller) error {
	// Build request header
	headers := map[string]string{
		"User-Agent": userAgent,
		"Range":      fmt.Sprintf("bytes=%d-%d", start, end),
	}

	retries := retryTimes
	var err error

	for retries > 0 {
		// Get header info
		var req *http.Request
		req, err = http.NewRequest("GET", url, nil)
		if err != nil {
			retries--
			time.Sleep(1 * time.Second)
			continue
		}
		for k, v := range headers {
			req.Header.Add(k, v)
		}

		// Read http header
		var resp *http.Response
		resp, err = client.Do(req)
		if err != nil || resp.StatusCode != http.StatusPartialContent {
			retries--
			time.Sleep(1 * time.Second)
			continue
		}

		defer func(Body io.ReadCloser) {
			err = Body.Close()
			if err != nil {
				panic(err)
			}
		}(resp.Body)

		// Pre-create file
		file, _ := os.OpenFile(destination, os.O_WRONLY, 0644)
		defer func(file *os.File) {
			err = file.Close()
			if err != nil {
				panic(err)
			}
		}(file)
		_, err = file.Seek(start, io.SeekStart)
		if err != nil {
			retries--
			time.Sleep(1 * time.Second)
			continue
		}

		// Create goroutines to download
		buffer := make([]byte, chunkSize)
		var downloaded int64
		for {
			// Controller pin
			if controller.paused {
				time.Sleep(1 * time.Second)
				continue
			}

			// Start read stream
			err = readWithTimeout(resp.Body, buffer, controller, file, &downloaded, fmt.Sprintf("%d-%d", start, end))
			if err == io.EOF {
				break
			} else if err != nil {
				break
			}
		}

		// Error handler
		if err != nil && err != io.EOF {
			retries--
			time.Sleep(1 * time.Second)
			continue
		}

		return nil
	}

	return err
}

// downloadSingle download a file directly
func downloadSingle(client *http.Client, url, destination string, controller *Controller) error {
	// Build request header
	headers := map[string]string{
		"User-Agent": userAgent,
	}
	retries := retryTimes

	for retries > 0 {
		// Get header info
		req, _ := http.NewRequest("GET", url, nil)
		for k, v := range headers {
			req.Header.Add(k, v)
		}

		// Read http header
		resp, err := client.Do(req)
		if err != nil || resp.StatusCode != http.StatusOK {
			retries--
			time.Sleep(1 * time.Second)
			continue
		}

		defer func(Body io.ReadCloser) {
			err = Body.Close()
			if err != nil {
				panic(err)
			}
		}(resp.Body)

		// Pre-create file
		file, _ := os.OpenFile(destination, os.O_WRONLY, 0644)
		defer func(file *os.File) {
			err = file.Close()
			if err != nil {
				panic(err)
			}
		}(file)

		// Start downloading
		buffer := make([]byte, chunkSize)
		var downloaded int64
		for {
			// Controller pin
			if controller.paused {
				time.Sleep(1 * time.Second)
				continue
			}

			// Start read stream
			var n int
			n, err = resp.Body.Read(buffer)
			if err != nil {
				break
			}
			if n > 0 {
				_, err = file.Write(buffer[:n])
				if err != nil {
					return err
				}
				downloaded += int64(n)
				controller.UpdateProgress(downloaded, "")
			}
		}

		// Error handler
		if err != nil {
			retries--
			time.Sleep(1 * time.Second)
			continue
		} else {
			return err
		}
	}

	return nil
}

// Download main download task creator
func Download(url, destination string, controller *Controller) (float64, int64, error) {
	// Create http client
	client := &http.Client{}

	// Get header info
	req, err := http.NewRequest("HEAD", url, nil)
	if err != nil {
		controller.excepted = true
		return 0, 0, err
	}
	req.Header.Add("User-Agent", userAgent)

	// Read http header
	resp, err := client.Do(req)
	if err != nil {
		controller.excepted = true
		return 0, 0, err
	}
	defer func(Body io.ReadCloser) {
		err = Body.Close()
		if err != nil {
			panic(err)
		}
	}(resp.Body)

	sizeStr := resp.Header.Get("Content-Length")
	fileSize, _ := strconv.ParseInt(sizeStr, 10, 64)
	controller.totalSize = fileSize

	// Calc slices size
	if fileSize <= sliceThreshold {
		file, _ := os.Create(destination)
		err = file.Truncate(fileSize)
		if err != nil {
			controller.excepted = true
			return 0, 0, err
		}
		err = file.Close()
		if err != nil {
			controller.excepted = true
			return 0, 0, err
		}

		// Too small
		startTime := time.Now()
		err = downloadSingle(client, url, destination, controller)
		if err != nil {
			controller.excepted = true
			return 0, 0, err
		}
		endTime := time.Now()
		return endTime.Sub(startTime).Seconds(), fileSize, nil
	}

	// Pre-create file
	partSize := fileSize / coroutineLimit
	file, _ := os.Create(destination)
	err = file.Truncate(fileSize)
	if err != nil {
		controller.excepted = true
		return 0, 0, err
	}
	err = file.Close()
	if err != nil {
		controller.excepted = true
		return 0, 0, err
	}

	// Start download goroutines
	group := new(errgroup.Group)
	for i := 0; i < coroutineLimit; i++ {
		start := int64(i) * partSize
		end := start + partSize - 1
		if i == coroutineLimit-1 {
			end = fileSize - 1
		}
		group.Go(func() error {
			err = downloadRange(client, url, start, end, destination, controller)
			return err
		})
	}

	// Start all tasks
	startTime := time.Now()
	if err = group.Wait(); err != nil {
		controller.excepted = true
		return 0, 0, err
	}
	endTime := time.Now()

	return endTime.Sub(startTime).Seconds(), fileSize, nil
}