mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-05-11 23:39:46 +08:00
refactor: complete core features
This commit is contained in:
krau
94
core/core.go
94
core/core.go
@@ -3,57 +3,73 @@ package core
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/bot"
|
||||
"github.com/krau/SaveAny-Bot/common"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/logger"
|
||||
"github.com/krau/SaveAny-Bot/queue"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
"github.com/krau/SaveAny-Bot/types"
|
||||
)
|
||||
|
||||
func processPendingTask(task types.Task) error {
|
||||
logger.L.Debugf("Start processing task: %s", task.FileName)
|
||||
time.Sleep(10 * time.Second)
|
||||
logger.L.Debugf("Task done: %s", task.FileName)
|
||||
func processPendingTask(task *types.Task) error {
|
||||
logger.L.Debugf("Start processing task: %s", task.String())
|
||||
|
||||
// os.MkdirAll(config.Cfg.Temp.BasePath, os.ModePerm)
|
||||
os.MkdirAll(config.Cfg.Temp.BasePath, os.ModePerm)
|
||||
|
||||
// message, err := bot.Client.GetMessageByID(task.ChatID, task.MessageID)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// logger.L.Debugf("Start downloading file: %s", task.FileName)
|
||||
// bot.Client.EditMessage(task.ChatID, task.ReplyMessageID, "正在下载文件...")
|
||||
// dest, err := message.Download(&telegram.DownloadOptions{
|
||||
// FileName: common.GetCacheFilePath(task.FileName),
|
||||
// Threads: config.Cfg.Threads,
|
||||
// ChunkSize: config.Cfg.ChunkSize,
|
||||
// // ProgressCallback: func(totalBytes, downloadedBytes int64) {},
|
||||
// })
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
logger.L.Debugf("Start downloading file: %s", task.String())
|
||||
|
||||
// defer func() {
|
||||
// if config.Cfg.Temp.CacheTTL > 0 {
|
||||
// common.RmFileAfter(dest, time.Duration(config.Cfg.Temp.CacheTTL)*time.Second)
|
||||
// } else {
|
||||
// if err := os.Remove(dest); err != nil {
|
||||
// logger.L.Errorf("Failed to purge file: %s", err)
|
||||
// }
|
||||
// }
|
||||
// }()
|
||||
// if task.StoragePath == "" {
|
||||
// task.StoragePath = task.FileName
|
||||
// }
|
||||
task.Ctx.(*ext.Context).EditMessage(task.ChatID, &tg.MessagesEditMessageRequest{
|
||||
Message: "开始下载文件...",
|
||||
ID: task.ReplyMessageID,
|
||||
})
|
||||
|
||||
// bot.Client.EditMessage(task.ChatID, task.ReplyMessageID, "下载完成, 正在转存文件...")
|
||||
// if err := storage.Save(task.Storage, task.Ctx, dest, task.StoragePath); err != nil {
|
||||
// return err
|
||||
// }
|
||||
readCloser, err := NewTelegramReader(task.Ctx, bot.Client, task.File.Location, 0, task.File.FileSize-1, task.File.FileSize)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to create reader: %w", err)
|
||||
}
|
||||
defer readCloser.Close()
|
||||
|
||||
dest, err := os.Create(common.GetCacheFilePath(task.FileName()))
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to create file: %w", err)
|
||||
}
|
||||
logger.L.Debug("Created file: ", dest.Name())
|
||||
defer dest.Close()
|
||||
|
||||
if _, err := io.CopyN(dest, readCloser, task.File.FileSize); err != nil {
|
||||
return fmt.Errorf("Failed to download file: %w", err)
|
||||
}
|
||||
|
||||
defer func() {
|
||||
if config.Cfg.Temp.CacheTTL > 0 {
|
||||
common.RmFileAfter(dest.Name(), time.Duration(config.Cfg.Temp.CacheTTL)*time.Second)
|
||||
} else {
|
||||
if err := os.Remove(dest.Name()); err != nil {
|
||||
logger.L.Errorf("Failed to purge file: %s", err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
if task.StoragePath == "" {
|
||||
task.StoragePath = task.File.FileName
|
||||
}
|
||||
|
||||
task.Ctx.(*ext.Context).EditMessage(task.ChatID, &tg.MessagesEditMessageRequest{
|
||||
Message: "下载完成, 正在转存文件...",
|
||||
ID: task.ReplyMessageID,
|
||||
})
|
||||
|
||||
if err := storage.Save(task.Storage, task.Ctx, dest.Name(), task.StoragePath); err != nil {
|
||||
return fmt.Errorf("Failed to save file: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -61,12 +77,12 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
|
||||
for {
|
||||
semaphore <- struct{}{}
|
||||
task := queue.GetTask()
|
||||
logger.L.Debugf("Got task: %s", task.FileName)
|
||||
logger.L.Debugf("Got task: %s", task.String())
|
||||
|
||||
switch task.Status {
|
||||
case types.Pending:
|
||||
logger.L.Infof("Processing task: %s", task.String())
|
||||
if err := processPendingTask(task); err != nil {
|
||||
if err := processPendingTask(&task); err != nil {
|
||||
logger.L.Errorf("Failed to do task: %s", err)
|
||||
task.Error = err
|
||||
if errors.Is(err, context.Canceled) {
|
||||
@@ -97,7 +113,7 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
|
||||
logger.L.Errorf("Unknown task status: %s", task.Status)
|
||||
}
|
||||
<-semaphore
|
||||
logger.L.Debugf("Task done: %s", task.FileName)
|
||||
logger.L.Debugf("Task done: %s", task.String())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
135
core/reader.go
Normal file
135
core/reader.go
Normal file
@@ -0,0 +1,135 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
|
||||
"github.com/celestix/gotgproto"
|
||||
"github.com/gotd/td/tg"
|
||||
)
|
||||
|
||||
type telegramReader struct {
|
||||
ctx context.Context
|
||||
client *gotgproto.Client
|
||||
location *tg.InputDocumentFileLocation
|
||||
start int64
|
||||
end int64
|
||||
next func() ([]byte, error)
|
||||
buffer []byte
|
||||
bytesread int64
|
||||
chunkSize int64
|
||||
i int64
|
||||
contentLength int64
|
||||
}
|
||||
|
||||
func (*telegramReader) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *telegramReader) Read(p []byte) (n int, err error) {
|
||||
if r.bytesread == r.contentLength {
|
||||
return 0, io.EOF
|
||||
}
|
||||
|
||||
if r.i >= int64(len(r.buffer)) {
|
||||
r.buffer, err = r.next()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(r.buffer) == 0 {
|
||||
r.next = r.partStream()
|
||||
r.buffer, err = r.next()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
}
|
||||
r.i = 0
|
||||
}
|
||||
n = copy(p, r.buffer[r.i:])
|
||||
r.i += int64(n)
|
||||
r.bytesread += int64(n)
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func NewTelegramReader(
|
||||
ctx context.Context,
|
||||
client *gotgproto.Client,
|
||||
location *tg.InputDocumentFileLocation,
|
||||
start int64,
|
||||
end int64,
|
||||
contentLength int64,
|
||||
) (io.ReadCloser, error) {
|
||||
|
||||
r := &telegramReader{
|
||||
ctx: ctx,
|
||||
location: location,
|
||||
client: client,
|
||||
start: start,
|
||||
end: end,
|
||||
chunkSize: int64(1024 * 1024),
|
||||
contentLength: contentLength,
|
||||
}
|
||||
|
||||
r.next = r.partStream()
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func (r *telegramReader) chunk(offset int64, limit int64) ([]byte, error) {
|
||||
|
||||
req := &tg.UploadGetFileRequest{
|
||||
Offset: offset,
|
||||
Limit: int(limit),
|
||||
Location: r.location,
|
||||
}
|
||||
|
||||
res, err := r.client.API().UploadGetFile(r.ctx, req)
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
switch result := res.(type) {
|
||||
case *tg.UploadFile:
|
||||
return result.Bytes, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unexpected type %T", r)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *telegramReader) partStream() func() ([]byte, error) {
|
||||
|
||||
start := r.start
|
||||
end := r.end
|
||||
offset := start - (start % r.chunkSize)
|
||||
|
||||
firstPartCut := start - offset
|
||||
lastPartCut := (end % r.chunkSize) + 1
|
||||
partCount := int((end - offset + r.chunkSize) / r.chunkSize)
|
||||
currentPart := 1
|
||||
|
||||
readData := func() ([]byte, error) {
|
||||
if currentPart > partCount {
|
||||
return make([]byte, 0), nil
|
||||
}
|
||||
res, err := r.chunk(offset, r.chunkSize)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(res) == 0 {
|
||||
return res, nil
|
||||
} else if partCount == 1 {
|
||||
res = res[firstPartCut:lastPartCut]
|
||||
} else if currentPart == 1 {
|
||||
res = res[firstPartCut:]
|
||||
} else if currentPart == partCount {
|
||||
res = res[:lastPartCut]
|
||||
}
|
||||
|
||||
currentPart++
|
||||
offset += r.chunkSize
|
||||
return res, nil
|
||||
}
|
||||
return readData
|
||||
}
|
||||
Reference in New Issue
Block a user