fix: prevent OOM when archiving large tasks
This commit is contained in:
1 parent
1fcc4b0588
commit
0ec45bb45d
11 files changed
+605
-87
No files matched your search
+13
-3
@@ -189,11 +189,21 @@ ALTER TABLE llm_records ADD COLUMN IF NOT EXISTS raw_response TEXT;
|
||||
|
||||
// EnsureLLMRecordsTable creates the llm_records table if it does not exist.
|
||||
func (d *DB) EnsureLLMRecordsTable() error {
|
||||
if _, err := d.Exec(llmRecordsSchema); err != nil {
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := d.Exec(llmRecordsMigrate)
|
||||
return err
|
||||
defer tx.Rollback() //nolint:errcheck
|
||||
if err := coordinateWithSchemaMigration(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(llmRecordsSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(llmRecordsMigrate); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// InsertLLMRecord stores one LLM call record.
|
||||
|
||||
@@ -3,19 +3,84 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
_ "embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/config"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
_ "github.com/jackc/pgx/v5/stdlib" // pgx database/sql driver ("pgx")
|
||||
)
|
||||
|
||||
//go:embed schema.sql
|
||||
var schemaSQL string
|
||||
|
||||
const schemaMigrationLockKey int64 = 7337741001
|
||||
|
||||
var schemaDeadlockRetryDelays = [...]time.Duration{
|
||||
100 * time.Millisecond,
|
||||
250 * time.Millisecond,
|
||||
500 * time.Millisecond,
|
||||
time.Second,
|
||||
}
|
||||
|
||||
type schemaExecer interface {
|
||||
ExecContext(context.Context, string, ...any) (sql.Result, error)
|
||||
}
|
||||
|
||||
func isPostgresDeadlock(err error) bool {
|
||||
var pgErr *pgconn.PgError
|
||||
return errors.As(err, &pgErr) && pgErr.Code == "40P01"
|
||||
}
|
||||
|
||||
func applySchemaWithRetry(ctx context.Context, execer schemaExecer, sleep func(time.Duration)) error {
|
||||
for attempt := 0; ; attempt++ {
|
||||
if _, err := execer.ExecContext(ctx, schemaSQL); err != nil {
|
||||
if !isPostgresDeadlock(err) || attempt >= len(schemaDeadlockRetryDelays) {
|
||||
return err
|
||||
}
|
||||
sleep(schemaDeadlockRetryDelays[attempt])
|
||||
continue
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// withSchemaMigrationLock pins the session-level lock to one checked-out
|
||||
// connection. Running pg_advisory_lock through *sql.DB is incorrect because a
|
||||
// later schema or unlock call may use a different pooled PostgreSQL session.
|
||||
func withSchemaMigrationLock(ctx context.Context, sqlDB *sql.DB, action func(*sql.Conn) error) (err error) {
|
||||
conn, err := sqlDB.Conn(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Close()
|
||||
if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1)`, schemaMigrationLockKey); err != nil {
|
||||
return fmt.Errorf("advisory lock: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if _, unlockErr := conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock($1)`, schemaMigrationLockKey); unlockErr != nil && err == nil {
|
||||
err = fmt.Errorf("advisory unlock: %w", unlockErr)
|
||||
}
|
||||
}()
|
||||
return action(conn)
|
||||
}
|
||||
|
||||
// coordinateWithSchemaMigration makes long, multi-table archive transactions
|
||||
// mutually exclusive with startup DDL while allowing ordinary runtime queries
|
||||
// to continue normally.
|
||||
func coordinateWithSchemaMigration(tx *sql.Tx) error {
|
||||
if _, err := tx.Exec(`SELECT pg_advisory_xact_lock($1)`, schemaMigrationLockKey); err != nil {
|
||||
return fmt.Errorf("coordinate with schema migration: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DSN resolves the PostgreSQL connection string and reports where it came from.
|
||||
// Precedence: env ARTEX_PG_DSN > config file (config.json). There is no
|
||||
// built-in default — it errors if neither source is configured.
|
||||
@@ -72,28 +137,22 @@ func Open(dsn string) (*DB, error) {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("ping postgres (%s): %w", config.Redact(dsn), err)
|
||||
}
|
||||
// Serialize schema migration + seeding across concurrent openers (e.g. parallel
|
||||
// test packages) using a session-level advisory lock. The lock is released
|
||||
// automatically when the connection is returned to the pool or closed.
|
||||
if _, err := sqlDB.Exec(`SELECT pg_advisory_lock(7337741001)`); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("advisory lock: %w", err)
|
||||
}
|
||||
// pgx runs multi-statement Exec via the simple protocol when there are no args.
|
||||
if _, err := sqlDB.Exec(schemaSQL); err != nil {
|
||||
sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("apply schema: %w", err)
|
||||
}
|
||||
d := &DB{sqlDB}
|
||||
if err := d.seedBuiltins(); err != nil {
|
||||
sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`) //nolint:errcheck
|
||||
// pgx runs multi-statement Exec via the simple protocol when there are no args.
|
||||
// Keep the dedicated lock connection checked out until both DDL and seeding
|
||||
// finish so concurrent application instances cannot initialize out of order.
|
||||
err = withSchemaMigrationLock(context.Background(), sqlDB, func(conn *sql.Conn) error {
|
||||
if err := applySchemaWithRetry(context.Background(), conn, time.Sleep); err != nil {
|
||||
return fmt.Errorf("apply schema: %w", err)
|
||||
}
|
||||
if err := d.seedBuiltins(); err != nil {
|
||||
return fmt.Errorf("seed builtins: %w", err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("seed builtins: %w", err)
|
||||
}
|
||||
if _, err := sqlDB.Exec(`SELECT pg_advisory_unlock(7337741001)`); err != nil {
|
||||
sqlDB.Close()
|
||||
return nil, fmt.Errorf("advisory unlock: %w", err)
|
||||
return nil, err
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
@@ -375,30 +434,30 @@ func (d *DB) seedDefaultInterceptRules() error {
|
||||
// 2. Python HTTP 客户端 .delete() 方法
|
||||
// 3. JS/通用脚本里的 method: 'DELETE' / method="DELETE"
|
||||
{
|
||||
name: "[内置] curl / wget 发送 DELETE 请求",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\bcurl\b[^|\n&;"]{0,300}(?:-X\s*DELETE|--request\s+DELETE|-XDELETE)|\bwget\b[^|\n&;"]{0,300}--method[=\s]+DELETE`,
|
||||
action: "deny",
|
||||
message: "禁止通过 curl/wget 发送 HTTP DELETE 请求,可能删除目标系统数据",
|
||||
name: "[内置] curl / wget 发送 DELETE 请求",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\bcurl\b[^|\n&;"]{0,300}(?:-X\s*DELETE|--request\s+DELETE|-XDELETE)|\bwget\b[^|\n&;"]{0,300}--method[=\s]+DELETE`,
|
||||
action: "deny",
|
||||
message: "禁止通过 curl/wget 发送 HTTP DELETE 请求,可能删除目标系统数据",
|
||||
priority: 80,
|
||||
},
|
||||
{
|
||||
name: "[内置] Python HTTP 客户端 DELETE(requests/httpx/aiohttp)",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\b(?:requests|httpx|aiohttp|urllib\.request)\.delete\s*\(|session\.delete\s*\(|client\.delete\s*\(`,
|
||||
action: "deny",
|
||||
message: "禁止使用 Python HTTP 客户端发送 DELETE 请求",
|
||||
name: "[内置] Python HTTP 客户端 DELETE(requests/httpx/aiohttp)",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)\b(?:requests|httpx|aiohttp|urllib\.request)\.delete\s*\(|session\.delete\s*\(|client\.delete\s*\(`,
|
||||
action: "deny",
|
||||
message: "禁止使用 Python HTTP 客户端发送 DELETE 请求",
|
||||
priority: 80,
|
||||
},
|
||||
{
|
||||
name: "[内置] 脚本中声明 HTTP DELETE 方法(JS/通用)",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)axios\.delete\s*\(|method\s*[:=]\s*['"]DELETE['"]`,
|
||||
action: "deny",
|
||||
message: "禁止在脚本中声明并发送 HTTP DELETE 请求",
|
||||
name: "[内置] 脚本中声明 HTTP DELETE 方法(JS/通用)",
|
||||
target: "tool_input",
|
||||
typ: "regex",
|
||||
pattern: `(?i)axios\.delete\s*\(|method\s*[:=]\s*['"]DELETE['"]`,
|
||||
action: "deny",
|
||||
message: "禁止在脚本中声明并发送 HTTP DELETE 请求",
|
||||
priority: 80,
|
||||
},
|
||||
{
|
||||
|
||||
@@ -1,9 +1,56 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
)
|
||||
|
||||
type schemaExecFunc func(context.Context, string, ...any) (sql.Result, error)
|
||||
|
||||
func (f schemaExecFunc) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
|
||||
return f(ctx, query, args...)
|
||||
}
|
||||
|
||||
func TestApplySchemaRetriesOnlyDeadlocks(t *testing.T) {
|
||||
deadlockAttempts := 0
|
||||
delays := []time.Duration{}
|
||||
err := applySchemaWithRetry(t.Context(), schemaExecFunc(func(context.Context, string, ...any) (sql.Result, error) {
|
||||
deadlockAttempts++
|
||||
if deadlockAttempts < 3 {
|
||||
return nil, &pgconn.PgError{Code: "40P01", Message: "deadlock detected"}
|
||||
}
|
||||
return nil, nil
|
||||
}), func(delay time.Duration) {
|
||||
delays = append(delays, delay)
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if deadlockAttempts != 3 {
|
||||
t.Fatalf("schema attempts=%d, want 3", deadlockAttempts)
|
||||
}
|
||||
if len(delays) != 2 || delays[0] != schemaDeadlockRetryDelays[0] || delays[1] != schemaDeadlockRetryDelays[1] {
|
||||
t.Fatalf("schema retry delays=%v", delays)
|
||||
}
|
||||
|
||||
nonRetryable := errors.New("permission denied")
|
||||
nonRetryableAttempts := 0
|
||||
err = applySchemaWithRetry(t.Context(), schemaExecFunc(func(context.Context, string, ...any) (sql.Result, error) {
|
||||
nonRetryableAttempts++
|
||||
return nil, nonRetryable
|
||||
}), func(time.Duration) {
|
||||
t.Fatal("non-deadlock error must not be retried")
|
||||
})
|
||||
if !errors.Is(err, nonRetryable) || nonRetryableAttempts != 1 {
|
||||
t.Fatalf("non-retryable result: attempts=%d err=%v", nonRetryableAttempts, err)
|
||||
}
|
||||
}
|
||||
|
||||
// testDSN returns the configured DSN, skipping the test when neither the env var
|
||||
// nor a config file supplies one (DSN no longer has a built-in default).
|
||||
func testDSN(t *testing.T) string {
|
||||
|
||||
+12
-2
@@ -49,8 +49,18 @@ CREATE INDEX IF NOT EXISTS idx_llm_usage_exp ON llm_usage(exploration_id);
|
||||
|
||||
// EnsureLLMUsageTable creates the llm_usage metering table if it does not exist.
|
||||
func (d *DB) EnsureLLMUsageTable() error {
|
||||
_, err := d.Exec(llmUsageSchema)
|
||||
return err
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck
|
||||
if err := coordinateWithSchemaMigration(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(llmUsageSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// InsertLLMUsage appends one metering row. Best-effort: callers log and continue on
|
||||
|
||||
+2
-1
@@ -452,7 +452,7 @@ CREATE TABLE IF NOT EXISTS task_archives (
|
||||
progress INTEGER NOT NULL DEFAULT 0 CHECK (progress BETWEEN 0 AND 100),
|
||||
error TEXT NOT NULL DEFAULT '',
|
||||
warnings JSONB NOT NULL DEFAULT '[]',
|
||||
format_version INTEGER NOT NULL DEFAULT 1,
|
||||
format_version INTEGER NOT NULL DEFAULT 2,
|
||||
archive_path TEXT NOT NULL DEFAULT '',
|
||||
sha256 TEXT NOT NULL DEFAULT '',
|
||||
original_size BIGINT NOT NULL DEFAULT 0,
|
||||
@@ -472,6 +472,7 @@ CREATE TABLE IF NOT EXISTS task_archives (
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
ALTER TABLE task_archives ALTER COLUMN format_version SET DEFAULT 2;
|
||||
CREATE INDEX IF NOT EXISTS idx_task_archives_state ON task_archives(state, requested_at, id);
|
||||
CREATE INDEX IF NOT EXISTS idx_task_archives_archived ON task_archives(archived_at DESC, id DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_task_archives_sources ON task_archives USING GIN(source_task_ids);
|
||||
|
||||
+106
-17
@@ -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()
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -36,6 +37,9 @@ func (d *DB) CompleteTaskArchive(
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck
|
||||
if err := coordinateWithSchemaMigration(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
var taskID, expID int64
|
||||
var state string
|
||||
if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state
|
||||
@@ -120,10 +124,11 @@ WHERE id=$1`, taskID); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`UPDATE task_archives SET
|
||||
state=$2,phase='ready',progress=100,error='',archive_path=$3,sha256=$4,
|
||||
original_size=$5,compressed_size=$6,data_counts=$7,aggregate_stats=$8,
|
||||
archived_at=now(),warnings='[]'
|
||||
WHERE id=$1`, archiveID, ArchiveReady, archivePath, sha256, originalSize, compressedSize, string(countsRaw), string(statsRaw)); err != nil {
|
||||
state=$2,phase='ready',progress=100,error='',archive_path=$3,sha256=$4,
|
||||
original_size=$5,compressed_size=$6,data_counts=$7,aggregate_stats=$8,
|
||||
format_version=$9,archived_at=now(),warnings='[]'
|
||||
WHERE id=$1`, archiveID, ArchiveReady, archivePath, sha256, originalSize, compressedSize,
|
||||
string(countsRaw), string(statsRaw), snapshot.FormatVersion); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
@@ -133,14 +138,49 @@ WHERE id=$1`, archiveID, ArchiveReady, archivePath, sha256, originalSize, compre
|
||||
// idempotent for accounting/traffic retry scenarios and returns non-fatal
|
||||
// warnings for global objects that intentionally are not recreated.
|
||||
func (d *DB) RestoreTaskArchive(archiveID int64, snapshot *TaskArchiveSnapshot, remainingTimeoutSeconds int64) ([]string, error) {
|
||||
if snapshot == nil || snapshot.FormatVersion != TaskArchiveFormatVersion {
|
||||
return d.restoreTaskArchive(archiveID, snapshot, remainingTimeoutSeconds, nil)
|
||||
}
|
||||
|
||||
// RestoreTaskArchiveWithLLMRecords restores a v2 package whose heavyweight LLM
|
||||
// record history is stored as a sequence of JSON objects outside manifest.json.
|
||||
func (d *DB) RestoreTaskArchiveWithLLMRecords(
|
||||
archiveID int64,
|
||||
snapshot *TaskArchiveSnapshot,
|
||||
remainingTimeoutSeconds int64,
|
||||
llmRecords io.Reader,
|
||||
) ([]string, error) {
|
||||
if llmRecords == nil {
|
||||
return nil, errors.New("nil streamed LLM record reader")
|
||||
}
|
||||
return d.restoreTaskArchive(archiveID, snapshot, remainingTimeoutSeconds, llmRecords)
|
||||
}
|
||||
|
||||
func (d *DB) restoreTaskArchive(
|
||||
archiveID int64,
|
||||
snapshot *TaskArchiveSnapshot,
|
||||
remainingTimeoutSeconds int64,
|
||||
llmRecords io.Reader,
|
||||
) ([]string, error) {
|
||||
if snapshot == nil || !IsTaskArchiveFormatSupported(snapshot.FormatVersion) {
|
||||
return nil, ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
streamedLLMRecords := snapshot.StreamedTables["llm_records"]
|
||||
if (len(snapshot.StreamedTables) > 0 && snapshot.FormatVersion < 2) ||
|
||||
len(snapshot.StreamedTables) > 1 ||
|
||||
(len(snapshot.StreamedTables) == 1 && streamedLLMRecords != TaskArchiveLLMRecordsPath) {
|
||||
return nil, fmt.Errorf("%w: unsupported streamed table metadata", ErrTaskArchiveFormatMismatch)
|
||||
}
|
||||
if streamedLLMRecords != "" && llmRecords == nil {
|
||||
return nil, fmt.Errorf("%w: streamed LLM records are missing", ErrTaskArchiveFormatMismatch)
|
||||
}
|
||||
tx, err := d.Begin()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck
|
||||
if err := coordinateWithSchemaMigration(tx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var taskID, expID int64
|
||||
var state string
|
||||
if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state
|
||||
@@ -215,7 +255,23 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID
|
||||
} else {
|
||||
warnings = append(warnings, warning...)
|
||||
}
|
||||
for _, table := range []string{"task_asset_links", "findings", "llm_records", "llm_usage", "skill_usage", "tool_usage"} {
|
||||
for _, table := range []string{"task_asset_links", "findings"} {
|
||||
if err := insertArchiveRows(tx, table, remappedTables[table]); err != nil {
|
||||
return nil, fmt.Errorf("restore %s: %w", table, err)
|
||||
}
|
||||
}
|
||||
if streamedLLMRecords != "" {
|
||||
count, err := insertArchiveJSONSequenceRows(tx, "llm_records", llmRecords)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("restore llm_records: %w", err)
|
||||
}
|
||||
if expected := snapshot.DataCounts["llm_records"]; count != expected {
|
||||
return nil, fmt.Errorf("restore llm_records: row count %d does not match manifest %d", count, expected)
|
||||
}
|
||||
} else if err := insertArchiveRows(tx, "llm_records", remappedTables["llm_records"]); err != nil {
|
||||
return nil, fmt.Errorf("restore llm_records: %w", err)
|
||||
}
|
||||
for _, table := range []string{"llm_usage", "skill_usage", "tool_usage"} {
|
||||
if err := insertArchiveRows(tx, table, remappedTables[table]); err != nil {
|
||||
return nil, fmt.Errorf("restore %s: %w", table, err)
|
||||
}
|
||||
@@ -306,6 +362,36 @@ func insertArchiveRows(tx *sql.Tx, table string, raw json.RawMessage) error {
|
||||
return err
|
||||
}
|
||||
|
||||
func insertArchiveJSONSequenceRows(tx *sql.Tx, table string, reader io.Reader) (int64, error) {
|
||||
if table != "llm_records" {
|
||||
return 0, fmt.Errorf("archive streamed restore table %q is not allowed", table)
|
||||
}
|
||||
statement, err := tx.Prepare(`INSERT INTO llm_records SELECT * FROM json_populate_record(NULL::llm_records,$1::json)`)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer statement.Close()
|
||||
decoder := json.NewDecoder(reader)
|
||||
var count int64
|
||||
for {
|
||||
var raw json.RawMessage
|
||||
if err := decoder.Decode(&raw); errors.Is(err, io.EOF) {
|
||||
break
|
||||
} else if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
trimmed := bytes.TrimSpace(raw)
|
||||
if len(trimmed) == 0 || trimmed[0] != '{' {
|
||||
return 0, errors.New("streamed archive row must be a JSON object")
|
||||
}
|
||||
if _, err := statement.Exec(string(trimmed)); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
count++
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func restoreArchiveAssets(tx *sql.Tx, taskID int64, raw json.RawMessage) (map[int64]int64, []string, error) {
|
||||
rows, err := decodeArchiveRows(raw)
|
||||
if err != nil {
|
||||
@@ -598,6 +684,9 @@ func (d *DB) DeleteTaskArchiveStub(archiveID int64) error {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback() //nolint:errcheck
|
||||
if err := coordinateWithSchemaMigration(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
var taskID, expID int64
|
||||
var state string
|
||||
if err := tx.QueryRow(`SELECT archive.task_id,task.exploration_id,archive.state
|
||||
|
||||
+150
-11
@@ -1,24 +1,55 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestQueryArchiveRowsStreamsJSON(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
type archiveJSONTestRows struct {
|
||||
values [][]byte
|
||||
next int
|
||||
}
|
||||
|
||||
func (r *archiveJSONTestRows) Next() bool {
|
||||
if r.next >= len(r.values) {
|
||||
return false
|
||||
}
|
||||
r.next++
|
||||
return true
|
||||
}
|
||||
|
||||
func (r *archiveJSONTestRows) Scan(dest ...any) error {
|
||||
if len(dest) != 1 || r.next == 0 || r.next > len(r.values) {
|
||||
return fmt.Errorf("invalid archive test row scan")
|
||||
}
|
||||
target, ok := dest[0].(*[]byte)
|
||||
if !ok {
|
||||
return fmt.Errorf("archive test row destination is %T", dest[0])
|
||||
}
|
||||
*target = append((*target)[:0], r.values[r.next-1]...)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *archiveJSONTestRows) Err() error { return nil }
|
||||
|
||||
func TestQueryArchiveRowsStreamsJSON(t *testing.T) {
|
||||
payload := strings.Repeat("large request/response payload ", 64*1024)
|
||||
raw, count, err := queryArchiveRows(d, `SELECT * FROM (VALUES
|
||||
($1::bigint,$2::text),($3::bigint,$4::text)) archived_row(id,body) ORDER BY id`,
|
||||
int64(1), payload, int64(2), "quoted: \"value\"\nline")
|
||||
values := make([][]byte, 2)
|
||||
var err error
|
||||
values[0], err = json.Marshal(map[string]any{"id": int64(1), "body": payload})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
values[1], err = json.Marshal(map[string]any{"id": int64(2), "body": "quoted: \"value\"\nline"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, count, err := encodeArchiveRows(&archiveJSONTestRows{values: values})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -40,6 +71,55 @@ func TestQueryArchiveRowsStreamsJSON(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteArchiveRowsProducesJSONSequence(t *testing.T) {
|
||||
values := [][]byte{
|
||||
json.RawMessage(`{"id":1,"body":"first"}`),
|
||||
json.RawMessage(`{"id":2,"body":"second\\nline"}`),
|
||||
}
|
||||
var output bytes.Buffer
|
||||
count, err := writeArchiveRows(&archiveJSONTestRows{values: values}, &output)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if count != 2 {
|
||||
t.Fatalf("streamed row count=%d, want 2", count)
|
||||
}
|
||||
decoder := json.NewDecoder(&output)
|
||||
for wantID := int64(1); wantID <= 2; wantID++ {
|
||||
var row struct {
|
||||
ID int64 `json:"id"`
|
||||
}
|
||||
if err := decoder.Decode(&row); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if row.ID != wantID {
|
||||
t.Fatalf("streamed row id=%d, want %d", row.ID, wantID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskArchiveFormatCompatibility(t *testing.T) {
|
||||
for _, version := range []int{TaskArchiveLegacyFormatVersion, TaskArchiveFormatVersion} {
|
||||
if !IsTaskArchiveFormatSupported(version) {
|
||||
t.Fatalf("archive format %d should be supported", version)
|
||||
}
|
||||
}
|
||||
for _, version := range []int{0, TaskArchiveFormatVersion + 1} {
|
||||
if IsTaskArchiveFormatSupported(version) {
|
||||
t.Fatalf("archive format %d should be rejected", version)
|
||||
}
|
||||
}
|
||||
invalidSnapshots := []*TaskArchiveSnapshot{
|
||||
{FormatVersion: TaskArchiveLegacyFormatVersion, StreamedTables: map[string]string{"llm_records": TaskArchiveLLMRecordsPath}},
|
||||
{FormatVersion: TaskArchiveFormatVersion, StreamedTables: map[string]string{"unknown": "database/unknown.ndjson"}},
|
||||
}
|
||||
for _, snapshot := range invalidSnapshots {
|
||||
if _, err := (&DB{}).RestoreTaskArchive(1, snapshot, 0); !errors.Is(err, ErrTaskArchiveFormatMismatch) {
|
||||
t.Fatalf("invalid streamed table metadata returned %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskArchiveDatabaseRoundTrip(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
@@ -111,6 +191,12 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro
|
||||
if err := d.InsertToolUsage(&ToolUsage{ToolKey: toolName, AgentKey: "worker", TaskID: task.ID, ExplorationID: task.ExplorationID}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := d.InsertLLMRecord(&LLMRecord{
|
||||
TaskID: fmt.Sprint(task.ID), Model: "archive-model", SessionID: "archive-session",
|
||||
Status: "ok", RawRequest: strings.Repeat("request", 1024), RawResponse: strings.Repeat("response", 1024),
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := d.AddFinding(task.ID, nodeID, vulnclass, "archive finding", SeverityCritical, "summary", "evidence", "worker", []int64{assetID}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -122,10 +208,17 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro
|
||||
if err != nil || claimed == nil || claimed.ID != archive.ID || claimed.State != Archiving {
|
||||
t.Fatalf("claim = %+v, %v", claimed, err)
|
||||
}
|
||||
snapshot, err := d.SnapshotTaskArchive(task.ID)
|
||||
var llmRecords bytes.Buffer
|
||||
snapshot, err := d.SnapshotTaskArchiveWithLLMRecords(task.ID, &llmRecords)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if snapshot.StreamedTables["llm_records"] != TaskArchiveLLMRecordsPath || snapshot.DataCounts["llm_records"] != 1 {
|
||||
t.Fatalf("unexpected streamed LLM metadata: paths=%v counts=%v", snapshot.StreamedTables, snapshot.DataCounts)
|
||||
}
|
||||
if rawRowCount(snapshot.Tables["llm_records"]) != 0 {
|
||||
t.Fatal("streamed LLM records were also retained in manifest memory")
|
||||
}
|
||||
if snapshot.DataCounts["assets"] != 1 || snapshot.DataCounts["exploration_nodes"] < 2 {
|
||||
t.Fatalf("unexpected snapshot counts: %#v", snapshot.DataCounts)
|
||||
}
|
||||
@@ -170,7 +263,9 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro
|
||||
if err != nil || claimed == nil || claimed.State != Restoring {
|
||||
t.Fatalf("restore claim = %+v, %v", claimed, err)
|
||||
}
|
||||
warnings, err := d.RestoreTaskArchive(archive.ID, snapshot, int64((10 * time.Minute).Seconds()))
|
||||
warnings, err := d.RestoreTaskArchiveWithLLMRecords(
|
||||
archive.ID, snapshot, int64((10 * time.Minute).Seconds()), bytes.NewReader(llmRecords.Bytes()),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -202,6 +297,14 @@ VALUES($1,$2,0,'quota_exhausted','balance exhausted',$3,$4,$3)`, task.ID, llmPro
|
||||
if nodes < 2 || assets != 1 || usage != 1 {
|
||||
t.Fatalf("restore incomplete nodes=%d assets=%d usage=%d", nodes, assets, usage)
|
||||
}
|
||||
var restoredRawRequest, restoredRawResponse string
|
||||
if err := d.QueryRow(`SELECT COALESCE(raw_request,''),COALESCE(raw_response,'') FROM llm_records WHERE task_id=$1`, fmt.Sprint(task.ID)).
|
||||
Scan(&restoredRawRequest, &restoredRawResponse); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if restoredRawRequest != strings.Repeat("request", 1024) || restoredRawResponse != strings.Repeat("response", 1024) {
|
||||
t.Fatal("restored streamed LLM record body was truncated")
|
||||
}
|
||||
var restoredCompanyID *int64
|
||||
var restoredStatus, restoredError string
|
||||
var restoredExhaustedAt, restoredCreatedAt time.Time
|
||||
@@ -274,6 +377,42 @@ func TestTaskArchiveBlockersIgnoreQueuedDependents(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecoverInterruptedArchiveRequiresManualRetry(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
task, err := d.CreateTaskWithOptions("interrupted archive", "must not restart automatically", TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() {
|
||||
_, _ = d.Exec(`DELETE FROM task_archives WHERE task_id=$1`, task.ID)
|
||||
_ = d.DeleteTask(task.ID)
|
||||
}()
|
||||
if err := d.SetPaused(task.ID, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
queued, err := d.QueueTaskArchive(task.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := d.Exec(`UPDATE task_archives SET state=$2,phase='snapshot_database',progress=10 WHERE id=$1`, queued.ID, Archiving); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := d.RecoverTaskArchiveJobs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
recovered, err := d.GetTaskArchive(queued.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if recovered == nil || recovered.State != ArchiveFailed || recovered.Phase != "interrupted" || recovered.Error == "" {
|
||||
t.Fatalf("recovered archive=%+v, want explicit-retry failure", recovered)
|
||||
}
|
||||
}
|
||||
|
||||
func assertArchiveGlobalStats(t *testing.T, d *DB, profileName, skillName, toolName, vulnclass string) {
|
||||
t.Helper()
|
||||
profiles, err := d.UsageByProfile()
|
||||
|
||||
@@ -480,11 +480,20 @@ func writeTaskArchivePackage(path, payloadDir string, snapshot *pgdb.TaskArchive
|
||||
if err = os.MkdirAll(filepath.Dir(path), archiveDirMode); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
manifest, err := json.Marshal(snapshot)
|
||||
manifestFile, err := os.OpenFile(
|
||||
filepath.Join(payloadDir, "manifest.json"),
|
||||
os.O_CREATE|os.O_TRUNC|os.O_WRONLY,
|
||||
archiveFileMode,
|
||||
)
|
||||
if err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
if err = os.WriteFile(filepath.Join(payloadDir, "manifest.json"), manifest, archiveFileMode); err != nil {
|
||||
encodeErr := json.NewEncoder(manifestFile).Encode(snapshot)
|
||||
var syncErr error
|
||||
if encodeErr == nil {
|
||||
syncErr = manifestFile.Sync()
|
||||
}
|
||||
if err := errors.Join(encodeErr, syncErr, manifestFile.Close()); err != nil {
|
||||
return 0, 0, "", err
|
||||
}
|
||||
temporary := path + ".partial"
|
||||
|
||||
@@ -47,7 +47,18 @@ func TestTaskArchivePackageFilesRoundTrip(t *testing.T) {
|
||||
}
|
||||
|
||||
archivePath := taskArchivePath(dataDir, 1, taskID)
|
||||
snapshot := &pgdb.TaskArchiveSnapshot{FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID}
|
||||
streamPath := filepath.Join(stage.payload, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath))
|
||||
if err := os.MkdirAll(filepath.Dir(streamPath), archiveDirMode); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
streamBody := []byte("{\"id\":1,\"raw_request\":\"large\"}\n")
|
||||
if err := os.WriteFile(streamPath, streamBody, archiveFileMode); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snapshot := &pgdb.TaskArchiveSnapshot{
|
||||
FormatVersion: pgdb.TaskArchiveFormatVersion, TaskID: 42, ExplorationID: explorationID,
|
||||
StreamedTables: map[string]string{"llm_records": pgdb.TaskArchiveLLMRecordsPath},
|
||||
}
|
||||
original, compressed, checksum, err := writeTaskArchivePackage(archivePath, stage.payload, snapshot)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
@@ -62,6 +73,9 @@ func TestTaskArchivePackageFilesRoundTrip(t *testing.T) {
|
||||
if err := extractTaskArchivePackage(archivePath, checksum, extracted); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, err := os.ReadFile(filepath.Join(extracted, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath))); err != nil || !bytes.Equal(got, streamBody) {
|
||||
t.Fatalf("streamed LLM archive payload=%q err=%v", got, err)
|
||||
}
|
||||
installed, err := installTaskArchiveFiles(dataDir, extracted, taskID, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
||||
+57
-6
@@ -131,13 +131,15 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
|
||||
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)
|
||||
task, err := s.m.pg.GetTask(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 task == nil {
|
||||
return pgdb.ErrTaskArchiveNotFound
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "stage_files", 10)
|
||||
fileStage, err := stageTaskArchiveFiles(s.m.dir, job.ID, taskID, task.ExplorationID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -146,6 +148,27 @@ func (s *Server) archiveTask(job *pgdb.TaskArchive) (runErr error) {
|
||||
runErr = errors.Join(runErr, fileStage.rollback())
|
||||
}
|
||||
}()
|
||||
streamPath := filepath.Join(fileStage.payload, filepath.FromSlash(pgdb.TaskArchiveLLMRecordsPath))
|
||||
if err := os.MkdirAll(filepath.Dir(streamPath), archiveDirMode); err != nil {
|
||||
return err
|
||||
}
|
||||
streamFile, err := os.OpenFile(streamPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, archiveFileMode)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "snapshot_database", 20)
|
||||
snapshot, snapshotErr := s.m.pg.SnapshotTaskArchiveWithLLMRecords(job.TaskID, streamFile)
|
||||
var syncErr error
|
||||
if snapshotErr == nil {
|
||||
syncErr = streamFile.Sync()
|
||||
}
|
||||
closeErr := streamFile.Close()
|
||||
if err := errors.Join(snapshotErr, syncErr, closeErr); err != nil {
|
||||
return err
|
||||
}
|
||||
if snapshot.ExplorationID != task.ExplorationID {
|
||||
return errors.New("任务探索记录在归档快照期间发生变化")
|
||||
}
|
||||
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"))
|
||||
@@ -244,9 +267,27 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
|
||||
if err := json.Unmarshal(manifestRaw, &snapshot); err != nil {
|
||||
return err
|
||||
}
|
||||
if snapshot.TaskID != job.TaskID || snapshot.FormatVersion != pgdb.TaskArchiveFormatVersion {
|
||||
if snapshot.TaskID != job.TaskID || !pgdb.IsTaskArchiveFormatSupported(snapshot.FormatVersion) {
|
||||
return pgdb.ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
var llmRecordsFile *os.File
|
||||
relative, hasStreamedLLMRecords := snapshot.StreamedTables["llm_records"]
|
||||
if len(snapshot.StreamedTables) > 0 &&
|
||||
(snapshot.FormatVersion < 2 || !hasStreamedLLMRecords || len(snapshot.StreamedTables) != 1 ||
|
||||
relative != pgdb.TaskArchiveLLMRecordsPath) {
|
||||
return pgdb.ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
if hasStreamedLLMRecords {
|
||||
llmRecordsFile, err = os.Open(filepath.Join(extracted, filepath.FromSlash(relative)))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
if llmRecordsFile != nil {
|
||||
_ = llmRecordsFile.Close()
|
||||
}
|
||||
}()
|
||||
}
|
||||
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 {
|
||||
@@ -265,7 +306,17 @@ func (s *Server) restoreTaskArchive(job *pgdb.TaskArchive) (runErr error) {
|
||||
}
|
||||
}()
|
||||
_ = s.m.pg.UpdateTaskArchiveProgress(job.ID, "restore_database", 70)
|
||||
warnings, err := s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
|
||||
var warnings []string
|
||||
if llmRecordsFile != nil {
|
||||
warnings, err = s.m.pg.RestoreTaskArchiveWithLLMRecords(
|
||||
job.ID, &snapshot, job.RemainingTimeoutSeconds, llmRecordsFile,
|
||||
)
|
||||
closeErr := llmRecordsFile.Close()
|
||||
llmRecordsFile = nil
|
||||
err = errors.Join(err, closeErr)
|
||||
} else {
|
||||
warnings, err = s.m.pg.RestoreTaskArchive(job.ID, &snapshot, job.RemainingTimeoutSeconds)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user