归档
This commit is contained in:
1 parent
594436fddd
commit
9da521ee4e
31 files changed
+5768
-267
No files matched your search
@@ -56,10 +56,24 @@ type TaskDTO struct {
|
||||
LLMFailoverState string `json:"llm_failover_state"`
|
||||
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
|
||||
SourceTaskIDs []string `json:"source_task_ids"`
|
||||
ArchiveBlockedBy string `json:"archive_blocked_by_task_id,omitempty"`
|
||||
CompanyIDs []int64 `json:"company_ids"`
|
||||
CoverageEnabled bool `json:"coverage_enabled"` // 资产覆盖度功能开关(创建时定)
|
||||
}
|
||||
|
||||
func applyTaskArchiveBlocker(dto *TaskDTO, blockers map[int64]int64) {
|
||||
if dto == nil || len(blockers) == 0 {
|
||||
return
|
||||
}
|
||||
taskID, err := strconv.ParseInt(dto.ID, 10, 64)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if dependentID := blockers[taskID]; dependentID > 0 {
|
||||
dto.ArchiveBlockedBy = strconv.FormatInt(dependentID, 10)
|
||||
}
|
||||
}
|
||||
|
||||
// TokenTotalDTO is a whole-task (all agents) token aggregate.
|
||||
type TokenTotalDTO struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
|
||||
@@ -37,6 +37,7 @@ func (s *Server) findingGroups(w http.ResponseWriter, r *http.Request) {
|
||||
Status: normFilter(q.Get("status")),
|
||||
VulnClass: normFilter(q.Get("vulnclass")),
|
||||
TaskID: normFilter(q.Get("task_id")),
|
||||
Query: q.Get("q"),
|
||||
Sort: q.Get("sort"),
|
||||
}
|
||||
groups, total, findingTotal, err := s.m.pg.ListFindingGroups(filter, page, limit)
|
||||
|
||||
@@ -102,6 +102,28 @@ func TestFindingGroupsReturnsTaskBucketsAndNormalizesPagination(t *testing.T) {
|
||||
t.Fatalf("unexpected first group: %+v", page.Items[0])
|
||||
}
|
||||
|
||||
request = httptest.NewRequest(http.MethodGet,
|
||||
"/api/exploration/findings/groups?vulnclass="+marker+"&q=finding+2&page=1&limit=10", nil)
|
||||
recorder = httptest.NewRecorder()
|
||||
s.findingGroups(recorder, request)
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("searched groups status=%d body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
page = struct {
|
||||
Items []db.FindingGroup `json:"items"`
|
||||
Total int `json:"total"`
|
||||
FindingTotal int `json:"finding_total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"page_size"`
|
||||
}{}
|
||||
if err := json.Unmarshal(recorder.Body.Bytes(), &page); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(page.Items) != 1 || page.Total != 1 || page.FindingTotal != 1 || page.Items[0].TaskID == nil ||
|
||||
*page.Items[0].TaskID != mustTaskID(t, taskB.ID) {
|
||||
t.Fatalf("unexpected searched group page: %+v", page)
|
||||
}
|
||||
|
||||
request = httptest.NewRequest(http.MethodGet,
|
||||
"/api/exploration/findings/groups?vulnclass="+marker+"&page=-4&limit=0", nil)
|
||||
recorder = httptest.NewRecorder()
|
||||
|
||||
@@ -358,6 +358,25 @@ func NewManager(dir, proxyAddr string) (*Manager, error) {
|
||||
if err != nil {
|
||||
log.Printf("[traffic] disabled: %v", err)
|
||||
} else {
|
||||
err = tr.RecoverHostDeleteStages(func(_ int64, taskID int64) (bool, error) {
|
||||
if taskID <= 0 {
|
||||
return false, errors.New("归档流量暂存日志缺少任务 ID")
|
||||
}
|
||||
task, taskErr := pg.GetTask(taskID)
|
||||
if taskErr != nil {
|
||||
return false, taskErr
|
||||
}
|
||||
// Archived/permanently deleted tasks are hidden from GetTask. A
|
||||
// restored task is visible and needs the staged traffic put back.
|
||||
return task == nil, nil
|
||||
})
|
||||
if err != nil {
|
||||
_ = tr.Close()
|
||||
_ = pg.Close()
|
||||
return nil, fmt.Errorf("recover traffic delete staging: %w", err)
|
||||
}
|
||||
}
|
||||
if tr != nil {
|
||||
m.traffic = tr
|
||||
go func() {
|
||||
log.Printf("[traffic] recording proxy on %s (set HTTP_PROXY=%s + trust _ca CA)", proxyAddr, tr.ProxyAddr())
|
||||
|
||||
+24
-3
@@ -116,6 +116,11 @@ type Server struct {
|
||||
// the persisted current profile at every new LLM call.
|
||||
taskAgentMu sync.Mutex
|
||||
taskAgents map[string]*taskAgentBundle
|
||||
|
||||
// Cold task archives run through one persistent FIFO worker. The buffered wake
|
||||
// channel coalesces enqueue bursts; the database remains the source of truth.
|
||||
archiveWake chan struct{}
|
||||
archiveWG sync.WaitGroup
|
||||
}
|
||||
|
||||
// profBundle is a planner/worker pair built from one LLM profile.
|
||||
@@ -151,7 +156,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
triggerActive: map[string]int{}, triggerCfg: map[string]triggerBehavior{},
|
||||
profAgents: map[int64]*profBundle{}, profChatAgents: map[int64]*agent.ChatAgent{},
|
||||
provByProfile: map[int64]*provEntry{}, llmHealth: newLLMHealthRegistry(m.pg),
|
||||
taskAgents: map[string]*taskAgentBundle{}}
|
||||
taskAgents: map[string]*taskAgentBundle{}, archiveWake: make(chan struct{}, 1)}
|
||||
// Every task uses a stable task router. An empty explicit chain is resolved by
|
||||
// that router through Agent bindings and then the global provider, so adding a
|
||||
// first chain to a running task takes effect on its very next LLM call.
|
||||
@@ -264,6 +269,7 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
|
||||
}
|
||||
}
|
||||
go s.reconcileConcurrency()
|
||||
s.startTaskArchiveWorker()
|
||||
s.wireInterceptReviewer() // LLM 兜底审批:未命中拦截规则的命令交给模型判定
|
||||
return s
|
||||
}
|
||||
@@ -690,6 +696,14 @@ func (s *Server) Handler() http.Handler {
|
||||
mux.HandleFunc("PATCH /api/tasks/{id}", s.updateTaskMetadata)
|
||||
mux.HandleFunc("PATCH /api/tasks/{id}/category", s.updateTaskCategory)
|
||||
mux.HandleFunc("POST /api/tasks/control/batch", s.controlTasksBatch)
|
||||
mux.HandleFunc("GET /api/task-archives", s.listTaskArchives)
|
||||
mux.HandleFunc("GET /api/task-archives/{id}", s.getTaskArchive)
|
||||
mux.HandleFunc("POST /api/tasks/{id}/archive", s.queueTaskArchive)
|
||||
mux.HandleFunc("POST /api/tasks/archive/batch", s.queueTaskArchivesBatch)
|
||||
mux.HandleFunc("POST /api/task-archives/{id}/restore", s.queueTaskArchiveRestore)
|
||||
mux.HandleFunc("POST /api/task-archives/restore/batch", s.restoreTaskArchivesBatch)
|
||||
mux.HandleFunc("DELETE /api/task-archives/{id}", s.queueTaskArchiveDelete)
|
||||
mux.HandleFunc("POST /api/task-archives/delete/batch", s.deleteTaskArchivesBatch)
|
||||
mux.HandleFunc("GET /api/tasks/{id}/coverage", s.taskCoverage)
|
||||
mux.HandleFunc("GET /api/tasks/{id}/coverage-graph", s.taskCoverageGraph)
|
||||
mux.HandleFunc("GET /api/tasks/{id}/asset-refs", s.taskAssetRefs)
|
||||
@@ -996,9 +1010,11 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
list := s.m.List()
|
||||
metrics, _ := s.m.PG().TaskListMetricsAll()
|
||||
archiveBlockers, _ := s.m.PG().TaskArchiveBlockers()
|
||||
dtos := make([]TaskDTO, 0, len(list))
|
||||
for _, t := range list {
|
||||
dto := taskDTO(t, s.resolvedTaskStatus(t))
|
||||
applyTaskArchiveBlocker(&dto, archiveBlockers)
|
||||
metric := metrics[t.ExpID]
|
||||
dto.Tokens = tokenTotalDTO(metric.Tokens)
|
||||
// prefer the live in-memory heartbeat (fresher) and fall back to the
|
||||
@@ -1640,7 +1656,10 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 404, "task not found")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, taskDTO(t, s.resolvedTaskStatus(t)))
|
||||
dto := taskDTO(t, s.resolvedTaskStatus(t))
|
||||
archiveBlockers, _ := s.m.PG().TaskArchiveBlockers()
|
||||
applyTaskArchiveBlocker(&dto, archiveBlockers)
|
||||
writeJSON(w, 200, dto)
|
||||
}
|
||||
|
||||
// taskCoverage returns a task's rough asset test coverage (denominator/tested/backlog).
|
||||
@@ -1842,6 +1861,7 @@ func (s *Server) findings(w http.ResponseWriter, r *http.Request) {
|
||||
VulnClass: normFilter(q.Get("vulnclass")),
|
||||
// task_id(独立于会切到「按任务节点」分支的 task 参数):全局表按任务筛选。
|
||||
TaskID: normFilter(q.Get("task_id")),
|
||||
Query: q.Get("q"),
|
||||
Sort: q.Get("sort"),
|
||||
}
|
||||
fs, total, err := s.m.pg.ListFindingsPage(filter, page, limit)
|
||||
@@ -2011,7 +2031,7 @@ func (s *Server) getFinding(w http.ResponseWriter, r *http.Request) {
|
||||
// format = md-single(整合一份 .md)| md-zip(一漏洞一 .md,打包 zip)
|
||||
// | csv | json
|
||||
// ids = 逗号分隔的 finding id(scope=selected 时必填)
|
||||
// 筛选参数 severity/status/vulnclass/task_id/sort 与列表接口一致(scope=filtered 用)。
|
||||
// 筛选参数 severity/status/vulnclass/task_id/q/sort 与列表接口一致(scope=filtered 用)。
|
||||
func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
q := r.URL.Query()
|
||||
scope := q.Get("scope")
|
||||
@@ -2045,6 +2065,7 @@ func (s *Server) findingsExport(w http.ResponseWriter, r *http.Request) {
|
||||
Status: normFilter(q.Get("status")),
|
||||
VulnClass: normFilter(q.Get("vulnclass")),
|
||||
TaskID: normFilter(q.Get("task_id")),
|
||||
Query: q.Get("q"),
|
||||
Sort: q.Get("sort"),
|
||||
}
|
||||
default:
|
||||
|
||||
@@ -0,0 +1,655 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
pgdb "github.com/Autumn-27/artex/db"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
)
|
||||
|
||||
const (
|
||||
archiveDirMode = 0o700
|
||||
archiveFileMode = 0o600
|
||||
maxArchiveFiles = 1_000_000
|
||||
maxArchiveBytes = int64(1 << 47) // 128 TiB safety ceiling for corrupt headers.
|
||||
)
|
||||
|
||||
type archiveFileMove struct {
|
||||
Source string `json:"source"`
|
||||
Relative string `json:"relative"`
|
||||
}
|
||||
|
||||
type archiveStageJournal struct {
|
||||
ArchiveID int64 `json:"archive_id"`
|
||||
TaskID string `json:"task_id"`
|
||||
Moves []archiveFileMove `json:"moves"`
|
||||
}
|
||||
|
||||
type archiveRestoreJournal struct {
|
||||
ArchiveID int64 `json:"archive_id"`
|
||||
TaskID int64 `json:"task_id"`
|
||||
Moves []restoredArchivePath `json:"moves"`
|
||||
}
|
||||
|
||||
type taskArchiveFileStage struct {
|
||||
root string
|
||||
payload string
|
||||
journal archiveStageJournal
|
||||
done bool
|
||||
}
|
||||
|
||||
type restoredArchivePath struct {
|
||||
Source string
|
||||
Destination string
|
||||
}
|
||||
|
||||
type taskArchiveRestoreFiles struct {
|
||||
extracted string
|
||||
moves []restoredArchivePath
|
||||
done bool
|
||||
}
|
||||
|
||||
func writeArchiveJournal(path string, value any) error {
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, archiveFileMode)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := file.Write(raw); err != nil {
|
||||
_ = file.Close()
|
||||
return err
|
||||
}
|
||||
if err := file.Sync(); err != nil {
|
||||
_ = file.Close()
|
||||
return err
|
||||
}
|
||||
return file.Close()
|
||||
}
|
||||
|
||||
func taskArchiveRoot(dataDir string) string {
|
||||
return filepath.Join(dataDir, "archives", "tasks")
|
||||
}
|
||||
|
||||
func taskArchivePath(dataDir string, archiveID int64, taskID string) string {
|
||||
return filepath.Join(taskArchiveRoot(dataDir), fmt.Sprintf("task-%s-%d.tar.zst", taskID, archiveID))
|
||||
}
|
||||
|
||||
func stageTaskArchiveFiles(dataDir string, archiveID int64, taskID string, explorationID int64) (*taskArchiveFileStage, error) {
|
||||
root := filepath.Join(taskArchiveRoot(dataDir), ".staging", strconv.FormatInt(archiveID, 10))
|
||||
stage := &taskArchiveFileStage{root: root, payload: filepath.Join(root, "payload")}
|
||||
journalPath := filepath.Join(root, "journal.json")
|
||||
if raw, err := os.ReadFile(journalPath); err == nil {
|
||||
if err := json.Unmarshal(raw, &stage.journal); err != nil {
|
||||
return nil, fmt.Errorf("read task archive staging journal: %w", err)
|
||||
}
|
||||
return stage, nil
|
||||
} else if !os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Join(stage.payload, "files"), archiveDirMode); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stage.journal = archiveStageJournal{ArchiveID: archiveID, TaskID: taskID}
|
||||
targets := []archiveFileMove{}
|
||||
taskDir := filepath.Join(dataDir, "tasks", taskID)
|
||||
if _, err := os.Lstat(taskDir); err == nil {
|
||||
targets = append(targets, archiveFileMove{Source: taskDir, Relative: filepath.Join("files", "tasks", taskID)})
|
||||
} else if !os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
transcriptDir := filepath.Join(dataDir, "transcripts")
|
||||
entries, err := os.ReadDir(transcriptDir)
|
||||
if err != nil && !os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
prefix := fmt.Sprintf("exp%d-", explorationID)
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if !strings.HasPrefix(name, prefix) || (!entry.IsDir() && !strings.HasSuffix(name, ".jsonl")) {
|
||||
continue
|
||||
}
|
||||
targets = append(targets, archiveFileMove{
|
||||
Source: filepath.Join(transcriptDir, name), Relative: filepath.Join("files", "transcripts", name),
|
||||
})
|
||||
}
|
||||
// Persist the complete plan before the first rename. Recovery can therefore
|
||||
// roll back any prefix of the moves after an abrupt process termination.
|
||||
stage.journal.Moves = targets
|
||||
if err := writeArchiveJournal(journalPath, stage.journal); err != nil {
|
||||
_ = stage.rollback()
|
||||
return nil, err
|
||||
}
|
||||
for _, move := range targets {
|
||||
destination := filepath.Join(stage.payload, move.Relative)
|
||||
if err := os.MkdirAll(filepath.Dir(destination), archiveDirMode); err != nil {
|
||||
_ = stage.rollback()
|
||||
return nil, err
|
||||
}
|
||||
if err := os.Rename(move.Source, destination); err != nil {
|
||||
_ = stage.rollback()
|
||||
return nil, fmt.Errorf("stage task archive path %s: %w", move.Source, err)
|
||||
}
|
||||
}
|
||||
return stage, nil
|
||||
}
|
||||
|
||||
func (s *taskArchiveFileStage) rollback() error {
|
||||
if s == nil || s.done {
|
||||
return nil
|
||||
}
|
||||
var errs []error
|
||||
for i := len(s.journal.Moves) - 1; i >= 0; i-- {
|
||||
move := s.journal.Moves[i]
|
||||
staged := filepath.Join(s.payload, move.Relative)
|
||||
if _, err := os.Lstat(staged); os.IsNotExist(err) {
|
||||
continue
|
||||
} else if err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if _, err := os.Lstat(move.Source); err == nil {
|
||||
errs = append(errs, fmt.Errorf("archive rollback destination exists: %s", move.Source))
|
||||
continue
|
||||
} else if !os.IsNotExist(err) {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(move.Source), 0o755); err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if err := os.Rename(staged, move.Source); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}
|
||||
if len(errs) == 0 {
|
||||
if err := os.RemoveAll(s.root); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}
|
||||
s.done = true
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
func (s *taskArchiveFileStage) commit() error {
|
||||
if s == nil || s.done {
|
||||
return nil
|
||||
}
|
||||
s.done = true
|
||||
return os.RemoveAll(s.root)
|
||||
}
|
||||
|
||||
func installTaskArchiveFiles(dataDir, extractedDir, taskID string, archiveID int64) (*taskArchiveRestoreFiles, error) {
|
||||
stage := &taskArchiveRestoreFiles{extracted: extractedDir}
|
||||
numericTaskID, err := strconv.ParseInt(taskID, 10, 64)
|
||||
if err != nil || numericTaskID <= 0 {
|
||||
return nil, fmt.Errorf("invalid restore task id %q", taskID)
|
||||
}
|
||||
sources := []restoredArchivePath{}
|
||||
workspace := filepath.Join(extractedDir, "files", "tasks", taskID)
|
||||
if _, err := os.Lstat(workspace); err == nil {
|
||||
sources = append(sources, restoredArchivePath{Source: workspace, Destination: filepath.Join(dataDir, "tasks", taskID)})
|
||||
} else if !os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
transcripts := filepath.Join(extractedDir, "files", "transcripts")
|
||||
entries, err := os.ReadDir(transcripts)
|
||||
if err != nil && !os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
for _, entry := range entries {
|
||||
sources = append(sources, restoredArchivePath{
|
||||
Source: filepath.Join(transcripts, entry.Name()), Destination: filepath.Join(dataDir, "transcripts", entry.Name()),
|
||||
})
|
||||
}
|
||||
for _, move := range sources {
|
||||
if _, err := os.Lstat(move.Destination); err == nil {
|
||||
return nil, fmt.Errorf("restore destination already exists: %s", move.Destination)
|
||||
} else if !os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
stage.moves = sources
|
||||
journal := archiveRestoreJournal{ArchiveID: archiveID, TaskID: numericTaskID, Moves: sources}
|
||||
if err := writeArchiveJournal(filepath.Join(extractedDir, "restore-journal.json"), journal); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, move := range sources {
|
||||
if err := os.MkdirAll(filepath.Dir(move.Destination), 0o755); err != nil {
|
||||
_ = stage.rollback()
|
||||
return nil, err
|
||||
}
|
||||
if err := os.Rename(move.Source, move.Destination); err != nil {
|
||||
_ = stage.rollback()
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return stage, nil
|
||||
}
|
||||
|
||||
func (s *taskArchiveRestoreFiles) rollback() error {
|
||||
if s == nil || s.done {
|
||||
return nil
|
||||
}
|
||||
var errs []error
|
||||
for i := len(s.moves) - 1; i >= 0; i-- {
|
||||
move := s.moves[i]
|
||||
if _, err := os.Lstat(move.Destination); os.IsNotExist(err) {
|
||||
continue
|
||||
} else if err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if _, err := os.Lstat(move.Source); err == nil {
|
||||
errs = append(errs, fmt.Errorf("restore rollback source exists: %s", move.Source))
|
||||
continue
|
||||
} else if !os.IsNotExist(err) {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(move.Source), archiveDirMode); err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if err := os.Rename(move.Destination, move.Source); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}
|
||||
s.done = true
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
func (s *taskArchiveRestoreFiles) commit() error {
|
||||
if s == nil || s.done {
|
||||
return nil
|
||||
}
|
||||
s.done = true
|
||||
return os.RemoveAll(s.extracted)
|
||||
}
|
||||
|
||||
// recoverTaskArchiveRestoreStages resolves file installs left by an interrupted
|
||||
// restore. If PostgreSQL committed, installed files are authoritative and only
|
||||
// the extraction directory is stale. Otherwise all completed renames are moved
|
||||
// back so the persistent restore job can retry from a clean destination.
|
||||
func recoverTaskArchiveRestoreStages(dataDir string, pg *pgdb.DB) error {
|
||||
parent := filepath.Join(taskArchiveRoot(dataDir), ".restore")
|
||||
entries, err := os.ReadDir(parent)
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var errs []error
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
root := filepath.Join(parent, entry.Name())
|
||||
raw, err := os.ReadFile(filepath.Join(root, "restore-journal.json"))
|
||||
if os.IsNotExist(err) {
|
||||
// Extraction was interrupted before any destination rename.
|
||||
if err := os.RemoveAll(root); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
var journal archiveRestoreJournal
|
||||
if err := json.Unmarshal(raw, &journal); err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
restored, err := pg.IsTaskArchiveRestored(journal.ArchiveID)
|
||||
if errors.Is(err, pgdb.ErrTaskArchiveNotFound) && journal.TaskID > 0 {
|
||||
task, taskErr := pg.GetTask(journal.TaskID)
|
||||
if taskErr != nil {
|
||||
errs = append(errs, taskErr)
|
||||
continue
|
||||
}
|
||||
restored = task != nil
|
||||
err = nil
|
||||
}
|
||||
if err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if restored {
|
||||
if err := os.RemoveAll(root); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
stage := &taskArchiveRestoreFiles{extracted: root, moves: journal.Moves}
|
||||
if err := stage.rollback(); err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if err := os.RemoveAll(root); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
// recoverTaskArchiveStages resolves file renames left by an interrupted archive.
|
||||
// A committed cold task already has a verified package, so stale staging can be
|
||||
// discarded; otherwise files are moved back before the persistent job retries.
|
||||
func recoverTaskArchiveStages(dataDir string, pg *pgdb.DB) error {
|
||||
parent := filepath.Join(taskArchiveRoot(dataDir), ".staging")
|
||||
entries, err := os.ReadDir(parent)
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var errs []error
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
root := filepath.Join(parent, entry.Name())
|
||||
raw, err := os.ReadFile(filepath.Join(root, "journal.json"))
|
||||
if err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
var journal archiveStageJournal
|
||||
if err := json.Unmarshal(raw, &journal); err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
archive, getErr := pg.GetTaskArchive(journal.ArchiveID)
|
||||
if getErr != nil {
|
||||
errs = append(errs, getErr)
|
||||
continue
|
||||
}
|
||||
if archive != nil && archive.State == pgdb.ArchiveReady {
|
||||
if err := os.RemoveAll(root); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
stage := &taskArchiveFileStage{root: root, payload: filepath.Join(root, "payload"), journal: journal}
|
||||
if err := stage.rollback(); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
if archive != nil && archive.ArchivePath != "" {
|
||||
_ = os.Remove(archive.ArchivePath)
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
func recoverTaskArchiveDeletePackages(dataDir string, pg *pgdb.DB) error {
|
||||
root := taskArchiveRoot(dataDir)
|
||||
entries, err := os.ReadDir(root)
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var errs []error
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := entry.Name()
|
||||
marker := strings.LastIndex(name, ".deleting-")
|
||||
if marker < 0 {
|
||||
continue
|
||||
}
|
||||
archiveID, err := strconv.ParseInt(name[marker+len(".deleting-"):], 10, 64)
|
||||
if err != nil || archiveID <= 0 {
|
||||
continue
|
||||
}
|
||||
staged := filepath.Join(root, name)
|
||||
original := filepath.Join(root, name[:marker])
|
||||
archive, err := pg.GetTaskArchive(archiveID)
|
||||
if err != nil {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if archive == nil {
|
||||
if err := os.Remove(staged); err != nil && !os.IsNotExist(err) {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if archive.State == pgdb.DeleteQueued || archive.State == pgdb.Deleting || archive.State == pgdb.DeleteFailed {
|
||||
// The idempotent delete worker consumes the staged path directly.
|
||||
continue
|
||||
}
|
||||
if _, err := os.Lstat(original); err == nil {
|
||||
errs = append(errs, fmt.Errorf("archive delete recovery destination exists: %s", original))
|
||||
continue
|
||||
} else if !os.IsNotExist(err) {
|
||||
errs = append(errs, err)
|
||||
continue
|
||||
}
|
||||
if err := os.Rename(staged, original); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
func stageTaskArchivePackageDelete(archivePath string, archiveID int64) (string, bool, error) {
|
||||
staged := archivePath + fmt.Sprintf(".deleting-%d", archiveID)
|
||||
if _, err := os.Lstat(staged); err == nil {
|
||||
if _, originalErr := os.Lstat(archivePath); originalErr == nil {
|
||||
return staged, false, errors.New("归档包原文件和删除暂存文件同时存在")
|
||||
} else if !os.IsNotExist(originalErr) {
|
||||
return staged, false, originalErr
|
||||
}
|
||||
return staged, true, nil
|
||||
} else if !os.IsNotExist(err) {
|
||||
return staged, false, err
|
||||
}
|
||||
if err := os.Rename(archivePath, staged); err == nil {
|
||||
return staged, true, nil
|
||||
} else if !os.IsNotExist(err) {
|
||||
return staged, false, err
|
||||
}
|
||||
return staged, false, nil
|
||||
}
|
||||
|
||||
func writeTaskArchivePackage(path, payloadDir string, snapshot *pgdb.TaskArchiveSnapshot) (originalSize, compressedSize int64, checksum string, err error) {
|
||||
if snapshot == nil {
|
||||
return 0, 0, "", errors.New("nil task archive snapshot")
|
||||
}
|
||||
if err = os.MkdirAll(filepath.Dir(path), archiveDirMode); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
manifest, err := json.Marshal(snapshot)
|
||||
if err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
if err = os.WriteFile(filepath.Join(payloadDir, "manifest.json"), manifest, archiveFileMode); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
temporary := path + ".partial"
|
||||
_ = os.Remove(temporary)
|
||||
file, err := os.OpenFile(temporary, os.O_CREATE|os.O_EXCL|os.O_WRONLY, archiveFileMode)
|
||||
if err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
committed := false
|
||||
defer func() {
|
||||
_ = file.Close()
|
||||
if !committed {
|
||||
_ = os.Remove(temporary)
|
||||
}
|
||||
}()
|
||||
hasher := sha256.New()
|
||||
multi := io.MultiWriter(file, hasher)
|
||||
encoder, err := zstd.NewWriter(multi,
|
||||
zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(10)),
|
||||
zstd.WithEncoderCRC(true),
|
||||
)
|
||||
if err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
tarWriter := tar.NewWriter(encoder)
|
||||
walkErr := filepath.WalkDir(payloadDir, func(current string, entry fs.DirEntry, walkErr error) error {
|
||||
if walkErr != nil {
|
||||
return walkErr
|
||||
}
|
||||
relative, err := filepath.Rel(payloadDir, current)
|
||||
if err != nil || relative == "." {
|
||||
return err
|
||||
}
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if info.Mode()&os.ModeSymlink != 0 {
|
||||
return fmt.Errorf("task archive refuses symlink: %s", current)
|
||||
}
|
||||
header, err := tar.FileInfoHeader(info, "")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
header.Name = filepath.ToSlash(relative)
|
||||
if err := tarWriter.WriteHeader(header); err != nil {
|
||||
return err
|
||||
}
|
||||
if entry.IsDir() {
|
||||
return nil
|
||||
}
|
||||
originalSize += info.Size()
|
||||
input, err := os.Open(current)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, copyErr := io.Copy(tarWriter, input)
|
||||
closeErr := input.Close()
|
||||
return errors.Join(copyErr, closeErr)
|
||||
})
|
||||
if walkErr != nil {
|
||||
_ = tarWriter.Close()
|
||||
_ = encoder.Close()
|
||||
return 0, 0, "", walkErr
|
||||
}
|
||||
if err := tarWriter.Close(); err != nil {
|
||||
_ = encoder.Close()
|
||||
return 0, 0, "", err
|
||||
}
|
||||
if err := encoder.Close(); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
if err := file.Sync(); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
if err := file.Close(); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
stat, err := os.Stat(temporary)
|
||||
if err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
compressedSize = stat.Size()
|
||||
checksum = hex.EncodeToString(hasher.Sum(nil))
|
||||
if err := os.Rename(temporary, path); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
if err := os.Chmod(path, archiveFileMode); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
committed = true
|
||||
return originalSize, compressedSize, checksum, nil
|
||||
}
|
||||
|
||||
func extractTaskArchivePackage(path, expectedSHA, destination string) error {
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer file.Close()
|
||||
if expectedSHA != "" {
|
||||
hasher := sha256.New()
|
||||
if _, err := io.Copy(hasher, file); err != nil {
|
||||
return err
|
||||
}
|
||||
if hex.EncodeToString(hasher.Sum(nil)) != strings.ToLower(expectedSHA) {
|
||||
return errors.New("task archive checksum mismatch")
|
||||
}
|
||||
if _, err := file.Seek(0, io.SeekStart); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
decoder, err := zstd.NewReader(file)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer decoder.Close()
|
||||
tarReader := tar.NewReader(decoder)
|
||||
if err := os.MkdirAll(destination, archiveDirMode); err != nil {
|
||||
return err
|
||||
}
|
||||
var entries int
|
||||
var total int64
|
||||
for {
|
||||
header, err := tarReader.Next()
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
entries++
|
||||
if entries > maxArchiveFiles || header.Size < 0 || total+header.Size > maxArchiveBytes {
|
||||
return errors.New("task archive exceeds extraction safety limits")
|
||||
}
|
||||
total += header.Size
|
||||
clean := filepath.Clean(filepath.FromSlash(header.Name))
|
||||
if clean == "." || filepath.IsAbs(clean) || clean == ".." || strings.HasPrefix(clean, ".."+string(filepath.Separator)) {
|
||||
return fmt.Errorf("unsafe task archive path %q", header.Name)
|
||||
}
|
||||
target := filepath.Join(destination, clean)
|
||||
if relative, err := filepath.Rel(destination, target); err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
|
||||
return fmt.Errorf("unsafe task archive target %q", header.Name)
|
||||
}
|
||||
switch header.Typeflag {
|
||||
case tar.TypeDir:
|
||||
if err := os.MkdirAll(target, archiveDirMode); err != nil {
|
||||
return err
|
||||
}
|
||||
case tar.TypeReg, tar.TypeRegA:
|
||||
if err := os.MkdirAll(filepath.Dir(target), archiveDirMode); err != nil {
|
||||
return err
|
||||
}
|
||||
output, err := os.OpenFile(target, os.O_CREATE|os.O_EXCL|os.O_WRONLY, fs.FileMode(header.Mode)&0o777)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, copyErr := io.CopyN(output, tarReader, header.Size)
|
||||
closeErr := output.Close()
|
||||
if err := errors.Join(copyErr, closeErr); err != nil {
|
||||
return err
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("unsupported task archive entry type %d", header.Typeflag)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"archive/tar"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
pgdb "github.com/Autumn-27/artex/db"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
)
|
||||
|
||||
func TestTaskArchivePackageFilesRoundTrip(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
taskID := "42"
|
||||
explorationID := int64(73)
|
||||
taskFile := filepath.Join(dataDir, "tasks", taskID, "uploads", "evidence.txt")
|
||||
transcriptFile := filepath.Join(dataDir, "transcripts", "exp73-worker-1.jsonl")
|
||||
unrelatedTranscript := filepath.Join(dataDir, "transcripts", "exp74-worker-1.jsonl")
|
||||
for path, body := range map[string]string{
|
||||
taskFile: "task evidence",
|
||||
transcriptFile: "transcript",
|
||||
unrelatedTranscript: "leave me hot",
|
||||
} {
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
stage, err := stageTaskArchiveFiles(dataDir, 1, taskID, explorationID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(taskFile); !os.IsNotExist(err) {
|
||||
t.Fatalf("task file still hot after staging: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(transcriptFile); !os.IsNotExist(err) {
|
||||
t.Fatalf("task transcript still hot after staging: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(unrelatedTranscript); err != nil {
|
||||
t.Fatalf("unrelated transcript was staged: %v", err)
|
||||
}
|
||||
|
||||
archivePath := taskArchivePath(dataDir, 1, taskID)
|
||||
snapshot := &pgdb.TaskArchiveSnapshot{FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID}
|
||||
original, compressed, checksum, err := writeTaskArchivePackage(archivePath, stage.payload, snapshot)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if original == 0 || compressed == 0 || checksum == "" {
|
||||
t.Fatalf("invalid package metrics original=%d compressed=%d checksum=%q", original, compressed, checksum)
|
||||
}
|
||||
if err := stage.commit(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
extracted := filepath.Join(dataDir, "restore")
|
||||
if err := extractTaskArchivePackage(archivePath, checksum, extracted); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
installed, err := installTaskArchiveFiles(dataDir, extracted, taskID, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := installed.commit(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for path, want := range map[string]string{taskFile: "task evidence", transcriptFile: "transcript"} {
|
||||
got, err := os.ReadFile(path)
|
||||
if err != nil || string(got) != want {
|
||||
t.Fatalf("restored %s = %q, %v; want %q", path, got, err, want)
|
||||
}
|
||||
}
|
||||
if err := extractTaskArchivePackage(archivePath, "deadbeef", filepath.Join(dataDir, "bad-checksum")); err == nil {
|
||||
t.Fatal("checksum mismatch was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskArchiveStageJournalRollsBackInterruptedMoves(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
taskFile := filepath.Join(dataDir, "tasks", "42", "evidence.txt")
|
||||
transcriptFile := filepath.Join(dataDir, "transcripts", "exp73-worker.jsonl")
|
||||
for _, path := range []string{taskFile, transcriptFile} {
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, []byte(path), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
stage, err := stageTaskArchiveFiles(dataDir, 9, "42", 73)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, err := os.ReadFile(filepath.Join(stage.root, "journal.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var journal archiveStageJournal
|
||||
if err := json.Unmarshal(raw, &journal); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(journal.Moves) != 2 {
|
||||
t.Fatalf("journal moves=%d, want 2", len(journal.Moves))
|
||||
}
|
||||
recovered := &taskArchiveFileStage{root: stage.root, payload: stage.payload, journal: journal}
|
||||
if err := recovered.rollback(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, path := range []string{taskFile, transcriptFile} {
|
||||
if _, err := os.Stat(path); err != nil {
|
||||
t.Fatalf("interrupted archive did not restore %s: %v", path, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskArchiveRestoreJournalRollsBackInterruptedInstall(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
extracted := filepath.Join(dataDir, "archives", "tasks", ".restore", "11-test")
|
||||
source := filepath.Join(extracted, "files", "tasks", "42", "evidence.txt")
|
||||
if err := os.MkdirAll(filepath.Dir(source), 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(source, []byte("evidence"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err := installTaskArchiveFiles(dataDir, extracted, "42", 11)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
destination := filepath.Join(dataDir, "tasks", "42", "evidence.txt")
|
||||
if _, err := os.Stat(destination); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, err := os.ReadFile(filepath.Join(extracted, "restore-journal.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var journal archiveRestoreJournal
|
||||
if err := json.Unmarshal(raw, &journal); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if journal.ArchiveID != 11 || len(journal.Moves) != 1 {
|
||||
t.Fatalf("unexpected restore journal: %+v", journal)
|
||||
}
|
||||
recovered := &taskArchiveRestoreFiles{extracted: extracted, moves: journal.Moves}
|
||||
if err := recovered.rollback(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(destination); !os.IsNotExist(err) {
|
||||
t.Fatalf("hot destination remains after restore rollback: %v", err)
|
||||
}
|
||||
if got, err := os.ReadFile(source); err != nil || string(got) != "evidence" {
|
||||
t.Fatalf("restore source=%q err=%v", got, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskArchiveDeletePackageCanResumeFromStagedPath(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
archivePath := taskArchivePath(dataDir, 21, "42")
|
||||
if err := os.MkdirAll(filepath.Dir(archivePath), archiveDirMode); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
staged := archivePath + ".deleting-21"
|
||||
if err := os.WriteFile(staged, []byte("archive"), archiveFileMode); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, moved, err := stageTaskArchivePackageDelete(archivePath, 21)
|
||||
if err != nil || !moved || got != staged {
|
||||
t.Fatalf("resume staged package path=%q moved=%v err=%v", got, moved, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskArchivePackageRejectsTraversal(t *testing.T) {
|
||||
var compressed bytes.Buffer
|
||||
encoder, err := zstd.NewWriter(&compressed)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tw := tar.NewWriter(encoder)
|
||||
content := []byte("escape")
|
||||
if err := tw.WriteHeader(&tar.Header{Name: "../escape.txt", Mode: 0o600, Size: int64(len(content))}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := tw.Write(content); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tw.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := encoder.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
path := filepath.Join(t.TempDir(), "malicious.tar.zst")
|
||||
if err := os.WriteFile(path, compressed.Bytes(), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
root := t.TempDir()
|
||||
if err := extractTaskArchivePackage(path, "", root); err == nil {
|
||||
t.Fatal("path traversal archive was accepted")
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(filepath.Dir(root), "escape.txt")); !os.IsNotExist(err) {
|
||||
t.Fatalf("traversal wrote outside destination: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,656 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
pgdb "github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
const taskArchivePollInterval = 2 * time.Second
|
||||
|
||||
type archiveBatchRequest struct {
|
||||
TaskIDs []string `json:"task_ids"`
|
||||
ArchiveIDs []int64 `json:"archive_ids"`
|
||||
}
|
||||
|
||||
type archiveBatchItem struct {
|
||||
ID string `json:"id"`
|
||||
OK bool `json:"ok"`
|
||||
Queued bool `json:"queued"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Archive int64 `json:"archive_id,omitempty"`
|
||||
}
|
||||
|
||||
func (s *Server) startTaskArchiveWorker() {
|
||||
if s.m == nil || s.m.pg == nil {
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.RecoverTaskArchiveJobs(); err != nil {
|
||||
log.Printf("[task-archive] recover jobs: %v", err)
|
||||
}
|
||||
recoveryFailed := false
|
||||
if err := recoverTaskArchiveStages(s.m.dir, s.m.pg); err != nil {
|
||||
log.Printf("[task-archive] recover file staging: %v", err)
|
||||
recoveryFailed = true
|
||||
}
|
||||
if err := recoverTaskArchiveRestoreStages(s.m.dir, s.m.pg); err != nil {
|
||||
log.Printf("[task-archive] recover restore staging: %v", err)
|
||||
recoveryFailed = true
|
||||
}
|
||||
if err := recoverTaskArchiveDeletePackages(s.m.dir, s.m.pg); err != nil {
|
||||
log.Printf("[task-archive] recover delete packages: %v", err)
|
||||
recoveryFailed = true
|
||||
}
|
||||
if recoveryFailed {
|
||||
log.Printf("[task-archive] worker disabled because startup recovery is incomplete")
|
||||
return
|
||||
}
|
||||
s.archiveWG.Add(1)
|
||||
go func() {
|
||||
defer s.archiveWG.Done()
|
||||
ticker := time.NewTicker(taskArchivePollInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
if err := s.runOneTaskArchiveJob(); err != nil {
|
||||
log.Printf("[task-archive] worker: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
case <-s.archiveWake:
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}()
|
||||
s.notifyTaskArchiveWorker()
|
||||
}
|
||||
|
||||
func (s *Server) notifyTaskArchiveWorker() {
|
||||
if s == nil || s.archiveWake == nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case s.archiveWake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) runOneTaskArchiveJob() error {
|
||||
job, err := s.m.pg.ClaimTaskArchiveJob(s.ctx)
|
||||
if err != nil || job == nil {
|
||||
return err
|
||||
}
|
||||
var runErr error
|
||||
switch job.State {
|
||||
case pgdb.Archiving:
|
||||
runErr = s.archiveTask(job)
|
||||
case pgdb.Restoring:
|
||||
runErr = s.restoreTaskArchive(job)
|
||||
case pgdb.Deleting:
|
||||
runErr = s.deleteTaskArchive(job)
|
||||
default:
|
||||
runErr = fmt.Errorf("unknown claimed archive state %q", job.State)
|
||||
}
|
||||
if runErr != nil {
|
||||
if failErr := s.m.pg.FailTaskArchiveJob(job.ID, job.State, runErr); failErr != nil {
|
||||
return errors.Join(runErr, failErr)
|
||||
}
|
||||
return fmt.Errorf("archive job %d (%s): %w", job.ID, job.State, runErr)
|
||||
}
|
||||
// Drain another queued item without waiting for the periodic poll.
|
||||
s.notifyTaskArchiveWorker()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
|
||||
taskID := strconv.FormatInt(job.TaskID, 10)
|
||||
if !s.beginTaskDelete(taskID) {
|
||||
return errors.New("任务正在进行其他归档或删除操作")
|
||||
}
|
||||
committed := false
|
||||
defer func() {
|
||||
if !committed {
|
||||
s.abortTaskDelete(taskID)
|
||||
}
|
||||
}()
|
||||
s.cancelTaskChat(taskID, agent.AbortTaskDeleted)
|
||||
drainCtx, cancel := context.WithTimeout(s.ctx, taskDeleteDrainTimeout)
|
||||
defer cancel()
|
||||
if err := s.waitTaskQuiescent(drainCtx, taskID); err != nil {
|
||||
return errors.New("任务仍有运行中的 Agent,请先暂停后重试归档")
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_database", 10)
|
||||
snapshot, err := s.m.pg.SnapshotTaskArchive(job.TaskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_files", 25)
|
||||
fileStage, err := stageTaskArchiveFiles(s.m.dir, job.ID, taskID, snapshot.ExplorationID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
if !committed {
|
||||
runErr = errors.Join(runErr, fileStage.rollback())
|
||||
}
|
||||
}()
|
||||
if s.m.traffic != nil && len(snapshot.Hosts) > 0 {
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_traffic", 38)
|
||||
trafficCount, err := s.m.traffic.ExportHosts(snapshot.Hosts, filepath.Join(fileStage.payload, "traffic"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
snapshot.DataCounts["traffic"] = trafficCount
|
||||
}
|
||||
archivePath := taskArchivePath(s.m.dir, job.ID, taskID)
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compress", 50)
|
||||
originalSize, compressedSize, checksum, err := writeTaskArchivePackage(archivePath, fileStage.payload, snapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
removePackage := true
|
||||
defer func() {
|
||||
if removePackage {
|
||||
_ = os.Remove(archivePath)
|
||||
}
|
||||
}()
|
||||
var trafficStage *traffic.HostDeleteStage
|
||||
if s.m.traffic != nil && len(snapshot.ExclusiveHosts) > 0 {
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compact_traffic", 68)
|
||||
trafficStage, err = s.m.traffic.StageDeleteHostsExactForArchive(snapshot.ExclusiveHosts, job.ID, job.TaskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
if !committed {
|
||||
runErr = errors.Join(runErr, trafficStage.Rollback())
|
||||
}
|
||||
}()
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "compact_database", 78)
|
||||
if err := s.m.pg.CompleteTaskArchive(job.ID, snapshot, archivePath, checksum, originalSize, compressedSize); err != nil {
|
||||
return err
|
||||
}
|
||||
committed = true
|
||||
removePackage = false
|
||||
if trafficStage != nil {
|
||||
if err := trafficStage.Commit(); err != nil {
|
||||
warning := "归档已完成,但独占流量热存储清理失败(归档包仍可恢复):" + err.Error()
|
||||
log.Printf("[task-archive] task %s: %s", taskID, warning)
|
||||
_ = s.m.pg.AppendTaskArchiveWarning(job.ID, warning)
|
||||
}
|
||||
}
|
||||
if err := fileStage.commit(); err != nil {
|
||||
warning := "归档已完成,但暂存目录清理失败:" + err.Error()
|
||||
log.Printf("[task-archive] task %s: %s", taskID, warning)
|
||||
_ = s.m.pg.AppendTaskArchiveWarning(job.ID, warning)
|
||||
}
|
||||
s.engine.StopTask(taskID)
|
||||
s.taskAgentMu.Lock()
|
||||
delete(s.taskAgents, taskID)
|
||||
s.taskAgentMu.Unlock()
|
||||
s.m.forgetTask(taskID, job.TaskID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
|
||||
if err := validateArchivePath(s.m.dir, job.ArchivePath); err != nil {
|
||||
return err
|
||||
}
|
||||
// A prior attempt may have committed PostgreSQL and files, then failed only
|
||||
// while consuming the package. Finish that cleanup without replaying rows.
|
||||
if restored, err := s.m.pg.IsTaskArchiveRestored(job.ID); err != nil {
|
||||
return err
|
||||
} else if restored {
|
||||
if err := os.Remove(job.ArchivePath); err != nil && !os.IsNotExist(err) {
|
||||
return err
|
||||
}
|
||||
if err := s.m.pg.CompleteTaskArchiveRestore(job.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
s.loadRestoredTask(job.TaskID)
|
||||
return nil
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "verify_package", 10)
|
||||
restoreParent := filepath.Join(taskArchiveRoot(s.m.dir), ".restore")
|
||||
if err := os.MkdirAll(restoreParent, archiveDirMode); err != nil {
|
||||
return err
|
||||
}
|
||||
extracted, err := os.MkdirTemp(restoreParent, fmt.Sprintf("%d-", job.ID))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = os.RemoveAll(extracted) }()
|
||||
if err := extractTaskArchivePackage(job.ArchivePath, job.SHA256, extracted); err != nil {
|
||||
return err
|
||||
}
|
||||
manifestRaw, err := os.ReadFile(filepath.Join(extracted, "manifest.json"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var snapshot pgdb.TaskArchiveSnapshot
|
||||
if err := json.Unmarshal(manifestRaw, &snapshot); err != nil {
|
||||
return err
|
||||
}
|
||||
if snapshot.TaskID != job.TaskID || snapshot.FormatVersion != pgdb.TaskArchiveFormatVersion {
|
||||
return pgdb.ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
if s.m.traffic != nil {
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_traffic", 35)
|
||||
if _, err := s.m.traffic.ImportArchive(filepath.Join(extracted, "traffic")); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_files", 55)
|
||||
files, err := installTaskArchiveFiles(s.m.dir, extracted, strconv.FormatInt(job.TaskID, 10), job.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
databaseRestored := false
|
||||
defer func() {
|
||||
if !databaseRestored {
|
||||
runErr = errors.Join(runErr, files.rollback())
|
||||
}
|
||||
}()
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_database", 70)
|
||||
warnings, err := s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, warning := range warnings {
|
||||
log.Printf("[task-archive] restored task %d with warning: %s", job.TaskID, warning)
|
||||
}
|
||||
if len(warnings) > 0 {
|
||||
if _, err := s.m.pg.Exploration(snapshot.ExplorationID).AppendActivity(pgdb.Activity{
|
||||
Worker: "system", Kind: "system", Summary: "任务还原完成,但部分关联已降级", Detail: strings.Join(warnings, "\n"),
|
||||
}); err != nil {
|
||||
log.Printf("[task-archive] persist restore warnings for task %d: %v", job.TaskID, err)
|
||||
}
|
||||
}
|
||||
databaseRestored = true
|
||||
if err := files.commit(); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "consume_package", 95)
|
||||
if err := os.Remove(job.ArchivePath); err != nil && !os.IsNotExist(err) {
|
||||
return err
|
||||
}
|
||||
if err := s.m.pg.CompleteTaskArchiveRestore(job.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
s.loadRestoredTask(job.TaskID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) loadRestoredTask(taskID int64) {
|
||||
id := strconv.FormatInt(taskID, 10)
|
||||
var restored *Task
|
||||
for _, task := range s.m.LoadExisting() {
|
||||
if task.ID == id {
|
||||
restored = task
|
||||
break
|
||||
}
|
||||
}
|
||||
if restored == nil {
|
||||
if task, ok := s.m.Task(id); ok {
|
||||
restored = task
|
||||
}
|
||||
}
|
||||
if restored == nil {
|
||||
return
|
||||
}
|
||||
lifecycle := restored.lifecycleSnapshot()
|
||||
if lifecycle.Paused {
|
||||
s.engine.Pause(id, agent.AbortPausedOnReload)
|
||||
}
|
||||
if !isTerminalStatus(lifecycle.Status) {
|
||||
s.engine.startDeadlineCoordinator(s.ctx, restored)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) deleteTaskArchive(job *pgdb.TaskArchive) error {
|
||||
if err := validateArchivePath(s.m.dir, job.ArchivePath); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_package_delete", 35)
|
||||
staged, moved, err := stageTaskArchivePackageDelete(job.ArchivePath, job.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "delete_metadata", 70)
|
||||
if err := s.m.pg.DeleteTaskArchiveStub(job.ID); err != nil {
|
||||
if moved {
|
||||
_ = os.Rename(staged, job.ArchivePath)
|
||||
}
|
||||
return err
|
||||
}
|
||||
if moved {
|
||||
if err := os.Remove(staged); err != nil && !os.IsNotExist(err) {
|
||||
log.Printf("[task-archive] archived task %d metadata deleted; stale package cleanup failed: %v", job.TaskID, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateArchivePath(dataDir, candidate string) error {
|
||||
if strings.TrimSpace(candidate) == "" {
|
||||
return errors.New("归档包路径为空")
|
||||
}
|
||||
root, err := filepath.Abs(taskArchiveRoot(dataDir))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
path, err := filepath.Abs(candidate)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
relative, err := filepath.Rel(root, path)
|
||||
if err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
|
||||
return errors.New("归档包路径不在受管目录内")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) listTaskArchives(w http.ResponseWriter, r *http.Request) {
|
||||
page, _ := strconv.Atoi(r.URL.Query().Get("page"))
|
||||
size, _ := strconv.Atoi(r.URL.Query().Get("size"))
|
||||
result, err := s.m.pg.ListTaskArchives(r.URL.Query().Get("q"), r.URL.Query().Get("state"), page, size)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, result)
|
||||
}
|
||||
|
||||
func (s *Server) getTaskArchive(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "归档 id 无效")
|
||||
return
|
||||
}
|
||||
item, err := s.m.pg.GetTaskArchive(id)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if item == nil {
|
||||
writeErr(w, 404, "归档不存在")
|
||||
return
|
||||
}
|
||||
writeJSON(w, 200, item)
|
||||
}
|
||||
|
||||
func (s *Server) queueTaskArchive(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := canonicalTaskID(r.PathValue("id"))
|
||||
if !ok {
|
||||
writeErr(w, 400, "任务 id 无效")
|
||||
return
|
||||
}
|
||||
numeric, _ := strconv.ParseInt(id, 10, 64)
|
||||
item, err := s.m.pg.QueueTaskArchive(numeric)
|
||||
if err != nil {
|
||||
writeArchiveError(w, err)
|
||||
return
|
||||
}
|
||||
s.notifyTaskArchiveWorker()
|
||||
writeJSON(w, http.StatusAccepted, item)
|
||||
}
|
||||
|
||||
func (s *Server) queueTaskArchiveRestore(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "归档 id 无效")
|
||||
return
|
||||
}
|
||||
item, err := s.m.pg.QueueTaskArchiveRestore(id)
|
||||
if err != nil {
|
||||
writeArchiveError(w, err)
|
||||
return
|
||||
}
|
||||
s.notifyTaskArchiveWorker()
|
||||
writeJSON(w, http.StatusAccepted, item)
|
||||
}
|
||||
|
||||
func (s *Server) queueTaskArchiveDelete(w http.ResponseWriter, r *http.Request) {
|
||||
id, ok := pathInt(r, "id")
|
||||
if !ok {
|
||||
writeErr(w, 400, "归档 id 无效")
|
||||
return
|
||||
}
|
||||
item, err := s.m.pg.QueueTaskArchiveDelete(id)
|
||||
if err != nil {
|
||||
writeArchiveError(w, err)
|
||||
return
|
||||
}
|
||||
s.notifyTaskArchiveWorker()
|
||||
writeJSON(w, http.StatusAccepted, item)
|
||||
}
|
||||
|
||||
func (s *Server) queueTaskArchivesBatch(w http.ResponseWriter, r *http.Request) {
|
||||
var request archiveBatchRequest
|
||||
if err := decode(r, &request); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
parsed := normalizeBatchTaskIDs(request.TaskIDs)
|
||||
if len(parsed) == 0 {
|
||||
writeErr(w, 400, "task_ids 不能为空")
|
||||
return
|
||||
}
|
||||
if len(parsed) > 100 {
|
||||
writeErr(w, 400, "一次最多处理 100 个任务")
|
||||
return
|
||||
}
|
||||
ids := make([]string, 0, len(parsed))
|
||||
items := make([]archiveBatchItem, 0, len(parsed))
|
||||
for _, item := range parsed {
|
||||
if !item.valid {
|
||||
items = append(items, archiveBatchItem{ID: item.id, Error: "任务 id 无效"})
|
||||
continue
|
||||
}
|
||||
ids = append(ids, item.id)
|
||||
}
|
||||
ids = s.orderLiveTaskArchiveBatch(ids)
|
||||
for _, id := range ids {
|
||||
numeric, _ := strconv.ParseInt(id, 10, 64)
|
||||
archive, err := s.m.pg.QueueTaskArchive(numeric)
|
||||
item := archiveBatchItem{ID: id, OK: err == nil, Queued: err == nil}
|
||||
if err != nil {
|
||||
item.Error = err.Error()
|
||||
} else {
|
||||
item.Archive = archive.ID
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
s.notifyTaskArchiveWorker()
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) queueTaskArchiveActionBatch(w http.ResponseWriter, r *http.Request, action string) {
|
||||
var request archiveBatchRequest
|
||||
if err := decode(r, &request); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
ids, err := normalizeArchiveIDs(request.ArchiveIDs)
|
||||
if err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
ids = s.orderColdTaskArchiveBatch(ids, action == "restore")
|
||||
items := make([]archiveBatchItem, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
var archive *pgdb.TaskArchive
|
||||
if action == "restore" {
|
||||
archive, err = s.m.pg.QueueTaskArchiveRestore(id)
|
||||
} else {
|
||||
archive, err = s.m.pg.QueueTaskArchiveDelete(id)
|
||||
}
|
||||
item := archiveBatchItem{ID: strconv.FormatInt(id, 10), Archive: id, OK: err == nil, Queued: err == nil}
|
||||
if err != nil {
|
||||
item.Error = err.Error()
|
||||
} else if archive != nil {
|
||||
item.Archive = archive.ID
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
s.notifyTaskArchiveWorker()
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func (s *Server) restoreTaskArchivesBatch(w http.ResponseWriter, r *http.Request) {
|
||||
s.queueTaskArchiveActionBatch(w, r, "restore")
|
||||
}
|
||||
|
||||
func (s *Server) deleteTaskArchivesBatch(w http.ResponseWriter, r *http.Request) {
|
||||
s.queueTaskArchiveActionBatch(w, r, "delete")
|
||||
}
|
||||
|
||||
func normalizeArchiveIDs(ids []int64) ([]int64, error) {
|
||||
if len(ids) == 0 {
|
||||
return nil, errors.New("archive_ids 不能为空")
|
||||
}
|
||||
if len(ids) > 100 {
|
||||
return nil, errors.New("一次最多处理 100 个归档")
|
||||
}
|
||||
seen := map[int64]bool{}
|
||||
out := make([]int64, 0, len(ids))
|
||||
for _, id := range ids {
|
||||
if id <= 0 {
|
||||
return nil, fmt.Errorf("归档 id %d 无效", id)
|
||||
}
|
||||
if !seen[id] {
|
||||
seen[id] = true
|
||||
out = append(out, id)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *Server) orderLiveTaskArchiveBatch(ids []string) []string {
|
||||
sources := map[string][]string{}
|
||||
for _, id := range ids {
|
||||
if task, ok := s.m.Task(id); ok {
|
||||
for _, sourceID := range task.lifecycleSnapshot().SourceTaskIDs {
|
||||
sources[id] = append(sources[id], strconv.FormatInt(sourceID, 10))
|
||||
}
|
||||
}
|
||||
}
|
||||
return topoTaskIDs(ids, sources, false)
|
||||
}
|
||||
|
||||
func (s *Server) orderColdTaskArchiveBatch(ids []int64, restore bool) []int64 {
|
||||
taskToArchive := map[string]int64{}
|
||||
sources := map[string][]string{}
|
||||
var taskIDs []string
|
||||
for _, id := range ids {
|
||||
archive, err := s.m.pg.GetTaskArchive(id)
|
||||
if err != nil || archive == nil {
|
||||
continue
|
||||
}
|
||||
taskID := strconv.FormatInt(archive.TaskID, 10)
|
||||
taskToArchive[taskID] = id
|
||||
taskIDs = append(taskIDs, taskID)
|
||||
for _, sourceID := range archive.SourceTaskIDs {
|
||||
sources[taskID] = append(sources[taskID], strconv.FormatInt(sourceID, 10))
|
||||
}
|
||||
}
|
||||
orderedTasks := topoTaskIDs(taskIDs, sources, restore)
|
||||
ordered := make([]int64, 0, len(ids))
|
||||
seen := map[int64]bool{}
|
||||
for _, taskID := range orderedTasks {
|
||||
id := taskToArchive[taskID]
|
||||
ordered = append(ordered, id)
|
||||
seen[id] = true
|
||||
}
|
||||
for _, id := range ids {
|
||||
if !seen[id] {
|
||||
ordered = append(ordered, id)
|
||||
}
|
||||
}
|
||||
return ordered
|
||||
}
|
||||
|
||||
// topoTaskIDs orders source->dependent for restore and dependent->source for
|
||||
// archive/delete. Cycles are impossible at creation, but stable fallback keeps a
|
||||
// damaged database operable.
|
||||
func topoTaskIDs(ids []string, sources map[string][]string, restore bool) []string {
|
||||
set := map[string]bool{}
|
||||
for _, id := range ids {
|
||||
set[id] = true
|
||||
}
|
||||
adjacency := map[string][]string{}
|
||||
indegree := map[string]int{}
|
||||
for id := range set {
|
||||
indegree[id] = 0
|
||||
}
|
||||
for dependent, sourceIDs := range sources {
|
||||
for _, source := range sourceIDs {
|
||||
if !set[dependent] || !set[source] {
|
||||
continue
|
||||
}
|
||||
from, to := dependent, source
|
||||
if restore {
|
||||
from, to = source, dependent
|
||||
}
|
||||
adjacency[from] = append(adjacency[from], to)
|
||||
indegree[to]++
|
||||
}
|
||||
}
|
||||
queue := []string{}
|
||||
for _, id := range ids {
|
||||
if indegree[id] == 0 {
|
||||
queue = append(queue, id)
|
||||
}
|
||||
}
|
||||
ordered := make([]string, 0, len(ids))
|
||||
for len(queue) > 0 {
|
||||
id := queue[0]
|
||||
queue = queue[1:]
|
||||
ordered = append(ordered, id)
|
||||
for _, next := range adjacency[id] {
|
||||
indegree[next]--
|
||||
if indegree[next] == 0 {
|
||||
queue = append(queue, next)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(ordered) != len(ids) {
|
||||
seen := map[string]bool{}
|
||||
for _, id := range ordered {
|
||||
seen[id] = true
|
||||
}
|
||||
for _, id := range ids {
|
||||
if !seen[id] {
|
||||
ordered = append(ordered, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
return ordered
|
||||
}
|
||||
|
||||
func writeArchiveError(w http.ResponseWriter, err error) {
|
||||
status := http.StatusConflict
|
||||
if errors.Is(err, pgdb.ErrTaskArchiveNotFound) {
|
||||
status = http.StatusNotFound
|
||||
} else if !errors.Is(err, pgdb.ErrTaskArchiveIneligible) &&
|
||||
!errors.Is(err, pgdb.ErrTaskArchiveQueued) &&
|
||||
!errors.Is(err, pgdb.ErrTaskArchiveDependent) &&
|
||||
!errors.Is(err, pgdb.ErrTaskArchiveState) &&
|
||||
!errors.Is(err, pgdb.ErrTaskArchiveDeleteBlocked) {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
writeErr(w, status, err.Error())
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTaskArchiveAPIQueueListAndLimits(t *testing.T) {
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer m.Close()
|
||||
task, err := m.CreateTask("archive api", "verify 202", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
taskID, err := strconv.ParseInt(task.ID, 10, 64)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() {
|
||||
_, _ = m.pg.Exec(`DELETE FROM task_archives WHERE task_id=$1`, taskID)
|
||||
_ = m.pg.DeleteTask(taskID)
|
||||
}()
|
||||
|
||||
// Start with an already-cancelled context so the background archive worker
|
||||
// exits before these route-contract assertions enqueue work.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
dataDir := t.TempDir()
|
||||
s := New(ctx, m, dataDir, dataDir, dataDir)
|
||||
s.archiveWG.Wait()
|
||||
token, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := func(method, path, body string) *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
s.Handler().ServeHTTP(recorder, req)
|
||||
return recorder
|
||||
}
|
||||
|
||||
recorder := request(http.MethodPost, "/api/tasks/"+task.ID+"/archive", "{}")
|
||||
if recorder.Code != http.StatusConflict {
|
||||
t.Fatalf("running archive status=%d body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
if err := m.pg.SetPaused(taskID, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
recorder = request(http.MethodPost, "/api/tasks/"+task.ID+"/archive", "{}")
|
||||
if recorder.Code != http.StatusAccepted {
|
||||
t.Fatalf("archive queue status=%d body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
var queued struct {
|
||||
ID int64 `json:"id"`
|
||||
State string `json:"state"`
|
||||
}
|
||||
if err := json.Unmarshal(recorder.Body.Bytes(), &queued); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if queued.ID <= 0 || queued.State != "archive_queued" {
|
||||
t.Fatalf("unexpected queue response: %+v", queued)
|
||||
}
|
||||
|
||||
recorder = request(http.MethodGet, "/api/task-archives?page=1&size=10&q="+task.ID, "")
|
||||
if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), fmt.Sprintf(`"task_id":%d`, taskID)) {
|
||||
t.Fatalf("archive list status=%d body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
ids := make([]string, 101)
|
||||
for index := range ids {
|
||||
ids[index] = strconv.Itoa(index + 1)
|
||||
}
|
||||
body, _ := json.Marshal(map[string]any{"task_ids": ids})
|
||||
recorder = request(http.MethodPost, "/api/tasks/archive/batch", string(body))
|
||||
if recorder.Code != http.StatusBadRequest || !strings.Contains(recorder.Body.String(), "100") {
|
||||
t.Fatalf("oversized batch status=%d body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"iter"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -148,6 +149,22 @@ func TestTaskDTOUsesEmptyArrays(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyTaskArchiveBlocker(t *testing.T) {
|
||||
t.Parallel()
|
||||
dto := taskDTO(&Task{ID: "7"}, "paused")
|
||||
applyTaskArchiveBlocker(&dto, map[int64]int64{7: 11})
|
||||
if dto.ArchiveBlockedBy != "11" {
|
||||
t.Fatalf("archive blocker=%q, want 11", dto.ArchiveBlockedBy)
|
||||
}
|
||||
raw, err := json.Marshal(dto)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(raw), `"archive_blocked_by_task_id":"11"`) {
|
||||
t.Fatalf("archive blocker missing from payload: %s", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
|
||||
t.Parallel()
|
||||
task := &Task{ID: "7"}
|
||||
|
||||
Reference in new issue
Block a user