mirror of
https://github.com/simon-ding/polaris.git
synced 2026-03-08 10:40:45 +08:00
WIP: upload with progress
This commit is contained in:
86
pkg/uploader/uploader.go
Normal file
86
pkg/uploader/uploader.go
Normal 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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user