admin/ARTEX
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: prevent OOM when archiving large tasks

Browse Source
This commit is contained in:
pphua committed 2026-08-28 15:25:20 +08:00
1 parent 1fcc4b0588
commit 0ec45bb45d
11 files changed
+605 -87
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

db/task_archives.go
+106 -17
View File
@@ -7,6 +7,7 @@ import (
"encoding/json"
"errors"
"fmt"
"io"
"net/url"
"sort"
"strconv"
@@ -14,7 +15,15 @@ import (
"time"
)
const TaskArchiveFormatVersion = 1
const (
TaskArchiveFormatVersion = 2
TaskArchiveLegacyFormatVersion = 1
TaskArchiveLLMRecordsPath = "database/llm_records.ndjson"
)
func IsTaskArchiveFormatSupported(version int) bool {
return version >= TaskArchiveLegacyFormatVersion && version <= TaskArchiveFormatVersion
}
const (
ArchiveQueued = "archive_queued"
@@ -103,8 +112,8 @@ type TaskArchivePage struct {
}
// TaskArchiveSnapshot is serialized into manifest.json inside the cold package.
// Each value in Tables is a JSON array generated by PostgreSQL row_to_json, which
// preserves exact IDs, timestamps, arrays and JSON payloads for restoration.
// Small tables remain JSON arrays in Tables. Large v2 tables are streamed to
// package files listed in StreamedTables so their size is not bounded by memory.
type TaskArchiveSnapshot struct {
FormatVersion int `json:"format_version"`
CreatedAt time.Time `json:"created_at"`
@@ -115,6 +124,7 @@ type TaskArchiveSnapshot struct {
ExclusiveHosts []string `json:"exclusive_hosts"`
ExclusiveAssetIDs []int64 `json:"exclusive_asset_ids"`
Tables map[string]json.RawMessage `json:"tables"`
StreamedTables map[string]string `json:"streamed_tables,omitempty"`
DataCounts map[string]int64 `json:"data_counts"`
AggregateStats map[string]any `json:"aggregate_stats"`
}
@@ -279,12 +289,13 @@ LIMIT 1`, taskID).Scan(&dependent)
task_description,task_goal,original_status,category_id_snapshot,
category_name_snapshot,source_task_ids,remaining_timeout_seconds,requested_at)
VALUES ($1,$2,'queued',0,'','[]',$3,$4,$5,$6,$7,$8,$9,$10,$11,now())
ON CONFLICT (task_id) DO UPDATE SET
state=CASE WHEN task_archives.state IN ('archive_failed') THEN EXCLUDED.state ELSE task_archives.state END,
phase=CASE WHEN task_archives.state IN ('archive_failed') THEN 'queued' ELSE task_archives.phase END,
progress=CASE WHEN task_archives.state IN ('archive_failed') THEN 0 ELSE task_archives.progress END,
error=CASE WHEN task_archives.state IN ('archive_failed') THEN '' ELSE task_archives.error END,
requested_at=CASE WHEN task_archives.state IN ('archive_failed') THEN now() ELSE task_archives.requested_at END`,
ON CONFLICT (task_id) DO UPDATE SET
state=CASE WHEN task_archives.state IN ('archive_failed') THEN EXCLUDED.state ELSE task_archives.state END,
phase=CASE WHEN task_archives.state IN ('archive_failed') THEN 'queued' ELSE task_archives.phase END,
progress=CASE WHEN task_archives.state IN ('archive_failed') THEN 0 ELSE task_archives.progress END,
error=CASE WHEN task_archives.state IN ('archive_failed') THEN '' ELSE task_archives.error END,
format_version=CASE WHEN task_archives.state IN ('archive_failed') THEN EXCLUDED.format_version ELSE task_archives.format_version END,
requested_at=CASE WHEN task_archives.state IN ('archive_failed') THEN now() ELSE task_archives.requested_at END`,
taskID, ArchiveQueued, TaskArchiveFormatVersion, name, description, goal, status,
categoryID, categoryName, sources, remaining)
if err != nil {
@@ -341,15 +352,17 @@ WHERE id=$1 RETURNING `+taskArchiveCols, id, DeleteQueued))
return item, tx.Commit()
}
// RecoverTaskArchiveJobs makes active operations retryable after an unclean
// shutdown. Idempotent restore/import semantics keep already-completed pieces safe.
// RecoverTaskArchiveJobs keeps restore/delete resumable after an unclean shutdown.
// An interrupted archive requires an explicit retry: automatic startup retries can
// otherwise form a crash loop when the prior process was killed by resource limits.
func (d *DB) RecoverTaskArchiveJobs() error {
_, err := d.Exec(`UPDATE task_archives SET
state=CASE state WHEN 'archiving' THEN 'archive_queued'
WHEN 'restoring' THEN 'restore_queued'
WHEN 'deleting' THEN 'delete_queued' ELSE state END,
phase='recovered', error=''
WHERE state IN ('archiving','restoring','deleting')`)
state=CASE state WHEN 'archiving' THEN 'archive_failed'
WHEN 'restoring' THEN 'restore_queued'
WHEN 'deleting' THEN 'delete_queued' ELSE state END,
phase='interrupted',
error=CASE WHEN state='archiving' THEN '上次归档进程异常退出，请手动重试' ELSE '' END
WHERE state IN ('archiving','restoring','deleting')`)
return err
}
@@ -435,7 +448,14 @@ func queryArchiveRows(q interface {
return nil, 0, err
}
defer rows.Close()
return encodeArchiveRows(rows)
}
func encodeArchiveRows(rows interface {
Next() bool
Scan(dest ...any) error
Err() error
}) (json.RawMessage, int64, error) {
var output bytes.Buffer
output.WriteByte('[')
var count int64
@@ -457,6 +477,46 @@ func queryArchiveRows(q interface {
return json.RawMessage(output.Bytes()), count, nil
}
func writeArchiveRows(rows interface {
Next() bool
Scan(dest ...any) error
Err() error
}, writer io.Writer) (int64, error) {
var count int64
for rows.Next() {
var raw []byte
if err := rows.Scan(&raw); err != nil {
return 0, err
}
if written, err := writer.Write(raw); err != nil {
return 0, err
} else if written != len(raw) {
return 0, io.ErrShortWrite
}
if written, err := io.WriteString(writer, "\n"); err != nil {
return 0, err
} else if written != 1 {
return 0, io.ErrShortWrite
}
count++
}
if err := rows.Err(); err != nil {
return 0, err
}
return count, nil
}
func streamArchiveRows(q interface {
Query(query string, args ...any) (*sql.Rows, error)
}, writer io.Writer, inner string, args ...any) (int64, error) {
rows, err := q.Query(`SELECT row_to_json(row_data)::text FROM (`+inner+`) row_data`, args...)
if err != nil {
return 0, err
}
defer rows.Close()
return writeArchiveRows(rows, writer)
}
func rawRowCount(raw json.RawMessage) int64 {
var rows []json.RawMessage
if json.Unmarshal(raw, &rows) != nil {
@@ -480,11 +540,28 @@ UNION SELECT value::bigint FROM findings finding
// writes are already quiescent at the server barrier; repeatable-read also keeps
// the asset and accounting views mutually consistent during serialization.
func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) {
return d.snapshotTaskArchive(taskID, nil)
}
// SnapshotTaskArchiveWithLLMRecords streams the heavyweight record history to
// llmRecords while all other task-owned data is read from the same repeatable
// PostgreSQL snapshot.
func (d *DB) SnapshotTaskArchiveWithLLMRecords(taskID int64, llmRecords io.Writer) (*TaskArchiveSnapshot, error) {
if llmRecords == nil {
return nil, errors.New("nil LLM record archive writer")
}
return d.snapshotTaskArchive(taskID, llmRecords)
}
func (d *DB) snapshotTaskArchive(taskID int64, llmRecords io.Writer) (*TaskArchiveSnapshot, error) {
tx, err := d.BeginTx(context.Background(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead, ReadOnly: true})
if err != nil {
return nil, err
}
defer tx.Rollback() //nolint:errcheck
if err := coordinateWithSchemaMigration(tx); err != nil {
return nil, err
}
var expID int64
if err := tx.QueryRow(`SELECT exploration_id FROM tasks WHERE id=$1 AND deleted_at IS NULL`, taskID).Scan(&expID); err != nil {
if errors.Is(err, sql.ErrNoRows) {
@@ -518,7 +595,18 @@ func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) {
{"assets", `SELECT asset.* FROM assets asset WHERE asset.id IN (` + archiveAssetIDsQuery() + `) ORDER BY asset.id`, []any{taskID, expID}},
}
counts := make(map[string]int64, len(queries))
streamedTables := map[string]string{}
for _, query := range queries {
if query.name == "llm_records" && llmRecords != nil {
count, err := streamArchiveRows(tx, llmRecords, query.query, query.args...)
if err != nil {
return nil, fmt.Errorf("snapshot %s: %w", query.name, err)
}
tables[query.name] = json.RawMessage("[]")
counts[query.name] = count
streamedTables[query.name] = TaskArchiveLLMRecordsPath
continue
}
raw, count, err := queryArchiveRows(tx, query.query, query.args...)
if err != nil {
return nil, fmt.Errorf("snapshot %s: %w", query.name, err)
@@ -555,7 +643,8 @@ func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) {
snapshot := &TaskArchiveSnapshot{
FormatVersion: TaskArchiveFormatVersion, CreatedAt: time.Now().UTC(), TaskID: taskID,
ExplorationID: expID, SourceTaskIDs: sources, Hosts: hosts, ExclusiveHosts: exclusiveHosts,
ExclusiveAssetIDs: exclusiveAssetIDs, Tables: tables, DataCounts: counts, AggregateStats: stats,
ExclusiveAssetIDs: exclusiveAssetIDs, Tables: tables, StreamedTables: streamedTables,
DataCounts: counts, AggregateStats: stats,
}
return snapshot, tx.Commit()
}