WIP: upload with progress

This commit is contained in:
Simon Ding
2024-11-05 18:54:40 +08:00
parent 66a307f202
commit 7b02eeac51
2 changed files with 103 additions and 3 deletions

86
pkg/uploader/uploader.go Normal file
View File

@@ -0,0 +1,86 @@
package uploader
import (
"fmt"
"io"
"os"
"polaris/pkg/utils"
"sync/atomic"
"time"
)
type StreamWriter interface {
WriteStream(path string, stream io.Reader, _ os.FileMode) error
}
type Uploader struct {
sw StreamWriter
progress atomic.Int64
dir string
size int64
}
func NewUploader(dir string, sw StreamWriter) (*Uploader, error) {
size, err := utils.DirSize(dir)
if err != nil {
return nil, err
}
return &Uploader{sw: sw, dir: dir, size: size, progress: atomic.Int64{}}, nil
}
func (u *Uploader) Upload() error {
return nil
}
type ProgressReader struct {
Reader io.Reader
Progress atomic.Int64
Size int64
Name string
Once bool
Done atomic.Bool
}
func (progressReader *ProgressReader) NewLoop() {
ticker := time.NewTicker(time.Second)
var op int64
for range ticker.C {
p := progressReader.Progress.Load()
KB := (p - op) / 1024
var percent int64
if progressReader.Size != 0 {
percent = p * 100 / progressReader.Size
} else {
percent = 100
}
if KB < 1024 {
fmt.Printf("%s: %dKB/s %d%%\n", progressReader.Name, KB, percent)
} else {
fmt.Printf("%s: %.2fMB/s %d%%\n", progressReader.Name, float64(KB)/1024, percent)
}
if progressReader.Done.Load() {
ticker.Stop()
return
}
}
}
func (progressReader *ProgressReader) Read(p []byte) (int, error) {
n, err := progressReader.Reader.Read(p)
progressReader.Progress.Add(int64(n))
if !progressReader.Once {
progressReader.Once = true
go progressReader.NewLoop()
}
if err != nil {
progressReader.Done.Store(true)
}
return n, err
}
func (progressReader *ProgressReader) Close() error {
progressReader.Done.Store(true)
return nil
}

View File

@@ -5,6 +5,7 @@ import (
"net/http"
"os"
"os/exec"
"path/filepath"
"regexp"
"runtime"
"strconv"
@@ -230,7 +231,7 @@ func Link2Magnet(link string) (string, error) {
return http.ErrUseLastResponse //do not follow redirects
},
}
resp, err := client.Get(link)
if err != nil {
return "", errors.Wrap(err, "get link")
@@ -252,7 +253,6 @@ func Link2Magnet(link string) (string, error) {
return mg.String(), nil
}
func MagnetHash(link string) (string, error) {
if mi, err := metainfo.ParseMagnetV2Uri(link); err != nil {
return "", errors.Errorf("magnet link is not valid: %v", err)
@@ -271,4 +271,18 @@ func MagnetHash(link string) (string, error) {
}
return hash, nil
}
}
}
func DirSize(path string) (int64, error) {
var size int64
err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
if err != nil {
return err
}
if !info.IsDir() {
size += info.Size()
}
return err
})
return size, err
}