feat: persist and dispatch client command input
This commit is contained in:
@@ -12,6 +12,7 @@ import (
|
||||
|
||||
rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
|
||||
"github.com/rvbox/rvbox/internal/domain"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -235,6 +236,18 @@ retained_compressed_bytes = retained_compressed_bytes + ?, output_charged_bytes
|
||||
}
|
||||
}
|
||||
}
|
||||
if err == nil && event.EventType == 7 {
|
||||
// Stdin acknowledgements are command events, but their durable delivery
|
||||
// cursor lives in stdin_writes. Mark the cumulative prefix in the same
|
||||
// transaction as the event so a crash cannot make an acknowledged write
|
||||
// replay forever or release it before its event is durable.
|
||||
var wire rvboxv1.CommandEvent
|
||||
if unmarshalErr := proto.Unmarshal(event.Payload, &wire); unmarshalErr != nil || wire.GetStdinAck() == nil || wire.GetStdinAck().GetWriteSeq() == 0 {
|
||||
err = ErrInvalidSegmentRecord
|
||||
} else {
|
||||
_, err = tx.ExecContext(ctx, `UPDATE stdin_writes SET acknowledged = 1 WHERE issue_uuid = ? AND write_seq <= ?`, event.IssueUUID[:], wire.GetStdinAck().GetWriteSeq())
|
||||
}
|
||||
}
|
||||
if err == nil {
|
||||
var update sql.Result
|
||||
update, err = tx.ExecContext(ctx, `UPDATE clients SET charged_bytes = ? WHERE client_id = ? AND charged_bytes = ?`, reservation.ClientTotalCharged, clientID, clientCharged)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
@@ -9,6 +10,7 @@ import (
|
||||
"math"
|
||||
"time"
|
||||
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/rvbox/rvbox/internal/domain"
|
||||
)
|
||||
|
||||
@@ -33,6 +35,84 @@ type StdinWriteResult struct {
|
||||
Duplicate bool
|
||||
}
|
||||
|
||||
// StdinIntent is a server-owned input record ready for delivery to a live
|
||||
// client session. It is returned in command/write order and contains a copy of
|
||||
// the decompressed payload; callers must not mutate storage-owned buffers.
|
||||
type StdinIntent struct {
|
||||
IssueUUID domain.UUID
|
||||
WriteSeq uint64
|
||||
Data []byte
|
||||
AppendNewline bool
|
||||
Close bool
|
||||
}
|
||||
|
||||
// PendingStdin returns unacknowledged input intents for one client. Delivery is
|
||||
// intentionally tracked by the session dispatcher, not by this durable query:
|
||||
// a lost connection simply causes the next session to replay the same write
|
||||
// sequence, which the client acknowledges idempotently.
|
||||
func (store *Store) PendingStdin(ctx context.Context, clientID string) ([]StdinIntent, error) {
|
||||
if clientID == "" {
|
||||
return nil, errors.New("client ID is required")
|
||||
}
|
||||
database, err := store.openDatabase()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := database.QueryContext(ctx, `SELECT w.issue_uuid, w.write_seq, w.payload, w.raw_bytes, w.stored_bytes, w.compression, w.sha256, w.append_newline, w.close_intent
|
||||
FROM stdin_writes w JOIN commands c ON c.issue_uuid = w.issue_uuid
|
||||
WHERE c.client_id = ? AND c.lifecycle BETWEEN 1 AND 4 AND w.acknowledged = 0
|
||||
ORDER BY c.issue_time, c.issue_uuid, w.write_seq`, clientID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var result []StdinIntent
|
||||
for rows.Next() {
|
||||
var encodedIssue, stored, digest []byte
|
||||
var writeSeq, rawBytes, storedBytes uint64
|
||||
var compression uint32
|
||||
var appendNewline, closeIntent bool
|
||||
if err := rows.Scan(&encodedIssue, &writeSeq, &stored, &rawBytes, &storedBytes, &compression, &digest, &appendNewline, &closeIntent); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(encodedIssue) != len(domain.UUID{}) || writeSeq == 0 || compression != 2 || storedBytes != uint64(len(stored)) || len(digest) != sha256.Size || rawBytes > maxStoredExecutionSpecBytes {
|
||||
return nil, ErrInvalidSegmentRecord
|
||||
}
|
||||
var data []byte
|
||||
data, err = decompressStdinPayload(stored, rawBytes)
|
||||
if err != nil || (!closeIntent && rawBytes == 0) || (closeIntent && rawBytes != 0) {
|
||||
return nil, ErrInvalidSegmentRecord
|
||||
}
|
||||
computed := sha256.Sum256(data)
|
||||
if string(computed[:]) != string(digest) {
|
||||
return nil, ErrInvalidSegmentRecord
|
||||
}
|
||||
var issue domain.UUID
|
||||
copy(issue[:], encodedIssue)
|
||||
result = append(result, StdinIntent{IssueUUID: issue, WriteSeq: writeSeq, Data: data, AppendNewline: appendNewline, Close: closeIntent})
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func decompressStdinPayload(stored []byte, expected uint64) ([]byte, error) {
|
||||
if expected > maxStoredExecutionSpecBytes {
|
||||
return nil, ErrInvalidSegmentRecord
|
||||
}
|
||||
decoder, err := zstd.NewReader(bytes.NewReader(stored), zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(maxStoredExecutionSpecBytes+1))
|
||||
if err != nil {
|
||||
return nil, ErrInvalidSegmentRecord
|
||||
}
|
||||
defer decoder.Close()
|
||||
data, err := decoder.DecodeAll(stored, nil)
|
||||
if err != nil || uint64(len(data)) != expected {
|
||||
return nil, ErrInvalidSegmentRecord
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// AppendStdin durably records one ordered stdin intent. It does not claim
|
||||
// delivery to a process; an agent acknowledgement is a later protocol event.
|
||||
// Reusing RequestUUID with the same immutable hash returns the original write
|
||||
|
||||
Reference in New Issue
Block a user