Add persistent job claim and file artifact reads

This commit is contained in:
npc0-hue
2026-08-26 22:01:55 +08:00
parent bb051f6da1
commit 3b4e857ef7
8 changed files with 388 additions and 35 deletions
+19
View File
@@ -8,6 +8,7 @@ import (
"crypto/sha256" "crypto/sha256"
"encoding/hex" "encoding/hex"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"io" "io"
"log" "log"
@@ -152,7 +153,25 @@ func (c PlatformClient) ReportLifecycle(ctx context.Context, request protocol.Ru
} }
func (c PlatformClient) ClaimJob(ctx context.Context, request protocol.RunJobClaimRequest) (protocol.RunJobClaimResponse, error) { func (c PlatformClient) ClaimJob(ctx context.Context, request protocol.RunJobClaimRequest) (protocol.RunJobClaimResponse, error) {
response, err := postPlatformJSON[protocol.RunJobClaimRequest, protocol.RunJobClaimResponse](ctx, c, "/api/v1/run/jobs/claim", request)
if request.WaitSeconds > 0 && claimWaitUnsupported(err) {
request.WaitSeconds = 0
return postPlatformJSON[protocol.RunJobClaimRequest, protocol.RunJobClaimResponse](ctx, c, "/api/v1/run/jobs/claim", request) return postPlatformJSON[protocol.RunJobClaimRequest, protocol.RunJobClaimResponse](ctx, c, "/api/v1/run/jobs/claim", request)
}
return response, err
}
func claimWaitUnsupported(err error) bool {
var requestErr PlatformRequestError
if !errors.As(err, &requestErr) || requestErr.Status != http.StatusBadRequest || requestErr.Code != "bad_request" {
return false
}
for _, detail := range requestErr.Details {
if strings.Contains(detail, "waitSeconds") {
return true
}
}
return false
} }
func (c PlatformClient) AckJob(ctx context.Context, request protocol.RunJobAckRequest) (protocol.RunJobAckResponse, error) { func (c PlatformClient) AckJob(ctx context.Context, request protocol.RunJobAckRequest) (protocol.RunJobAckResponse, error) {
+1
View File
@@ -321,6 +321,7 @@ type RunJobClaimRequest struct {
SessionToken string `json:"sessionToken"` SessionToken string `json:"sessionToken"`
Capabilities []string `json:"capabilities"` Capabilities []string `json:"capabilities"`
Capacity RunCapacityReport `json:"capacity"` Capacity RunCapacityReport `json:"capacity"`
WaitSeconds int `json:"waitSeconds,omitempty"`
} }
type RunJobClaimResponse struct { type RunJobClaimResponse struct {
+2 -2
View File
@@ -4,7 +4,7 @@ Jobs execute bounded server management work.
## Implemented Routes ## Implemented Routes
- `POST /api/v1/run/jobs/claim`: claims one queued job for the registered run endpoint. - `POST /api/v1/run/jobs/claim`: claims one queued job for the registered run endpoint, optionally holding the request for a bounded wait window so Platform can wake Run immediately when work arrives.
- `POST /api/v1/run/jobs/ack`: acknowledges an active leased job before execution. - `POST /api/v1/run/jobs/ack`: acknowledges an active leased job before execution.
- `POST /api/v1/run/jobs/progress`: reports bounded progress for an active leased job. - `POST /api/v1/run/jobs/progress`: reports bounded progress for an active leased job.
- `POST /api/v1/run/jobs/result`: submits a bounded terminal result for an active leased job. - `POST /api/v1/run/jobs/result`: submits a bounded terminal result for an active leased job.
@@ -17,7 +17,7 @@ Jobs execute bounded server management work.
## Payloads ## Payloads
- `RunJobClaimRequest`: session token, run ID, capacity, and supported capabilities. - `RunJobClaimRequest`: session token, run ID, capacity, supported capabilities, and optional `waitSeconds` for long-poll claim waiting.
- `RunJobClaimResponse`: optional job assignment with identity, capability, server instance, logical target key, scoped input ref, idempotency key, per-job attempt, max attempts, raw one-use lease token, ack deadline, execution lease deadline, and polling hint. Platform persists only the lease hash. - `RunJobClaimResponse`: optional job assignment with identity, capability, server instance, logical target key, scoped input ref, idempotency key, per-job attempt, max attempts, raw one-use lease token, ack deadline, execution lease deadline, and polling hint. Platform persists only the lease hash.
- `RunJobAckRequest`: job ID, run ID, session token, lease token, attempt, and bounded message. - `RunJobAckRequest`: job ID, run ID, session token, lease token, attempt, and bounded message.
- `RunJobProgressRequest`: job ID, run ID, session token, lease token, attempt, percent, sequence, and bounded message. - `RunJobProgressRequest`: job ID, run ID, session token, lease token, attempt, percent, sequence, and bounded message.
+30
View File
@@ -676,6 +676,36 @@ func TestDeploymentFileExecutorUsesServerRoot(t *testing.T) {
} }
} }
func TestDeploymentFileExecutorListsNestedPathsFromServerRoot(t *testing.T) {
workspaceRoot := t.TempDir()
serverRoot := t.TempDir()
if err := os.MkdirAll(filepath.Join(serverRoot, "SCUM", "Saved"), 0o700); err != nil {
t.Fatalf("mkdir server root fixture: %v", err)
}
if err := os.WriteFile(filepath.Join(serverRoot, "SCUM", "Saved", "ServerSettings.ini"), []byte("[/Script/SCUM.ServerSettings]\n"), 0o600); err != nil {
t.Fatalf("write server root fixture: %v", err)
}
executor, err := NewFileExecutor(workspaceRoot)
if err != nil {
t.Fatalf("new file executor: %v", err)
}
assignment := executionAssignment(protocol.RunCapabilityFilesList)
assignment.TargetKey = "server-root"
assignment.ExecutionInput.Deployment = &protocol.ServerDeploymentExecution{SchemaVersion: "1", Mode: "guided-install", ServerRoot: serverRoot, Revision: 7}
assignment.ExecutionInput.Inputs = map[string]string{"path": "SCUM", "recursive": "true", "query": ""}
result := executor.Execute(context.Background(), assignment)
if result.State != lifecycleResultStateSucceeded || result.ExecutionResult.Kind != "file.list" {
t.Fatalf("expected deployment root listing, got %+v", result)
}
var envelope fileListEnvelope
if err := json.Unmarshal([]byte(result.ExecutionResult.Content), &envelope); err != nil {
t.Fatalf("decode listing: %v", err)
}
if envelope.Path != "SCUM" || len(envelope.Entries) != 2 || envelope.Entries[0].RelativePath != "SCUM/Saved" || envelope.Entries[1].RelativePath != "SCUM/Saved/ServerSettings.ini" {
t.Fatalf("expected deployment-root relative paths, got %+v", envelope)
}
}
func TestScopedFileExecutorCancellationLeavesTargetUnchanged(t *testing.T) { func TestScopedFileExecutorCancellationLeavesTargetUnchanged(t *testing.T) {
root := t.TempDir() root := t.TempDir()
executor, err := NewFileExecutor(root) executor, err := NewFileExecutor(root)
+129
View File
@@ -0,0 +1,129 @@
package runtime
import (
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"io"
"os"
"browser.local/run/protocol"
)
func (worker *Worker) uploadFileArtifact(ctx context.Context, assignment protocol.RunJobAssignment, artifactID string, targetPath string, sizeBytes int64, checksum string) error {
state, err := worker.registeredState()
if err != nil {
return err
}
opened, err := worker.client.OpenArtifactTransfer(ctx, protocol.ArtifactTransferOpenRequest{
RunEndpointID: state.RunEndpointID,
SessionToken: state.SessionToken,
ArtifactID: artifactID,
Direction: "upload",
OwnerKind: "job",
OwnerID: assignment.JobID,
SizeBytes: sizeBytes,
ChunkSizeBytes: fileArtifactChunkSize,
Checksum: checksum,
IdempotencyKey: "file-read:" + assignment.JobID,
})
if err != nil {
return err
}
received := map[int]bool{}
for _, index := range opened.ReceivedChunkIndexes {
received[index] = true
}
file, err := os.Open(targetPath)
if err != nil {
return err
}
defer file.Close()
buffer := make([]byte, fileArtifactChunkSize)
for index, offset := 0, int64(0); offset < sizeBytes; index, offset = index+1, offset+int64(fileArtifactChunkSize) {
length := fileArtifactChunkSize
if remaining := sizeBytes - offset; remaining < int64(length) {
length = int(remaining)
}
if received[index] {
continue
}
if err := ctx.Err(); err != nil {
return err
}
read, err := file.ReadAt(buffer[:length], offset)
if err != nil && !(errors.Is(err, io.EOF) && read == length) {
return err
}
if read != length {
return fmt.Errorf("file artifact chunk is shorter than expected")
}
chunk := append([]byte(nil), buffer[:length]...)
state, err := worker.registeredState()
if err != nil {
return err
}
if _, err := worker.client.UploadArtifactChunk(ctx, protocol.ArtifactChunkUploadRequest{
RunEndpointID: state.RunEndpointID,
SessionToken: state.SessionToken,
TransferID: opened.TransferID,
ArtifactID: artifactID,
ChunkIndex: index,
Offset: offset,
SizeBytes: length,
Checksum: bytesChecksum(chunk),
Payload: chunk,
}); err != nil {
return err
}
}
state, err = worker.registeredState()
if err != nil {
return err
}
completed, err := worker.client.CompleteArtifactTransfer(ctx, protocol.ArtifactTransferCompleteRequest{
RunEndpointID: state.RunEndpointID,
SessionToken: state.SessionToken,
TransferID: opened.TransferID,
ArtifactID: artifactID,
Checksum: checksum,
SizeBytes: sizeBytes,
})
if err != nil {
return err
}
if !completed.Completed || completed.Artifact.State != "available" {
return fmt.Errorf("artifact transfer did not complete")
}
return nil
}
func checksumServerFileForArtifact(ctx context.Context, targetPath string) (string, error) {
file, err := os.Open(targetPath)
if err != nil {
return "", err
}
defer file.Close()
hash := sha256.New()
buffer := make([]byte, fileArtifactChunkSize)
for {
if err := ctx.Err(); err != nil {
return "", err
}
read, err := file.Read(buffer)
if read > 0 {
if _, writeErr := hash.Write(buffer[:read]); writeErr != nil {
return "", writeErr
}
}
if errors.Is(err, io.EOF) {
break
}
if err != nil {
return "", err
}
}
return "sha256:" + hex.EncodeToString(hash.Sum(nil)), nil
}
+19 -3
View File
@@ -101,6 +101,21 @@ func (executor *FileExecutor) scopeForAssignment(assignment protocol.RunJobAssig
return scope, false, err return scope, false, err
} }
func (executor *FileExecutor) existingReadTargetForAssignment(assignment protocol.RunJobAssignment) (string, string, bool, error) {
scope, deploymentRoot, err := executor.scopeForAssignment(assignment)
if err != nil {
return "", "", false, err
}
targetKey := assignment.TargetKey
if deploymentRoot {
targetKey = deploymentTargetKey(targetKey)
filePath, err := existingDeploymentTarget(scope, targetKey)
return scope, filePath, deploymentRoot, err
}
filePath, err := executor.resolver.ExistingTarget(scope, targetKey)
return scope, filePath, deploymentRoot, err
}
type fileListEntry struct { type fileListEntry struct {
Name string `json:"name"` Name string `json:"name"`
Kind string `json:"kind"` Kind string `json:"kind"`
@@ -175,11 +190,12 @@ func (executor *FileExecutor) list(ctx context.Context, scope string, deployment
if err != nil { if err != nil {
return err return err
} }
logicalKey := path.Join(directoryKey, rel) entryRelativePath := rel
if relativePath != "." { if relativePath != "." {
logicalKey = path.Join(directoryKey, relativePath, rel) entryRelativePath = path.Join(relativePath, rel)
} }
candidate := append(entries, fileListEntry{Name: name, Kind: kind, RelativePath: rel, LogicalKey: logicalKey, SizeBytes: entryInfo.Size(), ModifiedAt: entryInfo.ModTime().UTC().Format(time.RFC3339Nano)}) logicalKey := path.Join(directoryKey, entryRelativePath)
candidate := append(entries, fileListEntry{Name: name, Kind: kind, RelativePath: entryRelativePath, LogicalKey: logicalKey, SizeBytes: entryInfo.Size(), ModifiedAt: entryInfo.ModTime().UTC().Format(time.RFC3339Nano)})
body, marshalErr := json.Marshal(fileListEnvelope{DirectoryKey: directoryKey, Path: strings.TrimPrefix(relativePath, "."), Entries: candidate}) body, marshalErr := json.Marshal(fileListEnvelope{DirectoryKey: directoryKey, Path: strings.TrimPrefix(relativePath, "."), Entries: candidate})
if marshalErr != nil { if marshalErr != nil {
return marshalErr return marshalErr
+97 -11
View File
@@ -8,6 +8,7 @@ import (
"errors" "errors"
"fmt" "fmt"
"log" "log"
"net/http"
"os" "os"
"path/filepath" "path/filepath"
"runtime" "runtime"
@@ -46,7 +47,10 @@ type WorkerClient interface {
const ( const (
jobActivePollInterval = 10 * time.Second jobActivePollInterval = 10 * time.Second
jobClaimWaitSeconds = 25
durableUploaderFlushTimeout = 5 * time.Second durableUploaderFlushTimeout = 5 * time.Second
fileArtifactChunkSize = 1024 * 1024
maxFileArtifactBytes = int64(512 * 1024 * 1024)
) )
type Worker struct { type Worker struct {
@@ -436,6 +440,7 @@ func (worker *Worker) ClaimAndRunOnce(ctx context.Context) (bool, error) {
SessionToken: state.SessionToken, SessionToken: state.SessionToken,
Capabilities: state.Capabilities, Capabilities: state.Capabilities,
Capacity: worker.capacityReportFor(state), Capacity: worker.capacityReportFor(state),
WaitSeconds: jobClaimWaitSeconds,
}) })
if err != nil { if err != nil {
if sessionInvalidError(err) { if sessionInvalidError(err) {
@@ -728,7 +733,7 @@ func (worker *Worker) executeAssignment(ctx context.Context, assignment protocol
} }
if supportedCapability(SupportedFileCapabilities(), assignment.Capability) { if supportedCapability(SupportedFileCapabilities(), assignment.Capability) {
log.Printf("RUN phase=job.dispatch status=selected job=%s executor=file", assignment.JobID) log.Printf("RUN phase=job.dispatch status=selected job=%s executor=file", assignment.JobID)
return worker.executor.ExecuteContext(ctx, assignment) return worker.executeFileJob(ctx, assignment)
} }
if isSupportedDistributionCapability(assignment.Capability) { if isSupportedDistributionCapability(assignment.Capability) {
log.Printf("RUN phase=job.dispatch status=selected job=%s executor=distribution", assignment.JobID) log.Printf("RUN phase=job.dispatch status=selected job=%s executor=distribution", assignment.JobID)
@@ -742,6 +747,50 @@ func (worker *Worker) executeAssignment(ctx context.Context, assignment protocol
return lifecycleFailure("unsupported_run_capability", "unsupported run capability") return lifecycleFailure("unsupported_run_capability", "unsupported run capability")
} }
func (worker *Worker) executeFileJob(ctx context.Context, assignment protocol.RunJobAssignment) LifecycleExecutionResult {
if assignment.Capability != protocol.RunCapabilityFilesRead {
return worker.executor.ExecuteContext(ctx, assignment)
}
fileExecutor := worker.executor.fileExecutor
if fileExecutor == nil {
return lifecycleExecutionFailure("file_executor_unavailable", "file executor is unavailable", false)
}
scope, targetPath, deploymentRoot, err := fileExecutor.existingReadTargetForAssignment(assignment)
if err != nil {
return lifecycleExecutionFailure("file_read_failed", err.Error(), false)
}
info, err := os.Stat(targetPath)
if err != nil {
return lifecycleExecutionFailure("file_read_failed", err.Error(), false)
}
limit := assignment.ExecutionInput.MaxReadBytes
if limit <= 0 || limit > maxExecutionContentBytes {
limit = maxExecutionContentBytes
}
if info.Size() <= int64(limit) {
return fileExecutor.read(ctx, scope, deploymentRoot, assignment)
}
if info.Size() > maxFileArtifactBytes {
return lifecycleExecutionFailure("file_read_too_large", "file exceeds artifact transfer limit", false)
}
checksum, err := checksumServerFileForArtifact(ctx, targetPath)
if err != nil {
if errors.Is(err, context.Canceled) {
return lifecycleExecutionFailure("file_cancelled", "file read cancelled", false)
}
return lifecycleExecutionFailure("file_read_failed", err.Error(), false)
}
artifactID := "artifact-" + safeWorkspaceName(assignment.JobID) + "-file-read"
if err := worker.uploadFileArtifact(ctx, assignment, artifactID, targetPath, info.Size(), checksum); err != nil {
if errors.Is(err, context.Canceled) {
return lifecycleExecutionFailure("file_cancelled", "file read cancelled", false)
}
return lifecycleExecutionFailure("file_artifact_upload_failed", "file artifact upload failed", false)
}
metadata := fileExecutor.metadata(scope, assignment.TargetKey, checksum, info.Size())
return LifecycleExecutionResult{State: lifecycleResultStateSucceeded, Progress: protocol.RunJobProgressReport{Percent: 100, Message: "file read completed"}, ResultRef: "artifact://" + artifactID, Message: "file read completed", ExecutionResult: protocol.RunJobExecutionResult{Kind: "file.read", Version: metadata.Version, Checksum: checksum, SizeBytes: info.Size(), Summary: "large file transferred as artifact"}}
}
func (worker *Worker) executeProtectedRequestJob(ctx context.Context, assignment protocol.RunJobAssignment) LifecycleExecutionResult { func (worker *Worker) executeProtectedRequestJob(ctx context.Context, assignment protocol.RunJobAssignment) LifecycleExecutionResult {
if err := protocol.ValidateRunJobAssignment(assignment); err != nil { if err := protocol.ValidateRunJobAssignment(assignment); err != nil {
return protectedRequestFailure(assignment.Capability, ProtectedRequestStatusFailed, "protected_request_assignment_invalid") return protectedRequestFailure(assignment.Capability, ProtectedRequestStatusFailed, "protected_request_assignment_invalid")
@@ -988,31 +1037,38 @@ func (worker *Worker) reportRunUpdateHealth(ctx context.Context) error {
} }
func (worker *Worker) runJobLoop(ctx context.Context, interval time.Duration) error { func (worker *Worker) runJobLoop(ctx context.Context, interval time.Duration) error {
log.Printf("RUN phase=job_loop status=starting pollMs=%d", interval.Milliseconds()) log.Printf("RUN phase=job_loop status=starting fallbackPollMs=%d claimWaitSeconds=%d", interval.Milliseconds(), jobClaimWaitSeconds)
ticker := time.NewTicker(interval)
defer ticker.Stop()
for { for {
select { select {
case <-ctx.Done(): case <-ctx.Done():
log.Printf("RUN phase=job_loop status=context_done error=%s", RedactText(ctx.Err().Error())) log.Printf("RUN phase=job_loop status=context_done error=%s", RedactText(ctx.Err().Error()))
return ctx.Err() return ctx.Err()
case <-ticker.C: default:
}
if worker.journal.ActiveCount() > 0 { if worker.journal.ActiveCount() > 0 {
log.Printf("RUN phase=job_loop status=active_jobs activeJobs=%d", worker.journal.ActiveCount()) log.Printf("RUN phase=job_loop status=active_jobs activeJobs=%d", worker.journal.ActiveCount())
if err := worker.ReconcileOnce(ctx); err != nil { if err := worker.ReconcileOnce(ctx); err != nil {
log.Printf("RUN phase=job_loop status=reconcile_failed error=%s", RedactText(err.Error())) log.Printf("RUN phase=job_loop status=reconcile_failed error=%s", RedactText(err.Error()))
ticker.Reset(boundedRetryBackoff(worker.cfg.RetryBackoff)) if err := waitWorkerLoop(ctx, boundedRetryBackoff(worker.cfg.RetryBackoff)); err != nil {
return err
}
continue continue
} }
if err := worker.RecoverActiveJobs(ctx); err != nil { if err := worker.RecoverActiveJobs(ctx); err != nil {
log.Printf("RUN phase=job_loop status=recover_failed error=%s", RedactText(err.Error())) log.Printf("RUN phase=job_loop status=recover_failed error=%s", RedactText(err.Error()))
ticker.Reset(boundedRetryBackoff(worker.cfg.RetryBackoff)) if err := waitWorkerLoop(ctx, boundedRetryBackoff(worker.cfg.RetryBackoff)); err != nil {
return err
}
continue continue
} }
} }
if _, err := worker.ClaimAndRunOnce(ctx); err != nil { claimStartedAt := time.Now()
handled, err := worker.ClaimAndRunOnce(ctx)
if err != nil {
log.Printf("RUN phase=job_loop status=claim_failed error=%s", RedactText(err.Error())) log.Printf("RUN phase=job_loop status=claim_failed error=%s", RedactText(err.Error()))
ticker.Reset(boundedRetryBackoff(worker.cfg.RetryBackoff)) if err := waitWorkerLoop(ctx, boundedRetryBackoff(worker.cfg.RetryBackoff)); err != nil {
return err
}
continue continue
} }
worker.restartMu.Lock() worker.restartMu.Lock()
@@ -1022,9 +1078,26 @@ func (worker *Worker) runJobLoop(ctx context.Context, interval time.Duration) er
log.Printf("RUN phase=job_loop status=restart_requested") log.Printf("RUN phase=job_loop status=restart_requested")
return ErrSelfUpdateRestartRequested return ErrSelfUpdateRestartRequested
} }
ticker.Reset(interval) if !handled && time.Since(claimStartedAt) < interval {
if err := waitWorkerLoop(ctx, interval-time.Since(claimStartedAt)); err != nil {
return err
} }
} }
}
}
func waitWorkerLoop(ctx context.Context, delay time.Duration) error {
if delay <= 0 {
return nil
}
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
} }
type durableLogClient interface { type durableLogClient interface {
@@ -1112,6 +1185,11 @@ func logBatchSessionMetadataMismatchError(err error) bool {
return errors.As(err, &mismatch) && mismatch.LogBatchSessionMetadataMismatch() return errors.As(err, &mismatch) && mismatch.LogBatchSessionMetadataMismatch()
} }
func artifactChunkMissingOnPlatform(err error) bool {
var httpErr interface{ HTTPStatus() int }
return errors.As(err, &httpErr) && httpErr.HTTPStatus() == http.StatusNotFound
}
type sessionArtifactChunkClient struct { type sessionArtifactChunkClient struct {
client durableArtifactClient client durableArtifactClient
runEndpointID string runEndpointID string
@@ -1121,7 +1199,15 @@ type sessionArtifactChunkClient struct {
func (client sessionArtifactChunkClient) UploadArtifactChunk(ctx context.Context, chunk protocol.ArtifactChunkUploadRequest) (protocol.ArtifactChunkUploadResponse, error) { func (client sessionArtifactChunkClient) UploadArtifactChunk(ctx context.Context, chunk protocol.ArtifactChunkUploadRequest) (protocol.ArtifactChunkUploadResponse, error) {
chunk.RunEndpointID = client.runEndpointID chunk.RunEndpointID = client.runEndpointID
chunk.SessionToken = client.sessionToken chunk.SessionToken = client.sessionToken
return client.client.UploadArtifactChunk(ctx, chunk) response, err := client.client.UploadArtifactChunk(ctx, chunk)
if err != nil {
if artifactChunkMissingOnPlatform(err) {
log.Printf("RUN phase=durable_uploaders.artifacts status=drop_stale transfer=%s artifact=%s chunk=%d reason=platform_transfer_missing", safeOptional(chunk.TransferID), safeOptional(chunk.ArtifactID), chunk.ChunkIndex)
return protocol.ArtifactChunkUploadResponse{Accepted: true, TransferID: chunk.TransferID, ArtifactID: chunk.ArtifactID, ChunkIndex: chunk.ChunkIndex}, nil
}
return response, err
}
return response, nil
} }
func (worker *Worker) runDurableUploaders(ctx context.Context, done chan<- struct{}) { func (worker *Worker) runDurableUploaders(ctx context.Context, done chan<- struct{}) {
+73 -1
View File
@@ -7,6 +7,8 @@ import (
"fmt" "fmt"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"os"
"path/filepath"
"reflect" "reflect"
"runtime" "runtime"
"strings" "strings"
@@ -173,6 +175,9 @@ func TestWorkerClaimsAcksProgressAndCompletesJob(t *testing.T) {
if len(client.ackRequests) != 1 || len(client.progressRequests) != 1 || len(client.resultRequests) != 1 || len(client.cancelPollRequests) != 1 { if len(client.ackRequests) != 1 || len(client.progressRequests) != 1 || len(client.resultRequests) != 1 || len(client.cancelPollRequests) != 1 {
t.Fatalf("expected ack/progress/result/cancel calls, got ack=%d progress=%d result=%d cancel=%d", len(client.ackRequests), len(client.progressRequests), len(client.resultRequests), len(client.cancelPollRequests)) t.Fatalf("expected ack/progress/result/cancel calls, got ack=%d progress=%d result=%d cancel=%d", len(client.ackRequests), len(client.progressRequests), len(client.resultRequests), len(client.cancelPollRequests))
} }
if len(client.claimRequests) != 1 || client.claimRequests[0].WaitSeconds != jobClaimWaitSeconds {
t.Fatalf("expected persistent claim wait, got %+v", client.claimRequests)
}
if client.progressRequests[0].Progress.Percent != 10 || client.resultRequests[0].State != "succeeded" || client.resultRequests[0].ResultRef == "" { if client.progressRequests[0].Progress.Percent != 10 || client.resultRequests[0].State != "succeeded" || client.resultRequests[0].ResultRef == "" {
t.Fatalf("unexpected job channel payloads: progress=%+v result=%+v", client.progressRequests[0], client.resultRequests[0]) t.Fatalf("unexpected job channel payloads: progress=%+v result=%+v", client.progressRequests[0], client.resultRequests[0])
} }
@@ -214,6 +219,56 @@ func TestWorkerDispatchesSelfUpdateJob(t *testing.T) {
} }
} }
func TestWorkerFileReadLargeFileUploadsArtifact(t *testing.T) {
cfg := workerTestConfig(t)
client := newFakeWorkerClient()
assignment := workerJobAssignment(protocol.RunCapabilityFilesRead)
assignment.TargetKey = "logs/big.log"
assignment.ExecutionInput.WorkspaceScope = "run-local"
assignment.ExecutionInput.MaxReadBytes = 4
client.claimJob = assignment
client.buildInput = protocol.DistributionBuildInputResponse{JobID: assignment.JobID}
scope, err := NewWorkspaceResolver(cfg.WorkspaceRoot).Scope(assignment.ServerInstanceID, assignment.ExecutionInput.WorkspaceScope)
if err != nil {
t.Fatalf("scope: %v", err)
}
if err := os.MkdirAll(filepath.Join(scope, "logs"), 0o700); err != nil {
t.Fatalf("mkdir fixture: %v", err)
}
payload := []byte(strings.Repeat("A", fileArtifactChunkSize) + "tail")
if err := os.WriteFile(filepath.Join(scope, assignment.TargetKey), payload, 0o600); err != nil {
t.Fatalf("write fixture: %v", err)
}
worker, err := NewWorker(cfg, client)
if err != nil {
t.Fatalf("new worker: %v", err)
}
if err := worker.Register(context.Background()); err != nil {
t.Fatalf("register: %v", err)
}
handled, err := worker.ClaimAndRunOnce(context.Background())
if err != nil || !handled {
t.Fatalf("claim/run handled=%v err=%v", handled, err)
}
if len(client.artifactOpenRequests) != 1 {
t.Fatalf("expected artifact transfer, got %+v", client.artifactOpenRequests)
}
opened := client.artifactOpenRequests[0]
if opened.OwnerID != assignment.JobID || opened.SizeBytes != int64(len(payload)) || opened.Checksum != bytesChecksum(payload) || opened.ChunkSizeBytes != fileArtifactChunkSize {
t.Fatalf("unexpected artifact open: %+v", opened)
}
if string(client.artifactPayload) != string(payload) {
t.Fatalf("uploaded artifact payload mismatch")
}
if len(client.resultRequests) != 1 || client.resultRequests[0].State != "succeeded" || !strings.HasPrefix(client.resultRequests[0].ResultRef, "artifact://artifact-job-worker-file-read") {
t.Fatalf("expected artifact result ref, got %+v", client.resultRequests)
}
if client.resultRequests[0].ExecutionResult.Content != "" || client.resultRequests[0].ExecutionResult.SizeBytes != int64(len(payload)) {
t.Fatalf("large file result must not inline content: %+v", client.resultRequests[0].ExecutionResult)
}
}
func TestWorkerRegistersPackageIdentity(t *testing.T) { func TestWorkerRegistersPackageIdentity(t *testing.T) {
client := newFakeWorkerClient() client := newFakeWorkerClient()
cfg := workerTestConfig(t) cfg := workerTestConfig(t)
@@ -608,6 +663,19 @@ func TestSessionArtifactChunkClientOverridesSpooledIdentity(t *testing.T) {
} }
} }
func TestSessionArtifactChunkClientDropsPlatformMissingTransfer(t *testing.T) {
recorder := &recordingDurableArtifactClient{err: api.PlatformRequestError{Status: http.StatusNotFound, Code: "not_found"}}
chunk := protocol.ArtifactChunkUploadRequest{RunEndpointID: "old-endpoint", SessionToken: "old-token", TransferID: "transfer-stale", ArtifactID: "artifact-stale", ChunkIndex: 3}
client := sessionArtifactChunkClient{client: recorder, runEndpointID: "run-current", sessionToken: "token-current"}
response, err := client.UploadArtifactChunk(context.Background(), chunk)
if err != nil || !response.Accepted || response.TransferID != chunk.TransferID || response.ChunkIndex != chunk.ChunkIndex {
t.Fatalf("expected stale chunk ack for queue cleanup, response=%+v err=%v", response, err)
}
if recorder.chunk.RunEndpointID != "run-current" || recorder.chunk.SessionToken != "token-current" {
t.Fatalf("expected current identity before stale drop, got %+v", recorder.chunk)
}
}
func TestWorkerRetryBackoffIsBounded(t *testing.T) { func TestWorkerRetryBackoffIsBounded(t *testing.T) {
if got := boundedRetryBackoff(75 * time.Millisecond); got != 75*time.Millisecond { if got := boundedRetryBackoff(75 * time.Millisecond); got != 75*time.Millisecond {
t.Fatalf("expected configured backoff, got %s", got) t.Fatalf("expected configured backoff, got %s", got)
@@ -640,7 +708,7 @@ func TestWorkerIntegrationWithPlatformLikeServer(t *testing.T) {
case "/api/v1/run/jobs/claim": case "/api/v1/run/jobs/claim":
var request protocol.RunJobClaimRequest var request protocol.RunJobClaimRequest
decodeWorkerTestJSON(t, r, &request) decodeWorkerTestJSON(t, r, &request)
if request.SessionToken != "session-token" || len(request.Capabilities) == 0 { if request.SessionToken != "session-token" || len(request.Capabilities) == 0 || request.WaitSeconds != jobClaimWaitSeconds {
t.Fatalf("unexpected claim: %+v", request) t.Fatalf("unexpected claim: %+v", request)
} }
writeWorkerTestJSON(t, w, protocol.RunJobClaimResponse{Accepted: true, RunEndpointID: request.RunEndpointID, HasJob: true, Job: &assignment, NextPollSeconds: 2, ServerTime: workerTestTime()}) writeWorkerTestJSON(t, w, protocol.RunJobClaimResponse{Accepted: true, RunEndpointID: request.RunEndpointID, HasJob: true, Job: &assignment, NextPollSeconds: 2, ServerTime: workerTestTime()})
@@ -811,10 +879,14 @@ func (client *recordingDurableLogClient) IngestLogBatch(_ context.Context, batch
type recordingDurableArtifactClient struct { type recordingDurableArtifactClient struct {
chunk protocol.ArtifactChunkUploadRequest chunk protocol.ArtifactChunkUploadRequest
err error
} }
func (client *recordingDurableArtifactClient) UploadArtifactChunk(_ context.Context, chunk protocol.ArtifactChunkUploadRequest) (protocol.ArtifactChunkUploadResponse, error) { func (client *recordingDurableArtifactClient) UploadArtifactChunk(_ context.Context, chunk protocol.ArtifactChunkUploadRequest) (protocol.ArtifactChunkUploadResponse, error) {
client.chunk = chunk client.chunk = chunk
if client.err != nil {
return protocol.ArtifactChunkUploadResponse{}, client.err
}
return protocol.ArtifactChunkUploadResponse{Accepted: true, TransferID: chunk.TransferID, ArtifactID: chunk.ArtifactID, ChunkIndex: chunk.ChunkIndex}, nil return protocol.ArtifactChunkUploadResponse{Accepted: true, TransferID: chunk.TransferID, ArtifactID: chunk.ArtifactID, ChunkIndex: chunk.ChunkIndex}, nil
} }