feat: dispatch queued commands after reconciliation
This commit is contained in:
@@ -71,3 +71,56 @@ func TestQueueCommandDurableIdempotency_HP_DISPATCH_01(t *testing.T) {
|
||||
t.Fatalf("missing client error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimDispatchExpiresAndFencesRequeue_HP_DISPATCH_02(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
opened, err := Open(ctx, Options{DataDir: filepath.Join(t.TempDir(), "state"), BusyTimeout: time.Second})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = opened.Close() })
|
||||
if _, err := opened.RegisterClientSession(ctx, ClientRegistration{
|
||||
ClientID: "win-client", Platform: 2, Architecture: "amd64", DaemonVersion: "test", DaemonCWD: `C:\`, SupportedShells: []byte{1},
|
||||
ClientInstanceID: [16]byte{3}, SessionID: [16]byte{4}, ConnectedAt: time.Date(2026, time.September, 6, 12, 0, 0, 0, time.UTC),
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Date(2026, time.September, 6, 12, 0, 1, 0, time.UTC)
|
||||
expiredIssue, err := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000093")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
queuedIssue, err := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000094")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, input := range []QueueCommandInput{
|
||||
{IssueUUID: expiredIssue, ClientID: "win-client", IssueTime: now.Add(-time.Minute), ReceiptTime: now.Add(-time.Minute), QueueExpiryTime: timePtr(now.Add(-time.Second)), ImmutableSHA256: sha256.Sum256([]byte("expired")), ExecutionSpec: []byte("expired spec")},
|
||||
{IssueUUID: queuedIssue, ClientID: "win-client", IssueTime: now, ReceiptTime: now, ImmutableSHA256: sha256.Sum256([]byte("queued")), ExecutionSpec: []byte("queued spec")},
|
||||
} {
|
||||
if _, err := opened.QueueCommand(ctx, input); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
candidate, err := opened.ClaimNextDispatch(ctx, "win-client", 7, now)
|
||||
if err != nil || candidate == nil || candidate.IssueUUID != queuedIssue || candidate.Revision != 1 || string(candidate.ExecutionSpec) != "queued spec" {
|
||||
t.Fatalf("dispatch candidate = %#v, %v", candidate, err)
|
||||
}
|
||||
if requeued, err := opened.RequeueDispatch(ctx, queuedIssue, "win-client", 6); err != nil || requeued {
|
||||
t.Fatalf("stale requeue = %t, %v", requeued, err)
|
||||
}
|
||||
if requeued, err := opened.RequeueDispatch(ctx, queuedIssue, "win-client", 7); err != nil || !requeued {
|
||||
t.Fatalf("owned requeue = %t, %v", requeued, err)
|
||||
}
|
||||
candidate, err = opened.ClaimNextDispatch(ctx, "win-client", 8, now)
|
||||
if err != nil || candidate == nil || candidate.IssueUUID != queuedIssue {
|
||||
t.Fatalf("reclaimed candidate = %#v, %v", candidate, err)
|
||||
}
|
||||
var lifecycle, revision uint64
|
||||
if err := opened.DB().QueryRow(`SELECT lifecycle, revision FROM commands WHERE issue_uuid = ?`, expiredIssue[:]).Scan(&lifecycle, &revision); err != nil || lifecycle != 10 || revision != 2 {
|
||||
t.Fatalf("expired command lifecycle/revision = %d/%d, %v", lifecycle, revision, err)
|
||||
}
|
||||
}
|
||||
|
||||
func timePtr(value time.Time) *time.Time { return &value }
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"time"
|
||||
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/rvbox/rvbox/internal/domain"
|
||||
)
|
||||
|
||||
const maxStoredExecutionSpecBytes = 1 << 20
|
||||
|
||||
// DispatchCandidate is the immutable work reconstructed after an atomic
|
||||
// queued-to-dispatched transition. The session writer must use the supplied
|
||||
// generation and call RequeueDispatch if its network write never starts.
|
||||
type DispatchCandidate struct {
|
||||
IssueUUID domain.UUID
|
||||
Revision uint64
|
||||
IssueTime time.Time
|
||||
QueueExpiryTime *time.Time
|
||||
ExecutionSpec []byte
|
||||
}
|
||||
|
||||
// ClaimNextDispatch expires old queued rows and atomically assigns the oldest
|
||||
// remaining row to one live client-session generation. It deliberately records
|
||||
// the dispatch before the wire write so uncertain writes are reconciled rather
|
||||
// than accidentally delivered twice.
|
||||
func (store *Store) ClaimNextDispatch(ctx context.Context, clientID string, generation uint64, now time.Time) (*DispatchCandidate, error) {
|
||||
if clientID == "" || generation == 0 || now.IsZero() {
|
||||
return nil, errors.New("invalid dispatch claim")
|
||||
}
|
||||
store.writeMu.Lock()
|
||||
defer store.writeMu.Unlock()
|
||||
database, err := store.openDatabase()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tx, err := database.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
// Expiry is terminal. The revision advance distinguishes it from a delayed
|
||||
// queued snapshot or a late dispatch attempt.
|
||||
if _, err := tx.ExecContext(ctx, `UPDATE commands SET lifecycle = 10, terminal_time = ?, revision = revision + 1
|
||||
WHERE client_id = ? AND lifecycle = 1 AND queue_expiry_time IS NOT NULL AND queue_expiry_time <= ?`, now.UTC().UnixNano(), clientID, now.UTC().UnixNano()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var encodedIssue, stored []byte
|
||||
var revision uint64
|
||||
var issueTime int64
|
||||
var expiry sql.NullInt64
|
||||
var rawBytes uint64
|
||||
err = tx.QueryRowContext(ctx, `SELECT issue_uuid, revision, issue_time, queue_expiry_time, execution_spec, execution_spec_raw_bytes
|
||||
FROM commands WHERE client_id = ? AND lifecycle = 1 ORDER BY issue_time, issue_uuid LIMIT 1`, clientID).Scan(&encodedIssue, &revision, &issueTime, &expiry, &stored, &rawBytes)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
if err := tx.Commit(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(encodedIssue) != 16 {
|
||||
return nil, ErrInvalidSegmentRecord
|
||||
}
|
||||
var issue domain.UUID
|
||||
copy(issue[:], encodedIssue)
|
||||
spec, err := decompressCommandSpec(stored, rawBytes)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result, err := tx.ExecContext(ctx, `UPDATE commands SET lifecycle = 2, target_session_generation = ? WHERE issue_uuid = ? AND client_id = ? AND lifecycle = 1`, generation, issue[:], clientID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
changed, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if changed != 1 {
|
||||
return nil, errors.New("queued command changed during dispatch claim")
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
candidate := &DispatchCandidate{IssueUUID: issue, Revision: revision, IssueTime: time.Unix(0, issueTime).UTC(), ExecutionSpec: spec}
|
||||
if expiry.Valid {
|
||||
value := time.Unix(0, expiry.Int64).UTC()
|
||||
candidate.QueueExpiryTime = &value
|
||||
}
|
||||
return candidate, nil
|
||||
}
|
||||
|
||||
// RequeueDispatch reverses only a known pre-acceptance write failure. A later
|
||||
// session generation, acceptance, or reconciliation decision cannot be rolled
|
||||
// back by a stale writer.
|
||||
func (store *Store) RequeueDispatch(ctx context.Context, issueUUID domain.UUID, clientID string, generation uint64) (bool, error) {
|
||||
if isZeroUUID([16]byte(issueUUID)) || clientID == "" || generation == 0 {
|
||||
return false, errors.New("invalid dispatch requeue")
|
||||
}
|
||||
store.writeMu.Lock()
|
||||
defer store.writeMu.Unlock()
|
||||
database, err := store.openDatabase()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
result, err := database.ExecContext(ctx, `UPDATE commands SET lifecycle = 1, target_session_generation = NULL
|
||||
WHERE issue_uuid = ? AND client_id = ? AND lifecycle = 2 AND target_session_generation = ?`, issueUUID[:], clientID, generation)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
changed, err := result.RowsAffected()
|
||||
return changed == 1, err
|
||||
}
|
||||
|
||||
func decompressCommandSpec(stored []byte, rawBytes uint64) ([]byte, error) {
|
||||
if rawBytes == 0 || rawBytes > maxStoredExecutionSpecBytes || rawBytes > uint64(math.MaxInt) {
|
||||
return nil, errors.New("invalid stored execution spec size")
|
||||
}
|
||||
decoder, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(maxStoredExecutionSpecBytes+1))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer decoder.Close()
|
||||
decoded, err := decoder.DecodeAll(stored, nil)
|
||||
if err != nil || uint64(len(decoded)) != rawBytes || uint64(len(decoded)) > maxStoredExecutionSpecBytes {
|
||||
return nil, fmt.Errorf("invalid stored execution spec")
|
||||
}
|
||||
return bytes.Clone(decoded), nil
|
||||
}
|
||||
@@ -75,7 +75,7 @@ CREATE TABLE commands (
|
||||
client_id TEXT NOT NULL REFERENCES clients(client_id) ON DELETE RESTRICT,
|
||||
issue_time INTEGER NOT NULL, server_receipt_time INTEGER NOT NULL, queue_expiry_time INTEGER, terminal_time INTEGER,
|
||||
lifecycle INTEGER NOT NULL CHECK(lifecycle BETWEEN 1 AND 11),
|
||||
revision INTEGER NOT NULL CHECK(revision > 0), exit_code INTEGER,
|
||||
revision INTEGER NOT NULL CHECK(revision > 0), target_session_generation INTEGER CHECK(target_session_generation IS NULL OR target_session_generation > 0), exit_code INTEGER,
|
||||
retention_status INTEGER NOT NULL DEFAULT 1 CHECK(retention_status BETWEEN 1 AND 3),
|
||||
last_event_seq INTEGER NOT NULL DEFAULT 0 CHECK(last_event_seq >= 0),
|
||||
retained_compressed_bytes INTEGER NOT NULL DEFAULT 0 CHECK(retained_compressed_bytes >= 0),
|
||||
|
||||
Reference in New Issue
Block a user