192 lines
3.6 KiB
Go
192 lines
3.6 KiB
Go
package internal
|
|
|
|
import (
|
|
"archive/zip"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
func humanFormat(bytes int64) string {
|
|
n := float64(bytes)
|
|
for _, unit := range []string{"", "Ki", "Mi", "Gi"} {
|
|
if math.Abs(n) < 1024.0 {
|
|
return fmt.Sprintf("%3.1f%sB", n, unit)
|
|
}
|
|
n /= 1024.0
|
|
}
|
|
return fmt.Sprintf("%.1fTiB", n)
|
|
}
|
|
|
|
type WriteCounter struct {
|
|
Total uint64
|
|
StartTime time.Time
|
|
OnEmit func(percent float64, speedMBps float64)
|
|
TotalSize int64
|
|
lastLoggedPercent int
|
|
}
|
|
|
|
func NewWriteCounter(total int64, onEmit func(percent float64, speedMBps float64)) *WriteCounter {
|
|
return &WriteCounter{
|
|
StartTime: time.Now(),
|
|
TotalSize: total,
|
|
lastLoggedPercent: -1,
|
|
OnEmit: onEmit,
|
|
}
|
|
}
|
|
|
|
func (wc *WriteCounter) Write(p []byte) (int, error) {
|
|
n := len(p)
|
|
wc.Total += uint64(n)
|
|
wc.PrintProgress()
|
|
return n, nil
|
|
}
|
|
|
|
func (wc *WriteCounter) PrintProgress() {
|
|
elapsed := time.Since(wc.StartTime).Seconds()
|
|
if elapsed < 0.001 {
|
|
elapsed = 0.001
|
|
}
|
|
|
|
speed := float64(wc.Total) / 1024 / 1024 / elapsed // MB/s
|
|
percent := float64(wc.Total) / float64(wc.TotalSize) * 100
|
|
if wc.OnEmit != nil {
|
|
wc.OnEmit(percent, speed)
|
|
}
|
|
}
|
|
|
|
func DownloadFile(filepath string, url string, onEmit func(percent float64, speed float64)) error {
|
|
tmpPath := filepath + ".tmp"
|
|
|
|
resp, err := http.Get(url)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get file: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return fmt.Errorf("bad status: %s", resp.Status)
|
|
}
|
|
|
|
out, err := os.Create(tmpPath)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create tmp file: %w", err)
|
|
}
|
|
|
|
counter := NewWriteCounter(resp.ContentLength, onEmit)
|
|
_, err = io.Copy(out, io.TeeReader(resp.Body, counter))
|
|
if closeErr := out.Close(); closeErr != nil {
|
|
return fmt.Errorf("failed to close tmp file: %w", closeErr)
|
|
}
|
|
if err != nil {
|
|
return fmt.Errorf("failed to download file: %w", err)
|
|
}
|
|
|
|
// Delete destination file if it exists
|
|
if _, err := os.Stat(filepath); err == nil {
|
|
if err := os.Remove(filepath); err != nil {
|
|
return fmt.Errorf("failed to remove existing file: %w", err)
|
|
}
|
|
}
|
|
|
|
for i := 0; i < 3; i++ {
|
|
err = os.Rename(tmpPath, filepath)
|
|
if err == nil {
|
|
break
|
|
}
|
|
time.Sleep(300 * time.Millisecond)
|
|
}
|
|
if err != nil {
|
|
return fmt.Errorf("failed to rename after retries: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
|
|
|
|
func unzipParallel(src string, dest string) error {
|
|
numCPU := runtime.NumCPU()
|
|
|
|
reserved := 1
|
|
if numCPU > 4 {
|
|
reserved = 2
|
|
}
|
|
maxWorkers := numCPU - reserved
|
|
if maxWorkers < 1 {
|
|
maxWorkers = 1
|
|
}
|
|
r, err := zip.OpenReader(src)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer r.Close()
|
|
|
|
err = os.MkdirAll(dest, 0755)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
type job struct {
|
|
f *zip.File
|
|
}
|
|
jobs := make(chan job)
|
|
var wg sync.WaitGroup
|
|
|
|
// Worker pool
|
|
for i := 0; i < maxWorkers; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for j := range jobs {
|
|
err := extractFile(j.f, dest)
|
|
if err != nil {
|
|
fmt.Printf("Error extracting %s: %v\n", j.f.Name, err)
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
// Feed jobs
|
|
for _, f := range r.File {
|
|
jobs <- job{f}
|
|
}
|
|
close(jobs)
|
|
wg.Wait()
|
|
|
|
return nil
|
|
}
|
|
|
|
func extractFile(f *zip.File, dest string) error {
|
|
fp := filepath.Join(dest, f.Name)
|
|
|
|
if f.FileInfo().IsDir() {
|
|
return os.MkdirAll(fp, f.Mode())
|
|
}
|
|
|
|
err := os.MkdirAll(filepath.Dir(fp), 0755)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
rc, err := f.Open()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer rc.Close()
|
|
|
|
out, err := os.OpenFile(fp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer out.Close()
|
|
|
|
_, err = io.Copy(out, rc)
|
|
return err
|
|
} |