feat: persist queued server commands
This commit is contained in:
@@ -0,0 +1,142 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/rvbox/rvbox/internal/domain"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrCommandConflict = errors.New("command UUID conflicts with immutable request")
|
||||
ErrClientNotFound = errors.New("target client does not exist")
|
||||
)
|
||||
|
||||
// QueueCommandInput is the already-validated immutable command payload. The
|
||||
// control boundary validates and serializes the ExecutionSpec before it enters
|
||||
// this storage layer; the store owns only durable admission and quota charging.
|
||||
type QueueCommandInput struct {
|
||||
IssueUUID domain.UUID
|
||||
ClientID string
|
||||
IssueTime time.Time
|
||||
ReceiptTime time.Time
|
||||
QueueExpiryTime *time.Time
|
||||
ImmutableSHA256 [32]byte
|
||||
ExecutionSpec []byte
|
||||
}
|
||||
|
||||
type QueueCommandResult struct {
|
||||
Duplicate bool
|
||||
Lifecycle uint32
|
||||
Revision uint64
|
||||
}
|
||||
|
||||
// QueueCommand atomically persists a queued command before it is visible to a
|
||||
// dispatcher. Repeating the same UUID/hash is idempotent; a different hash is
|
||||
// a hard conflict and cannot replace a previously admitted command.
|
||||
func (store *Store) QueueCommand(ctx context.Context, input QueueCommandInput) (QueueCommandResult, error) {
|
||||
if !validCommandInput(input) {
|
||||
return QueueCommandResult{}, errors.New("invalid queued command")
|
||||
}
|
||||
stored, err := compressCommandSpec(input.ExecutionSpec)
|
||||
if err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
charge, err := EstimateCharge(ChargeInput{EncodedBytes: uint64(len(stored)), SQLiteRows: 1, IndexEntries: 2})
|
||||
if err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
|
||||
store.writeMu.Lock()
|
||||
defer store.writeMu.Unlock()
|
||||
database, err := store.openDatabase()
|
||||
if err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
tx, err := database.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
var storedHash []byte
|
||||
var lifecycle uint32
|
||||
var revision uint64
|
||||
err = tx.QueryRowContext(ctx, `SELECT immutable_request_sha256, lifecycle, revision FROM commands WHERE issue_uuid = ?`, input.IssueUUID[:]).Scan(&storedHash, &lifecycle, &revision)
|
||||
if err == nil {
|
||||
if len(storedHash) != len(input.ImmutableSHA256) || subtle.ConstantTimeCompare(storedHash, input.ImmutableSHA256[:]) != 1 {
|
||||
return QueueCommandResult{}, ErrCommandConflict
|
||||
}
|
||||
return QueueCommandResult{Duplicate: true, Lifecycle: lifecycle, Revision: revision}, tx.Commit()
|
||||
}
|
||||
if !errors.Is(err, sql.ErrNoRows) {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
|
||||
var clientCharged, serverCharged uint64
|
||||
err = tx.QueryRowContext(ctx, `SELECT clients.charged_bytes, storage_counters.command_charged_bytes
|
||||
FROM clients JOIN storage_counters ON storage_counters.singleton = 1 WHERE clients.client_id = ?`, input.ClientID).Scan(&clientCharged, &serverCharged)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return QueueCommandResult{}, ErrClientNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
freeBytes, err := store.freeSpaceProbe.AvailableBytes(store.dataDir)
|
||||
if err != nil {
|
||||
return QueueCommandResult{}, fmt.Errorf("check filesystem capacity: %w", err)
|
||||
}
|
||||
decision, err := CheckReservation(store.quotaLimits, ReservationState{
|
||||
ClientTotalCharged: clientCharged, ServerTotalCharged: serverCharged,
|
||||
CloseoutRemaining: store.quotaLimits.CloseoutReserveBytes, FilesystemFreeBytes: freeBytes,
|
||||
}, ReservationRequest{ChargedBytes: charge, PhysicalBytes: uint64(len(stored))})
|
||||
if err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
var expiry any
|
||||
if input.QueueExpiryTime != nil {
|
||||
expiry = input.QueueExpiryTime.UTC().UnixNano()
|
||||
}
|
||||
_, err = tx.ExecContext(ctx, `INSERT INTO commands (
|
||||
issue_uuid, client_id, issue_time, server_receipt_time, queue_expiry_time,
|
||||
lifecycle, revision, retained_compressed_bytes, charged_bytes,
|
||||
closeout_remaining_bytes, immutable_request_sha256, execution_spec,
|
||||
execution_spec_raw_bytes, execution_spec_stored_bytes, execution_spec_compression
|
||||
) VALUES (?, ?, ?, ?, ?, 1, 1, ?, ?, ?, ?, ?, ?, ?, 2)`,
|
||||
input.IssueUUID[:], input.ClientID, input.IssueTime.UTC().UnixNano(), input.ReceiptTime.UTC().UnixNano(), expiry,
|
||||
len(stored), charge, decision.CloseoutRemaining, input.ImmutableSHA256[:], stored, len(input.ExecutionSpec), len(stored))
|
||||
if err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `UPDATE clients SET charged_bytes = ? WHERE client_id = ?`, decision.ClientTotalCharged, input.ClientID); err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `UPDATE storage_counters SET command_charged_bytes = ? WHERE singleton = 1`, decision.ServerTotalCharged); err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return QueueCommandResult{}, err
|
||||
}
|
||||
return QueueCommandResult{Lifecycle: 1, Revision: 1}, nil
|
||||
}
|
||||
|
||||
func validCommandInput(input QueueCommandInput) bool {
|
||||
if isZeroUUID([16]byte(input.IssueUUID)) || input.ClientID == "" || len(input.ClientID) > 128 || input.IssueTime.IsZero() || input.ReceiptTime.IsZero() || len(input.ExecutionSpec) == 0 || allZero(input.ImmutableSHA256[:]) {
|
||||
return false
|
||||
}
|
||||
return input.QueueExpiryTime == nil || input.QueueExpiryTime.After(input.IssueTime)
|
||||
}
|
||||
|
||||
func compressCommandSpec(spec []byte) ([]byte, error) {
|
||||
encoder, err := zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer encoder.Close()
|
||||
return encoder.EncodeAll(spec, nil), nil
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/rvbox/rvbox/internal/domain"
|
||||
)
|
||||
|
||||
func TestQueueCommandDurableIdempotency_HP_DISPATCH_01(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{1}, SessionID: [16]byte{2}, ConnectedAt: time.Date(2026, time.September, 6, 12, 0, 0, 0, time.UTC),
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
issue, err := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000091")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
immutable := sha256.Sum256([]byte("queue immutable request"))
|
||||
now := time.Date(2026, time.September, 6, 12, 0, 1, 0, time.UTC)
|
||||
input := QueueCommandInput{IssueUUID: issue, ClientID: "win-client", IssueTime: now, ReceiptTime: now, ImmutableSHA256: immutable, ExecutionSpec: []byte("serialized execution spec")}
|
||||
queued, err := opened.QueueCommand(ctx, input)
|
||||
if err != nil || queued.Duplicate || queued.Lifecycle != 1 || queued.Revision != 1 {
|
||||
t.Fatalf("initial queue = %#v, %v", queued, err)
|
||||
}
|
||||
duplicate, err := opened.QueueCommand(ctx, input)
|
||||
if err != nil || !duplicate.Duplicate || duplicate.Lifecycle != 1 || duplicate.Revision != 1 {
|
||||
t.Fatalf("idempotent queue = %#v, %v", duplicate, err)
|
||||
}
|
||||
var stored []byte
|
||||
var compression, rawBytes, storedBytes uint64
|
||||
if err := opened.DB().QueryRow(`SELECT execution_spec, execution_spec_compression, execution_spec_raw_bytes, execution_spec_stored_bytes FROM commands WHERE issue_uuid = ?`, issue[:]).Scan(&stored, &compression, &rawBytes, &storedBytes); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
decoder, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
decompressed, err := decoder.DecodeAll(stored, nil)
|
||||
decoder.Close()
|
||||
if err != nil || !bytes.Equal(decompressed, input.ExecutionSpec) || compression != 2 || rawBytes != uint64(len(input.ExecutionSpec)) || storedBytes != uint64(len(stored)) {
|
||||
t.Fatalf("stored command = raw=%q compression=%d sizes=%d/%d err=%v", decompressed, compression, rawBytes, storedBytes, err)
|
||||
}
|
||||
different := input
|
||||
different.ImmutableSHA256 = sha256.Sum256([]byte("different request"))
|
||||
if _, err := opened.QueueCommand(ctx, different); !errors.Is(err, ErrCommandConflict) {
|
||||
t.Fatalf("conflicting UUID error = %v", err)
|
||||
}
|
||||
different = input
|
||||
different.IssueUUID, err = domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000092")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
different.ClientID = "missing-client"
|
||||
if _, err := opened.QueueCommand(ctx, different); !errors.Is(err, ErrClientNotFound) {
|
||||
t.Fatalf("missing client error = %v", err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user