~mendelmaleh/download

ref: d367236e0766f5511a5845f6fd923c6742c0c8e3 download/tar.go -rw-r--r-- 1.4 KiB
d367236eMendel E Initial commit 1 year, 1 month ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
package download

import (
	"archive/tar"
	"errors"
	"io"
	"os"
	"sync"
)

var (
	ErrNoData       = errors.New("nothing to download")
	ErrNoGoroutines = errors.New("MaxGoroutines should be positive")
)

// Tar downloads data to a tarball at file.
func (dl *DL) Tar(data Interface, file *os.File, opt *Options) error {
	// set default options
	if opt == nil {
		opt = DefaultOptions
	}

	// ensure opt.MaxGoroutines is positive, otherwise it will hang
	// forever in concurrent mode.
	if opt.Concurrent && opt.MaxGoroutines <= 0 {
		return ErrNoGoroutines
	}

	// get files
	files := data.URLs(opt.Separator)
	if len(files) == 0 {
		return ErrNoData
	}

	// prep tar
	archive := tar.NewWriter(file)
	defer archive.Close()

	// download not concurrent
	if !opt.Concurrent {
		for _, f := range files {
			dl.tar(f, archive, opt)
		}

		return nil
	}

	// download concurrent
	var wg sync.WaitGroup
	defer wg.Wait()

	mg := make(chan bool, opt.MaxGoroutines)

	for _, f := range files {
		mg <- true
		wg.Add(1)

		go func(f File) {
			defer wg.Done()
			defer func() { <-mg }()

			dl.tar(f, archive, opt)
		}(f)
	}

	return nil
}

func (dl *DL) tar(f File, a *tar.Writer, opt *Options) {
	buf, err := dl.try(f.URL, opt.MaxAttempts)
	if err != nil {
		panic(err)
	}

	dl.Write.Lock()
	defer dl.Write.Unlock()

	a.WriteHeader(&tar.Header{
		Name: f.Path,
		Mode: 0600,
		Size: int64(buf.Len()),
	})

	_, err = io.Copy(a, buf)
	if err != nil {
		panic(err)
	}
}