Merge pull request #33 from krau/dev-stream

impl webdav stream mode & progress callback for stream mode

core/download.go

@@ -8,8 +8,6 @@ import (
 
 	"github.com/celestix/gotgproto/ext"
 	"github.com/duke-git/lancet/v2/fileutil"
-	"github.com/gotd/td/telegram/message/entity"
-	"github.com/gotd/td/telegram/message/styling"
 	"github.com/gotd/td/tg"
 	"github.com/krau/SaveAny-Bot/bot"
 	"github.com/krau/SaveAny-Bot/config"
@@ -56,28 +54,12 @@ func processPendingTask(task *types.Task) error {
 
 	downloadBuider := Downloader.Download(bot.Client.API(), task.File.Location).WithThreads(getTaskThreads(task.File.FileSize))
 
-	// TODO: show progress for stream storage
 	taskStreamStorage, isStreamStorage := taskStorage.(storage.StreamStorage)
 	if config.Cfg.Stream {
 		if !isStreamStorage {
 			logger.L.Warnf("存储 %s 不支持流式上传", taskStorage.Name())
 		} else {
-			entityBuilder := entity.Builder{}
+			text, entities := buildProgressMessageEntity(task, 0, task.StartTime, 0)
-			text := fmt.Sprintf("正在处理下载任务 (流式)\n文件名: %s\n保存路径: %s",
-				task.FileName(),
-				fmt.Sprintf("[%s]:%s", task.StorageName, task.StoragePath),
-			)
-			var entities []tg.MessageEntityClass
-			if err := styling.Perform(&entityBuilder,
-				styling.Plain("正在处理下载任务 (流式)\n文件名: "),
-				styling.Code(task.FileName()),
-				styling.Plain("\n保存路径: "),
-				styling.Code(fmt.Sprintf("[%s]:%s", task.StorageName, task.StoragePath)),
-			); err != nil {
-				logger.L.Errorf("Failed to build entities: %s", err)
-			} else {
-				text, entities = entityBuilder.Complete()
-			}
 			ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
 				Message:     text,
 				Entities:    entities,
@@ -89,7 +71,13 @@ func processPendingTask(task *types.Task) error {
 				return fmt.Errorf("创建上传流失败: %w", err)
 			}
 			defer uploadStream.Close()
-			_, err = downloadBuider.Stream(cancelCtx, uploadStream)
+
+			task.StartTime = time.Now()
+			progressCallback := buildProgressCallback(ctx, task, getProgressUpdateCount(task.File.FileSize))
+
+			progressStream := NewProgressStream(uploadStream, task.File.FileSize, progressCallback)
+
+			_, err = downloadBuider.Stream(cancelCtx, progressStream)
 			if err != nil {
 				return fmt.Errorf("下载文件失败: %w", err)
 			}
@@ -128,5 +116,4 @@ func processPendingTask(task *types.Task) error {
 	})
 
 	return saveFileWithRetry(cancelCtx, task, taskStorage, cacheDestPath)
-
 }

core/utils.go

@@ -3,6 +3,7 @@ package core
 import (
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"path"
 	"time"
@@ -68,19 +69,6 @@ func processPhoto(task *types.Task, taskStorage storage.Storage, cachePath strin
 	return saveFileWithRetry(task.Ctx, task, taskStorage, cachePath)
 }
 
-// func getProgressBar(progress float64, updateCount int) string {
-// 	bar := ""
-// 	barSize := 100 / updateCount
-// 	for i := 0; i < updateCount; i++ {
-// 		if progress >= float64(barSize*(i+1)) {
-// 			bar += "█"
-// 		} else {
-// 			bar += "░"
-// 		}
-// 	}
-// 	return bar
-// }
-
 func cleanCacheFile(destPath string) {
 	if config.Cfg.Temp.CacheTTL > 0 {
 		common.RmFileAfter(destPath, time.Duration(config.Cfg.Temp.CacheTTL)*time.Second)
@@ -233,3 +221,40 @@ func NewTaskLocalFile(filePath string, fileSize int64, progressCallback func(byt
 		callbackInterval: callbackInterval,
 	}, nil
 }
+
+type ProgressStream struct {
+	writer   io.Writer
+	size     int64
+	done     int64
+	callback func(bytesRead, contentLength int64)
+	nextAt   int64
+	interval int64
+}
+
+func (ps *ProgressStream) Write(p []byte) (n int, err error) {
+	n, err = ps.writer.Write(p)
+	if err != nil {
+		return n, err
+	}
+	ps.done += int64(n)
+	if ps.callback != nil && ps.done >= ps.nextAt {
+		ps.callback(ps.done, ps.size)
+		ps.nextAt += ps.interval
+	}
+	return n, nil
+}
+
+func NewProgressStream(writer io.Writer, size int64, callback func(bytesRead, contentLength int64)) *ProgressStream {
+	var interval int64
+	interval = size / 100
+	if interval == 0 {
+		interval = 1
+	}
+	return &ProgressStream{
+		writer:   writer,
+		size:     size,
+		callback: callback,
+		nextAt:   interval,
+		interval: interval,
+	}
+}

go.mod

@@ -12,7 +12,6 @@ require (
 	github.com/rhysd/go-github-selfupdate v1.2.3
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/viper v1.19.0
-	github.com/studio-b12/gowebdav v0.10.0
 	golang.org/x/net v0.35.0
 	golang.org/x/time v0.10.0
 )

go.sum

@@ -172,8 +172,6 @@ github.com/spf13/viper v1.19.0 h1:RWq5SEjt8o25SROyN3z2OrDB9l7RPd3lwTWU8EcEdcI=
 github.com/spf13/viper v1.19.0/go.mod h1:GQUN9bilAbhU/jgc1bKs99f/suXKeUMct8Adx5+Ntkg=
 github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
 github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/studio-b12/gowebdav v0.10.0 h1:Yewz8FFiadcGEu4hxS/AAJQlHelndqln1bns3hcJIYc=
-github.com/studio-b12/gowebdav v0.10.0/go.mod h1:bHA7t77X/QFExdeAnDzK6vKM34kEZAcE1OX4MfiwjkE=
 github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
 github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
 github.com/tcnksm/go-gitconfig v0.1.2 h1:iiDhRitByXAEyjgBqsKi9QU4o2TNtv9kPP3RgPgXBPw=

storage/webdav/client.go

@@ -0,0 +1,70 @@
+package webdav
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+)
+
+type Client struct {
+	BaseURL    string
+	Username   string
+	Password   string
+	httpClient *http.Client
+}
+
+func NewClient(baseURL, username, password string, httpClient *http.Client) *Client {
+	if !strings.HasSuffix(baseURL, "/") {
+		baseURL += "/"
+	}
+	if httpClient == nil {
+		httpClient = http.DefaultClient
+	}
+	return &Client{
+		BaseURL:    baseURL,
+		Username:   username,
+		Password:   password,
+		httpClient: httpClient,
+	}
+}
+
+func (c *Client) doRequest(ctx context.Context, method, url string, body io.Reader) (*http.Response, error) {
+	req, err := http.NewRequestWithContext(ctx, method, url, body)
+	if err != nil {
+		return nil, err
+	}
+	if c.Username != "" && c.Password != "" {
+		req.SetBasicAuth(c.Username, c.Password)
+	}
+	return c.httpClient.Do(req)
+}
+
+func (c *Client) MkDir(ctx context.Context, dirPath string) error {
+	url := c.BaseURL + dirPath
+	resp, err := c.doRequest(ctx, "MKCOL", url, nil)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
+		return nil
+	}
+	return fmt.Errorf("MKCOL: %s", resp.Status)
+}
+
+func (c *Client) WriteFile(ctx context.Context, remotePath string, content io.Reader) error {
+	url := c.BaseURL + remotePath
+	resp, err := c.doRequest(ctx, "PUT", url, content)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
+		return nil
+	}
+	return fmt.Errorf("PUT: %s", resp.Status)
+}

storage/webdav/stream.go

@@ -1,50 +1,58 @@
 package webdav
 
-// TODO: gowebdav's WriteStream impl cause high memory usage, need to implement our own WriteStream
+import (
-// type WebdavWriter struct {
+	"context"
-// 	pipeWriter *io.PipeWriter
+	"fmt"
-// 	done       chan error
+	"io"
-// 	path       string
+	"path"
-// }
 
-// func (w *WebdavWriter) Write(p []byte) (n int, err error) {
+	"github.com/krau/SaveAny-Bot/logger"
-// 	return w.pipeWriter.Write(p)
+)
-// }
 
-// func (w *WebdavWriter) Close() error {
+type WebdavWriter struct {
-// 	if err := w.pipeWriter.Close(); err != nil {
+	pipeWriter *io.PipeWriter
-// 		return err
+	done       chan error
-// 	}
+	path       string
-// 	if err := <-w.done; err != nil {
+}
-// 		return fmt.Errorf("upload failed: %w", err)
-// 	}
 
-// 	return nil
+func (w *WebdavWriter) Write(p []byte) (n int, err error) {
-// }
+	return w.pipeWriter.Write(p)
+}
 
-// func (w *Webdav) NewUploadStream(ctx context.Context, storagePath string) (io.WriteCloser, error) {
+func (w *WebdavWriter) Close() error {
-// 	if err := w.client.MkdirAll(path.Dir(storagePath), os.ModePerm); err != nil {
+	if err := w.pipeWriter.Close(); err != nil {
-// 		logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
+		return err
-// 		return nil, ErrFailedToCreateDirectory
+	}
-// 	}
+	if err := <-w.done; err != nil {
-// 	pipeReader, pipeWriter := io.Pipe()
+		return fmt.Errorf("upload failed: %w", err)
-// 	done := make(chan error, 1)
+	}
-// 	go func() {
-// 		defer func() {
-// 			if err := recover(); err != nil {
-// 				done <- fmt.Errorf("panic during upload: %v", err)
-// 			}
-// 		}()
 
-// 		err := w.client.WriteStream(storagePath, pipeReader, os.ModePerm)
+	return nil
+}
 
-// 		pipeReader.Close()
+func (w *Webdav) NewUploadStream(ctx context.Context, storagePath string) (io.WriteCloser, error) {
-// 		done <- err
+	if err := w.client.MkDir(ctx, path.Dir(storagePath)); err != nil {
-// 	}()
+		logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
+		return nil, ErrFailedToCreateDirectory
+	}
+	pipeReader, pipeWriter := io.Pipe()
+	done := make(chan error, 1)
+	go func() {
+		defer func() {
+			if err := recover(); err != nil {
+				done <- fmt.Errorf("panic during upload: %v", err)
+			}
+		}()
 
-// 	return &WebdavWriter{
+		err := w.client.WriteFile(ctx, storagePath, pipeReader)
-// 		pipeWriter: pipeWriter,
+
-// 		done:       done,
+		pipeReader.Close()
-// 		path:       storagePath,
+		done <- err
-// 	}, nil
+	}()
-// }
+
+	return &WebdavWriter{
+		pipeWriter: pipeWriter,
+		done:       done,
+		path:       storagePath,
+	}, nil
+}

storage/webdav/webdav.go

@@ -3,6 +3,7 @@ package webdav
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"os"
 	"path"
 	"time"
@@ -10,12 +11,11 @@ import (
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/krau/SaveAny-Bot/logger"
 	"github.com/krau/SaveAny-Bot/types"
-	"github.com/studio-b12/gowebdav"
 )
 
 type Webdav struct {
 	config config.WebdavStorageConfig
-	client *gowebdav.Client
+	client *Client
 }
 
 func (w *Webdav) Init(cfg config.StorageConfig) error {
@@ -27,12 +27,9 @@ func (w *Webdav) Init(cfg config.StorageConfig) error {
 		return err
 	}
 	w.config = *webdavConfig
-	client := gowebdav.NewClient(webdavConfig.URL, webdavConfig.Username, webdavConfig.Password)
+	w.client = NewClient(w.config.URL, w.config.Username, w.config.Password, &http.Client{
-	if err := client.Connect(); err != nil {
+		Timeout: time.Hour * 12,
-		return fmt.Errorf("failed to connect to webdav server: %w", err)
+	})
-	}
-	client.SetTimeout(12 * time.Hour)
-	w.client = client
 	return nil
 }
 
@@ -46,7 +43,7 @@ func (w *Webdav) Name() string {
 
 func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error {
 	logger.L.Infof("Saving file %s to %s", filePath, storagePath)
-	if err := w.client.MkdirAll(path.Dir(storagePath), os.ModePerm); err != nil {
+	if err := w.client.MkDir(ctx, path.Dir(storagePath)); err != nil {
 		logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
 		return ErrFailedToCreateDirectory
 	}
@@ -57,7 +54,7 @@ func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error {
 	}
 	defer file.Close()
 
-	if err := w.client.WriteStream(storagePath, file, os.ModePerm); err != nil {
+	if err := w.client.WriteFile(ctx, storagePath, file); err != nil {
 		logger.L.Errorf("Failed to write file %s: %v", storagePath, err)
 		return ErrFailedToWriteFile
 	}