mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-08-30 04:36:41 +08:00
115 lines
3.1 KiB
Go
115 lines
3.1 KiB
Go
package core
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/charmbracelet/log"
|
|
"github.com/gotd/td/telegram/downloader"
|
|
|
|
"github.com/krau/SaveAny-Bot/database"
|
|
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
|
|
)
|
|
|
|
// TaskCodec serializes and rebuilds a task from its persisted payload.
|
|
// Task types without a registered codec are dropped with a warning on
|
|
// recovery instead of being silently re-enqueued.
|
|
type TaskCodec interface {
|
|
Marshal(task Executable) ([]byte, error)
|
|
Unmarshal(payload []byte) (Executable, error)
|
|
}
|
|
|
|
var (
|
|
taskCodecsMu sync.RWMutex
|
|
taskCodecs = make(map[tasktype.TaskType]TaskCodec)
|
|
|
|
dlerMu sync.RWMutex
|
|
dlerProvider func() downloader.Client
|
|
)
|
|
|
|
func RegisterTaskCodec(t tasktype.TaskType, codec TaskCodec) {
|
|
taskCodecsMu.Lock()
|
|
defer taskCodecsMu.Unlock()
|
|
taskCodecs[t] = codec
|
|
}
|
|
|
|
func TaskCodecFor(t tasktype.TaskType) (TaskCodec, bool) {
|
|
taskCodecsMu.RLock()
|
|
defer taskCodecsMu.RUnlock()
|
|
codec, ok := taskCodecs[t]
|
|
return codec, ok
|
|
}
|
|
|
|
// SetDownloaderProvider registers the download client factory used to
|
|
// rebuild tfile.TGFile values when recovering tasks.
|
|
func SetDownloaderProvider(f func() downloader.Client) {
|
|
dlerMu.Lock()
|
|
defer dlerMu.Unlock()
|
|
dlerProvider = f
|
|
}
|
|
|
|
// DownloaderClient returns the registered download client, or nil.
|
|
func DownloaderClient() downloader.Client {
|
|
dlerMu.RLock()
|
|
defer dlerMu.RUnlock()
|
|
if dlerProvider == nil {
|
|
return nil
|
|
}
|
|
return dlerProvider()
|
|
}
|
|
|
|
func persistTask(ctx context.Context, task Executable) error {
|
|
codec, ok := TaskCodecFor(task.Type())
|
|
if !ok {
|
|
return nil
|
|
}
|
|
payload, err := codec.Marshal(task)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return database.UpsertTask(ctx, &database.Task{
|
|
ID: task.TaskID(),
|
|
Type: string(task.Type()),
|
|
Payload: payload,
|
|
Status: string(database.TaskStatusQueued),
|
|
})
|
|
}
|
|
|
|
// RecoverTasks re-enqueues tasks that were unfinished when the process last
|
|
// exited. Must be called after storages are loaded and before Run. Tasks of
|
|
// types without a registered codec are dropped with a warning.
|
|
func RecoverTasks(ctx context.Context) {
|
|
logger := log.FromContext(ctx)
|
|
tasks, err := database.GetUnfinishedTasks(ctx)
|
|
if err != nil {
|
|
logger.Errorf("Failed to load unfinished tasks: %v", err)
|
|
return
|
|
}
|
|
for _, t := range tasks {
|
|
codec, ok := TaskCodecFor(tasktype.TaskType(t.Type))
|
|
if !ok {
|
|
logger.Warnf("Dropping unrecoverable task %s (type %s): no codec registered", t.ID, t.Type)
|
|
if err := database.DeleteTask(ctx, t.ID); err != nil {
|
|
logger.Errorf("Failed to delete task %s: %v", t.ID, err)
|
|
}
|
|
continue
|
|
}
|
|
task, err := codec.Unmarshal(t.Payload)
|
|
if err != nil {
|
|
logger.Errorf("Dropping task %s: failed to rebuild: %v", t.ID, err)
|
|
if err := database.DeleteTask(ctx, t.ID); err != nil {
|
|
logger.Errorf("Failed to delete task %s: %v", t.ID, err)
|
|
}
|
|
continue
|
|
}
|
|
if err := AddTask(ctx, task); err != nil {
|
|
logger.Errorf("Dropping task %s: failed to re-enqueue: %v", t.ID, err)
|
|
if err := database.DeleteTask(ctx, t.ID); err != nil {
|
|
logger.Errorf("Failed to delete task %s: %v", t.ID, err)
|
|
}
|
|
continue
|
|
}
|
|
logger.Infof("Recovered task %s (%s)", t.ID, t.Type)
|
|
}
|
|
}
|