diff --git a/db/commands.go b/db/commands.go index 27172da..9e6e2e6 100644 --- a/db/commands.go +++ b/db/commands.go @@ -133,11 +133,21 @@ ALTER TABLE llm_records ADD COLUMN IF NOT EXISTS profile_name 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. diff --git a/db/db.go b/db/db.go index de0a738..7bf2c36 100644 --- a/db/db.go +++ b/db/db.go @@ -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 } diff --git a/db/db_test.go b/db/db_test.go index d549cdd..41c5c05 100644 --- a/db/db_test.go +++ b/db/db_test.go @@ -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 { diff --git a/db/llm_usage.go b/db/llm_usage.go index 21b3b8c..25d3abe 100644 --- a/db/llm_usage.go +++ b/db/llm_usage.go @@ -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 diff --git a/db/schema.sql b/db/schema.sql index de0b7cd..5c6356d 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -448,7 +448,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, @@ -468,6 +468,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); diff --git a/db/task_archives.go b/db/task_archives.go index bef9f3a..2093975 100644 --- a/db/task_archives.go +++ b/db/task_archives.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/url" "sort" "strconv" @@ -13,7 +14,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" @@ -102,8 +111,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"` @@ -114,6 +123,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"` } @@ -278,12 +288,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 { @@ -340,15 +351,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 } @@ -423,11 +436,92 @@ func (d *DB) FailTaskArchiveJob(id int64, activeState string, cause error) error } func queryArchiveRows(q interface { +<<<<<<< Updated upstream QueryRow(query string, args ...any) *sql.Row }, inner string, args ...any) (json.RawMessage, error) { var raw []byte err := q.QueryRow(`SELECT COALESCE(jsonb_agg(to_jsonb(row_data)), '[]'::jsonb) FROM (`+inner+`) row_data`, args...).Scan(&raw) return json.RawMessage(raw), err +======= + Query(query string, args ...any) (*sql.Rows, error) +}, inner string, args ...any) (json.RawMessage, int64, error) { + // Do not aggregate the result in PostgreSQL. A jsonb array has a hard limit + // of 256 MiB for its elements, which large LLM request/response histories can + // exceed even though every individual record is valid. Reading row JSON in + // order also avoids building a second copy of the full table in PostgreSQL. + rows, err := q.Query(`SELECT row_to_json(row_data)::text FROM (`+inner+`) row_data`, args...) + if err != nil { + 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 + for rows.Next() { + var raw []byte + if err := rows.Scan(&raw); err != nil { + return nil, 0, err + } + if count > 0 { + output.WriteByte(',') + } + output.Write(raw) + count++ + } + if err := rows.Err(); err != nil { + return nil, 0, err + } + output.WriteByte(']') + return json.RawMessage(output.Bytes()), count, nil +>>>>>>> Stashed changes +} + +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 { @@ -453,11 +547,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) { @@ -491,8 +602,23 @@ 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 { +<<<<<<< Updated upstream raw, err := queryArchiveRows(tx, query.query, query.args...) +======= + 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...) +>>>>>>> Stashed changes if err != nil { return nil, fmt.Errorf("snapshot %s: %w", query.name, err) } @@ -528,7 +654,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() } diff --git a/db/task_archives_restore.go b/db/task_archives_restore.go index d3f15b0..71e916a 100644 --- a/db/task_archives_restore.go +++ b/db/task_archives_restore.go @@ -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 diff --git a/db/task_archives_test.go b/db/task_archives_test.go index dbe282f..d2496a3 100644 --- a/db/task_archives_test.go +++ b/db/task_archives_test.go @@ -1,13 +1,128 @@ package db import ( + "bytes" "encoding/json" + "errors" "fmt" "strings" "testing" "time" ) +<<<<<<< Updated upstream +======= +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) + 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) + } + if count != 2 { + t.Fatalf("row count=%d, want 2", count) + } + var rows []struct { + ID int64 `json:"id"` + Body string `json:"body"` + } + if err := json.Unmarshal(raw, &rows); err != nil { + t.Fatal(err) + } + if len(rows) != 2 { + t.Fatalf("decoded row count=%d, want 2", len(rows)) + } + if rows[0].ID != 1 || rows[0].Body != payload || rows[1].ID != 2 || rows[1].Body != "quoted: \"value\"\nline" { + t.Fatal("streamed rows were reordered or truncated") + } +} + +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) + } + } +} + +>>>>>>> Stashed changes func TestTaskArchiveDatabaseRoundTrip(t *testing.T) { d, err := Open(testDSN(t)) if err != nil { @@ -79,6 +194,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) } @@ -90,10 +211,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) } @@ -138,7 +266,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) } @@ -170,6 +300,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 @@ -242,6 +380,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() diff --git a/server/task_archive_package.go b/server/task_archive_package.go index ce4494d..bc55617 100644 --- a/server/task_archive_package.go +++ b/server/task_archive_package.go @@ -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" diff --git a/server/task_archive_package_test.go b/server/task_archive_package_test.go index 74e86ad..9846673 100644 --- a/server/task_archive_package_test.go +++ b/server/task_archive_package_test.go @@ -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) diff --git a/server/task_archives.go b/server/task_archives.go index c953a1a..30cea69 100644 --- a/server/task_archives.go +++ b/server/task_archives.go @@ -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 }