From 749032396e7b5cb152d9ec4b1920cc8b98436dc2 Mon Sep 17 00:00:00 2001 From: Ada Lundhe Date: Fri, 23 Oct 2026 12:31:02 -0700 Subject: [PATCH] fix(session): add durable recovery primitives Signed-off-by: Ada Lundhe --- go/core/internal/database/session_recovery.go | 197 ++++++++++++++ .../database/session_recovery_test.go | 255 ++++++++++++++++++ .../internal/database/task_finalization.go | 96 +++---- go/core/internal/substrate/client.go | 10 + go/core/internal/substrate/snapshot.go | 11 + go/core/internal/substrate/snapshot_test.go | 20 ++ 6 files changed, 544 insertions(+), 45 deletions(-) create mode 100644 go/core/internal/database/session_recovery.go create mode 100644 go/core/internal/database/session_recovery_test.go diff --git a/go/core/internal/database/session_recovery.go b/go/core/internal/database/session_recovery.go new file mode 100644 index 0000000000..a803907ab9 --- /dev/null +++ b/go/core/internal/database/session_recovery.go @@ -0,0 +1,197 @@ +package database + +import ( + "context" + "fmt" + "time" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "google.golang.org/protobuf/types/known/timestamppb" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +// SessionRuntimeCrashed identifies a recovery admitted after Substrate confirmed +// that the recorded Actor crashed. Transport failures cannot admit this work. +const SessionRuntimeCrashed = "RuntimeCrashed" + +// ListSessionRecoveryCandidates finds recoveries admitted by an explicit request. +// Retrying durable work cannot admit recovery for a READY Session. +func (c *Client) ListSessionRecoveryCandidates(ctx context.Context, afterID string, limit int) ([]*apiv1.Session, error) { + rows, err := queryMany(ctx, c.db, ` + SELECT id, user_id, prepared_revision, state, data, operation, context_id, + source_checkpoint_id, history_id, operation_id, executor_id + FROM session_record i + WHERE ($1::text = '' OR i.id > $1::uuid) + AND i.state = 'RUNTIME_STATE_FAILED' AND i.operation = 'RUNTIME_OPERATION_RESUME' + ORDER BY i.id LIMIT $2 + `, pgx.RowToStructByName[sessionRow], afterID, limit) + if err != nil { + return nil, err + } + result := make([]*apiv1.Session, 0, len(rows)) + for _, row := range rows { + session, err := toSession(row) + if err != nil { + return nil, err + } + result = append(result, session) + } + return result, nil +} + +// SessionLifecycleGeneration returns the ID of the Session's latest lifecycle +// operation, or uuid.Nil when none is recorded. Every operation replaces it, and +// completion keeps it. Recovery reads it before observing Substrate. +func (c *Client) SessionLifecycleGeneration(ctx context.Context, id string) (uuid.UUID, error) { + generation, err := queryOne(ctx, c.db, `SELECT operation_id FROM session_record WHERE id = $1`, pgx.RowTo[*uuid.UUID], id) + if err != nil { + return uuid.Nil, notFoundOr(err) + } + if generation == nil { + return uuid.Nil, nil + } + return *generation, nil +} + +// BeginCrashedSessionRecovery fences confirmed-crash recovery under the Session lock. +// It preserves durable results; obsolete Actors and conflicting operations fail. +// generation is the caller's SessionLifecycleGeneration read from before it +// observed the crash. A lifecycle operation begun since then makes the observation +// stale, because a concurrent revert keeps the Actor UID and returns to READY. +func (c *Client) BeginCrashedSessionRecovery(ctx context.Context, id string, generation uuid.UUID, actorUID string, snapshot SessionTaskSnapshot) (*SessionOperation, error) { + var operation *SessionOperation + err := c.withTx(ctx, func(tx pgx.Tx) error { + row, err := lockSession(ctx, tx, id) + if err != nil { + return notFoundOr(err) + } + matches, err := queryOne(ctx, tx, `SELECT actor_uid = $2 FROM session WHERE id = $1`, pgx.RowTo[bool], row.ID, actorUID) + if err != nil { + return err + } + if !matches || actorUID == "" { + return ErrConflict + } + session, err := toSession(row) + if err != nil { + return err + } + if session.State == apiv1.RuntimeState_RUNTIME_STATE_FAILED && session.Operation == apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME && session.Failure.GetReason() == SessionRuntimeCrashed { + operation, err = toSessionOperation(row) + return err + } + if session.State != apiv1.RuntimeState_RUNTIME_STATE_READY || session.Operation != apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE { + return ErrConflict + } + current := uuid.Nil + if row.OperationID != nil { + current = *row.OperationID + } + if current != generation { + return ErrConflict + } + // An issued idle suspension does not exclude recovery here. It cannot finish + // against the crashed Actor, and its worker keeps the claim when suspension + // fails, so waiting for it would block recovery permanently. Settlement below + // retires the claim in this transaction; a late finish for it changes nothing. + // The generation check keeps the crash observation current: without another + // lifecycle operation, nothing returns the crashed Actor to service. + busy, err := queryOne(ctx, tx, ` + SELECT EXISTS (SELECT 1 FROM session_checkpoint WHERE source_session_id = $1 AND state = 'CREATING') + `, pgx.RowTo[bool], row.ID) + if err != nil { + return err + } + if busy { + return ErrConflict + } + row.State = apiv1.RuntimeState_RUNTIME_STATE_FAILED.String() + row.runtimeInstanceRow, err = beginRuntimeOperation(ctx, tx, row.runtimeInstanceRow, runtimeKindAgent, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME) + if err != nil { + return err + } + session.State, session.Operation = apiv1.RuntimeState_RUNTIME_STATE_FAILED, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME + session.Failure = &apiv1.Failure{Reason: SessionRuntimeCrashed, Message: "Runtime crashed; recovery to its last snapshot is pending"} + session.UpdatedAt = timestamppb.Now() + row.Data, err = marshalSession(session) + if err != nil { + return err + } + if err := execSQL(ctx, tx, `UPDATE session SET data = $2, dispatch_id = NULL, dispatch_expires_at = NULL WHERE id = $1`, row.ID, row.Data); err != nil { + return err + } + if err := settleCrashedSessionTask(ctx, tx, row, snapshot); err != nil { + return err + } + operation, err = toSessionOperation(row) + return err + }) + return operation, err +} + +// settleCrashedSessionTask retains captured waits and fails discarded nonterminal work. +// Callers hold the Session lock; recovery cannot checkpoint discarded state. +func settleCrashedSessionTask(ctx context.Context, tx pgx.Tx, session sessionRow, snapshot SessionTaskSnapshot) error { + if err := execSQL(ctx, tx, `UPDATE session_task_event SET quiescence_pending = FALSE WHERE history_id = $1 AND quiescence_pending`, session.HistoryID); err != nil { + return err + } + taskIDs, err := queryMany(ctx, tx, ` + SELECT t.id FROM session_task t WHERE t.history_id = $1 + AND (t.state IN ('TASK_STATE_SUBMITTED', 'TASK_STATE_WORKING') + OR (t.state IN ('TASK_STATE_INPUT_REQUIRED', 'TASK_STATE_AUTH_REQUIRED') + AND NOT EXISTS (SELECT 1 FROM session_task_event captured + WHERE captured.history_id = t.history_id AND $2::text <> '' + AND captured.snapshot_uri = $2 AND captured.snapshot_fidelity = $3 + AND captured.snapshot_atespace = $4 + AND captured.sequence >= (SELECT MAX(sequence) FROM session_task_event latest + WHERE latest.history_id = t.history_id AND latest.task_id = t.id)))) + `, pgx.RowTo[string], session.HistoryID, snapshot.URI, snapshot.Fidelity, snapshot.Atespace) + if err != nil { + return fmt.Errorf("read interrupted Session tasks: %w", err) + } + for _, taskID := range taskIDs { + version, err := taskVersion(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + pending, err := queryOne(ctx, tx, ` + SELECT EXISTS (SELECT 1 FROM session_task_event + WHERE sequence = $1 AND NOT published AND quiescence_pending IS NOT NULL) + `, pgx.RowTo[bool], version) + if err != nil { + return err + } + if pending { + if err := settleSessionTask(ctx, tx, session, taskID, version); err != nil { + return err + } + } + stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + task, err := unmarshalSessionTask(stored.Data) + if err != nil { + return err + } + if !task.Status.State.Terminal() { + task.Status = a2a.TaskStatus{State: a2a.TaskStateFailed, Timestamp: new(time.Now().UTC()), + Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("The runtime crashed before this turn completed. Send a new message to continue."))} + event := &a2a.TaskStatusUpdateEvent{TaskID: task.ID, ContextID: task.ContextID, Status: task.Status} + if err := storeSessionTaskEvent(ctx, tx, session, task, event, false); err != nil { + return err + } + version, err = taskVersion(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + if err := publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task); err != nil { + return err + } + } + } + return nil +} diff --git a/go/core/internal/database/session_recovery_test.go b/go/core/internal/database/session_recovery_test.go new file mode 100644 index 0000000000..68827b938d --- /dev/null +++ b/go/core/internal/database/session_recovery_test.go @@ -0,0 +1,255 @@ +package database + +import ( + "testing" + + "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + apiv1 "github.com/kagent-dev/kagent/go/api/gen/kagent/api/v1" +) + +func TestCrashedRecoverySettlesTasksAtomically(t *testing.T) { + for _, test := range []struct { + name string + state, want a2a.TaskState + published, captured bool + }{ + {name: "submitted", state: a2a.TaskStateSubmitted, want: a2a.TaskStateFailed}, + {name: "working", state: a2a.TaskStateWorking, want: a2a.TaskStateFailed}, + {name: "completed before publication", state: a2a.TaskStateCompleted, want: a2a.TaskStateCompleted}, + {name: "unpublished input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateFailed}, + {name: "unpublished auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateFailed}, + {name: "published input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateFailed, published: true}, + {name: "published auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateFailed, published: true}, + {name: "captured input wait", state: a2a.TaskStateInputRequired, want: a2a.TaskStateInputRequired, published: true, captured: true}, + {name: "captured auth wait", state: a2a.TaskStateAuthRequired, want: a2a.TaskStateAuthRequired, published: true, captured: true}, + } { + t.Run(test.name, func(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, _, err := client.CreateSession(t.Context(), newSessionRequest(uuid.NewString(), "assistant", "kagent", ""), uuid.NewString()) + require.NoError(t, err) + session, err = markSessionReady(t.Context(), client, session.Id, "runtime.example") + require.NoError(t, err) + task := newSessionTask(uuid.NewString(), "original input") + task.ContextID, task.Status.State = session.ContextId, test.state + version, err := client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("create"), task, "") + require.NoError(t, err) + if test.published { + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + } + snapshot := SessionTaskSnapshot{} + if test.captured { + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + snapshot = SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/captured", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + require.NoError(t, client.FinishSessionQuiescence(t.Context(), work, &snapshot)) + } + before, _, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + page, err := client.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Empty(t, page) + + operation, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, snapshot) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_FAILED, operation.Instance.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, operation.Instance.Operation) + require.Equal(t, SessionRuntimeCrashed, operation.Instance.Failure.Reason) + stored, err := client.GetSettledSessionTask(t.Context(), session.Id, string(task.ID), nil) + require.NoError(t, err) + require.Equal(t, test.want, stored.Status.State) + require.Equal(t, before.History, stored.History) + require.Equal(t, before.Artifacts, stored.Artifacts) + if test.want == test.state { + require.Equal(t, before.Status, stored.Status) + } + _, settledVersion, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + replayed, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, snapshot) + require.NoError(t, err) + require.Equal(t, operation.ID, replayed.ID) + _, replayedVersion, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, settledVersion, replayedVersion) + _, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash("late"), task, task, "") + if test.state.Terminal() { + require.ErrorIs(t, err, ErrFailedPrecondition) + } else { + require.ErrorIs(t, err, ErrConflict) + } + require.ErrorIs(t, client.ReserveSessionDispatch(t.Context(), session.Id, uuid.New(), "next"), ErrConflict) + _, _, err = client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.ErrorIs(t, err, ErrConflict) + page, err = client.ListSessionRecoveryCandidates(t.Context(), "", 100) + require.NoError(t, err) + require.Len(t, page, 1) + require.Equal(t, session.Id, page[0].Id) + }) + } +} + +func TestCrashedRecoveryAdmissionRespectsFences(t *testing.T) { + for _, test := range []struct { + name string + actorUID string + checkpoint, superseded, lifecycle bool + }{ + {name: "replacement actor", actorUID: "replacement"}, + {name: "checkpoint", checkpoint: true}, + {name: "revert after the crash was observed", superseded: true}, + {name: "lifecycle", lifecycle: true}, + } { + t.Run(test.name, func(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + uid := "actor-" + session.Id + // Callers read the generation before they observe the crash. + generation := lifecycleGeneration(t, client, session.Id) + switch { + case test.actorUID != "": + uid = test.actorUID + case test.checkpoint: + task.Status.State = a2a.TaskStateCompleted + require.NoError(t, saveRuntimeTask(t, client, session.Id, task, task, &SessionTaskSnapshot{ + Atespace: "team-a", URI: "s3://snapshots/completed", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES", + })) + _, _, err := client.ReserveSessionCheckpoint(t.Context(), &apiv1.Checkpoint{Id: uuid.NewString(), SessionId: session.Id, HeadTaskId: string(task.ID)}, session.Creator, uuid.NewString()) + require.NoError(t, err) + case test.superseded: + // Another revert finishes on the same Actor UID, and a new turn + // starts once ingress wakes the restored Actor. + completeCrashedRecovery(t, client, session.Id, uid) + task = newSessionTask(uuid.NewString(), "next input") + task.ContextID, task.Status.State = session.ContextId, a2a.TaskStateWorking + _, err := client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("next"), task, "") + require.NoError(t, err) + case test.lifecycle: + _, err := client.BeginSessionOperation(t.Context(), session.Id, apiv1.RuntimeOperation_RUNTIME_OPERATION_SUSPEND) + require.NoError(t, err) + } + before, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + _, err = client.BeginCrashedSessionRecovery(t.Context(), session.Id, generation, uid, SessionTaskSnapshot{}) + require.ErrorIs(t, err, ErrConflict) + after, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.True(t, proto.Equal(before, after)) + if test.superseded { + running, _, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + require.Equal(t, a2a.TaskStateWorking, running.Status.State, "a stale observation must not fail the new turn") + } + }) + } +} + +// A crash during post-turn suspension leaves the worker's claim issued, because +// the worker cannot tell whether its suspension took effect. The confirmed crash +// settles that: admission retires the claim instead of waiting for it. +func TestCrashedRecoveryRetiresIssuedIdleSuspension(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + uid := "actor-" + session.Id + + operation, err := client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), uid, SessionTaskSnapshot{}) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_RESUME, operation.Instance.Operation) + pending, err := queryOne(t.Context(), client.db, `SELECT quiescence_pending FROM session_task_event WHERE sequence = $1`, pgx.RowTo[bool], work.Version) + require.NoError(t, err) + require.False(t, pending, "admission retires the claim in its own transaction") + finishCrashedRecovery(t, client, session.Id, operation, uid) + + late := SessionTaskSnapshot{Atespace: "team-a", URI: "s3://snapshots/late", Fidelity: "SNAPSHOT_FIDELITY_VOLUMES"} + require.NoError(t, client.FinishSessionQuiescence(t.Context(), work, &late)) + row, err := readSession(t.Context(), client.db, session.Id) + require.NoError(t, err) + recorded, err := queryOne(t.Context(), client.db, ` + SELECT EXISTS (SELECT 1 FROM session_task WHERE history_id = $1 AND snapshot_uri = $2) + OR EXISTS (SELECT 1 FROM session_task_event WHERE history_id = $1 AND snapshot_uri = $2) + `, pgx.RowTo[bool], row.HistoryID, late.URI) + require.NoError(t, err) + require.False(t, recorded, "a late finish of the retired claim records nothing") + _, err = client.ClaimSessionQuiescence(t.Context()) + require.ErrorIs(t, err, ErrNotFound) + next := newSessionTask(uuid.NewString(), "next input") + next.ContextID, next.Status.State = session.ContextId, a2a.TaskStateWorking + _, err = client.CreateRuntimeTask(t.Context(), session.Id, taskMutationHash("next"), next, "") + require.NoError(t, err, "the retired claim no longer blocks a new turn") +} + +func lifecycleGeneration(t *testing.T, client *Client, sessionID string) uuid.UUID { + t.Helper() + generation, err := client.SessionLifecycleGeneration(t.Context(), sessionID) + require.NoError(t, err) + return generation +} + +// completeCrashedRecovery admits and finishes a revert as a successful concurrent +// caller would, leaving the Session READY on the same Actor UID. +func completeCrashedRecovery(t *testing.T, client *Client, sessionID, actorUID string) { + t.Helper() + operation, err := client.BeginCrashedSessionRecovery(t.Context(), sessionID, lifecycleGeneration(t, client, sessionID), actorUID, SessionTaskSnapshot{}) + require.NoError(t, err) + finishCrashedRecovery(t, client, sessionID, operation, actorUID) +} + +func finishCrashedRecovery(t *testing.T, client *Client, sessionID string, operation *SessionOperation, actorUID string) { + t.Helper() + executor := uuid.New() + claimed, err := client.ClaimSessionOperation(t.Context(), sessionID, operation.ID, executor) + require.NoError(t, err) + require.True(t, claimed) + recovered, err := client.FinishSessionOperation(t.Context(), sessionID, operation.ID, executor, "", actorUID, "") + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, recovered.State) +} + +func TestCrashedRecoveryAdmissionRollsBackOnMalformedTask(t *testing.T) { + client := NewClient(setupTestDB(t), "public") + sessionFixture(t, client, t.Context(), "team-a", "revision", "assistant", "kagent") + session, task := waitingTaskFixture(t, client) + _, version, err := client.GetVersionedSessionTask(t.Context(), session.Id, string(task.ID)) + require.NoError(t, err) + for _, state := range []a2a.TaskState{a2a.TaskStateWorking, a2a.TaskStateInputRequired} { + task.Status.State = state + version, err = client.UpdateSessionTask(t.Context(), session.Id, version, taskMutationHash(string(state)), task, task, "") + require.NoError(t, err) + } + require.NoError(t, client.SettleSessionTask(t.Context(), session.Id, string(task.ID), version)) + row, err := readSession(t.Context(), client.db, session.Id) + require.NoError(t, err) + before, err := readSessionTask(t.Context(), client.db, row.HistoryID, string(task.ID)) + require.NoError(t, err) + // A corrupt projection makes settlement fail after the lifecycle write. + _, err = client.db.Exec(t.Context(), `UPDATE session_task SET data = $2 WHERE history_id = $1 AND id = $3`, row.HistoryID, []byte{0xff}, string(task.ID)) + require.NoError(t, err) + _, err = client.BeginCrashedSessionRecovery(t.Context(), session.Id, lifecycleGeneration(t, client, session.Id), "actor-"+session.Id, SessionTaskSnapshot{}) + require.Error(t, err) + current, err := client.GetSessionByID(t.Context(), session.Id) + require.NoError(t, err) + require.Equal(t, apiv1.RuntimeState_RUNTIME_STATE_READY, current.State) + require.Equal(t, apiv1.RuntimeOperation_RUNTIME_OPERATION_NONE, current.Operation) + require.Nil(t, current.Failure) + _, err = client.db.Exec(t.Context(), `UPDATE session_task SET data = $2 WHERE history_id = $1 AND id = $3`, row.HistoryID, before.Data, string(task.ID)) + require.NoError(t, err) + work, err := client.ClaimSessionQuiescence(t.Context()) + require.NoError(t, err) + require.Equal(t, string(task.ID), work.TaskID) +} diff --git a/go/core/internal/database/task_finalization.go b/go/core/internal/database/task_finalization.go index 9907257b21..ac209c8dfe 100644 --- a/go/core/internal/database/task_finalization.go +++ b/go/core/internal/database/task_finalization.go @@ -23,54 +23,60 @@ func (c *Client) SettleSessionTask(ctx context.Context, sessionID, taskID string if session.State == "RUNTIME_STATE_DELETED" { return ErrNotFound } - type boundary struct { - Data []byte - ExpectedVersion int64 - Published bool - } - row, err := queryOne(ctx, tx, ` - SELECT data, expected_version, published FROM session_task_event - WHERE history_id = $1 AND task_id = $2 AND sequence = $3 - AND expected_version IS NOT NULL AND quiescence_pending IS NOT NULL - `, pgx.RowToStructByName[boundary], session.HistoryID, taskID, version) - if errors.Is(err, pgx.ErrNoRows) { - return ErrConflict - } - if err != nil { - return err - } - if row.Published { - return nil - } - stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) - if err != nil { - return err - } - wire, err := applyStoredBoundary(stored.Data, row.Data) - if err != nil { - return err - } - task, err := pbconv.FromProtoTask(wire) - if err != nil { - return err - } - data, err := proto.Marshal(wire) - if err != nil { - return err - } - if _, err := saveTaskProjection(ctx, tx, session.HistoryID, taskID, string(task.Status.State), task.Status.Timestamp, data); err != nil { - return err - } - if err := execSQL(ctx, tx, ` - UPDATE session_task_event SET published = TRUE - WHERE history_id = $1 AND task_id = $2 AND sequence > $3 AND sequence <= $4 - `, session.HistoryID, taskID, row.ExpectedVersion, version); err != nil { - return err - } - return publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task) + return settleSessionTask(ctx, tx, session, taskID, version) }) } +// settleSessionTask publishes an already saved boundary under the Session row lock. +// Recovery uses the same publication after confirming the runtime crashed. +func settleSessionTask(ctx context.Context, tx pgx.Tx, session sessionRow, taskID string, version int64) error { + type boundary struct { + Data []byte + ExpectedVersion int64 + Published bool + } + row, err := queryOne(ctx, tx, ` + SELECT data, expected_version, published FROM session_task_event + WHERE history_id = $1 AND task_id = $2 AND sequence = $3 + AND expected_version IS NOT NULL AND quiescence_pending IS NOT NULL + `, pgx.RowToStructByName[boundary], session.HistoryID, taskID, version) + if errors.Is(err, pgx.ErrNoRows) { + return ErrConflict + } + if err != nil { + return err + } + if row.Published { + return nil + } + stored, err := readSessionTask(ctx, tx, session.HistoryID, taskID) + if err != nil { + return err + } + wire, err := applyStoredBoundary(stored.Data, row.Data) + if err != nil { + return err + } + task, err := pbconv.FromProtoTask(wire) + if err != nil { + return err + } + data, err := proto.Marshal(wire) + if err != nil { + return err + } + if _, err := saveTaskProjection(ctx, tx, session.HistoryID, taskID, string(task.Status.State), task.Status.Timestamp, data); err != nil { + return err + } + if err := execSQL(ctx, tx, ` + UPDATE session_task_event SET published = TRUE + WHERE history_id = $1 AND task_id = $2 AND sequence > $3 AND sequence <= $4 + `, session.HistoryID, taskID, row.ExpectedVersion, version); err != nil { + return err + } + return publishPushBoundary(ctx, tx, session.HistoryID, taskID, version, task) +} + func applyStoredBoundary(data, eventData []byte) (*a2apb.Task, error) { task, event := &a2apb.Task{}, &a2apb.StreamResponse{} if err := proto.Unmarshal(data, task); err != nil { diff --git a/go/core/internal/substrate/client.go b/go/core/internal/substrate/client.go index 4494b87c6e..798e6f72f1 100644 --- a/go/core/internal/substrate/client.go +++ b/go/core/internal/substrate/client.go @@ -200,6 +200,16 @@ func (c *Client) ResumeActor(ctx context.Context, atespace, actorID string) (*at return resp.GetActor(), nil } +func (c *Client) RevertActor(ctx context.Context, atespace, actorID string) (*ateapipb.Actor, error) { + ctx, cancel := c.callCtx(ctx) + defer cancel() + resp, err := c.ControlClient.RevertActor(ctx, &ateapipb.RevertActorRequest{Actor: actorRef(atespace, actorID)}) + if err != nil { + return nil, fmt.Errorf("failed to revert Actor %s/%s: %w", atespace, actorID, err) + } + return resp.GetActor(), nil +} + // SuspendActor waits for durable storage, including an upload left by an earlier // attempt. Substrate's synchronous mode returns already-suspended Actors as-is. func (c *Client) SuspendActor(ctx context.Context, atespace, actorID string) (*ateapipb.Actor, error) { diff --git a/go/core/internal/substrate/snapshot.go b/go/core/internal/substrate/snapshot.go index 02434ab83d..de2597052b 100644 --- a/go/core/internal/substrate/snapshot.go +++ b/go/core/internal/substrate/snapshot.go @@ -22,3 +22,14 @@ func SnapshotURI(snapshot *ateapipb.Snapshot) string { } return storage.GetObject().GetSnapshotUri() } + +// Checkpoints require the newest generation; revert retains the newest durable copy. +func LatestDurableActorSnapshot(actor *ateapipb.Actor) *ateapipb.Snapshot { + var latest *ateapipb.Snapshot + for _, snapshot := range actor.GetStatus().GetSnapshots() { + if SnapshotURI(snapshot) != "" && snapshot.GetGeneration() > latest.GetGeneration() { + latest = snapshot + } + } + return latest +} diff --git a/go/core/internal/substrate/snapshot_test.go b/go/core/internal/substrate/snapshot_test.go index 0d29f25f65..647e4aa903 100644 --- a/go/core/internal/substrate/snapshot_test.go +++ b/go/core/internal/substrate/snapshot_test.go @@ -40,3 +40,23 @@ func TestActorSnapshotGenerationAndDurability(t *testing.T) { require.Nil(t, LatestActorSnapshot(nil)) require.Empty(t, SnapshotURI(nil)) } + +func TestRevertSnapshotSkipsIncompleteGenerations(t *testing.T) { + durable := &ateapipb.Snapshot{Generation: 2, DurableSnapshot: &ateapipb.SnapshotStorage{ + Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED, + Object: &ateapipb.ObjectSnapshot{SnapshotUri: "s3://snapshots/durable"}, + }} + for _, latest := range []*ateapipb.Snapshot{ + {Generation: 3, Locality: "node-1"}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING}}, + {Generation: 3, DurableSnapshot: &ateapipb.SnapshotStorage{Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED}}, + } { + actor := &ateapipb.Actor{Status: &ateapipb.ActorStatus{Snapshots: []*ateapipb.Snapshot{latest, durable}}} + require.Same(t, durable, LatestDurableActorSnapshot(actor)) + require.Same(t, latest, LatestActorSnapshot(actor)) + } + require.Nil(t, LatestDurableActorSnapshot(nil)) + require.Nil(t, LatestDurableActorSnapshot(&ateapipb.Actor{Status: &ateapipb.ActorStatus{ + Snapshots: []*ateapipb.Snapshot{{Generation: 1, Locality: "node-1"}}, + }})) +}