eric/distributor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 7 Wiki Activity

Stage uploads before accepting HTTP runs

Browse Source
This commit is contained in:
Eric Rakestraw
2026-06-04 00:14:32 +00:00
parent 1340418a2b
commit 0d346dcdf5
12 changed files with 248 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files

112
internal/app/upload_coordinator.go
View File

@@ -78,6 +78,7 @@ type UploadCoordinator struct {
queueSize int
maxConcurrency int
runningCount int
reservedCount int
activePipeline map[string]bool
pending []*uploadJob
records map[UploadRunID]UploadRunRecord
@@ -88,9 +89,10 @@ type uploadStageFunc func(context.Context, ingest.StageOptions) (ingest.StagedBu
type uploadRunFunc func(context.Context, config.Config, RunPipelineWithLocalSourceOptions) (RunReport, error)
type uploadJob struct {
recordID UploadRunID
request UploadRequest
pipeline config.Pipeline
recordID UploadRunID
request UploadRequest
pipeline config.Pipeline
stagedRoot string
}
type uploadCoordinatorHooks struct {
@@ -164,24 +166,49 @@ func (coordinator *UploadCoordinator) Submit(ctx context.Context, request Upload
if err != nil {
return UploadRunRecord{}, err
}
if err := ingest.ValidateContentType(request.ContentType); err != nil {
return UploadRunRecord{}, err
}
coordinator.mu.Lock()
coordinator.expireLocked(coordinator.now().UTC())
if coordinator.queueFullLocked() {
coordinator.mu.Unlock()
return UploadRunRecord{}, UploadQueueFullError{QueueSize: coordinator.queueSize}
}
coordinator.reservedCount++
coordinator.mu.Unlock()
staged, err := coordinator.stage(ctx, ingest.StageOptions{
Body: request.Body,
ContentType: request.ContentType,
PipelineStagingPath: pipeline.Source.Upload.StagingPath,
RunID: string(runID),
MaxUploadSize: int64(*pipeline.Source.Upload.MaxUploadSize),
MaxExtractedSize: int64(*pipeline.Source.Upload.MaxUploadSize),
MaxFileCount: uploadMaxFileCount(request.MaxFileCount),
})
if err != nil {
coordinator.releaseReservation()
return UploadRunRecord{}, err
}
coordinator.mu.Lock()
defer coordinator.mu.Unlock()
coordinator.expireLocked(coordinator.now().UTC())
if len(coordinator.pending) >= coordinator.queueSize {
return UploadRunRecord{}, UploadQueueFullError{QueueSize: coordinator.queueSize}
}
coordinator.reservedCount--
record := UploadRunRecord{
ID: runID,
PipelineID: pipeline.ID,
Status: UploadStatusAccepted,
AcceptedAt: coordinator.now().UTC(),
StagedRoot: staged.Root,
}
coordinator.records[runID] = record
coordinator.pending = append(coordinator.pending, &uploadJob{
recordID: runID,
request: request,
pipeline: pipeline,
recordID: runID,
request: request,
pipeline: pipeline,
stagedRoot: staged.Root,
})
coordinator.notify()
return record, nil
@@ -205,13 +232,13 @@ func (coordinator *UploadCoordinator) CanAccept() bool {
coordinator.mu.Lock()
defer coordinator.mu.Unlock()
coordinator.expireLocked(coordinator.now().UTC())
return len(coordinator.pending) < coordinator.queueSize
return !coordinator.queueFullLocked()
}
func (coordinator *UploadCoordinator) QueueDepth() int {
coordinator.mu.Lock()
defer coordinator.mu.Unlock()
return len(coordinator.pending)
return len(coordinator.pending) + coordinator.reservedCount
}
func (coordinator *UploadCoordinator) RunningCount() int {
@@ -290,47 +317,30 @@ func (coordinator *UploadCoordinator) markPendingQueuedLocked() {
}
func (coordinator *UploadCoordinator) runJob(job *uploadJob) {
record := coordinator.currentRecord(job.recordID)
maxFileCount := job.request.MaxFileCount
if maxFileCount <= 0 {
maxFileCount = DefaultUploadMaxFileCount
}
staged, err := coordinator.stage(coordinator.ctx, ingest.StageOptions{
Body: job.request.Body,
ContentType: job.request.ContentType,
PipelineStagingPath: job.pipeline.Source.Upload.StagingPath,
RunID: string(record.ID),
MaxUploadSize: int64(*job.pipeline.Source.Upload.MaxUploadSize),
MaxExtractedSize: int64(*job.pipeline.Source.Upload.MaxUploadSize),
MaxFileCount: maxFileCount,
report, err := coordinator.run(coordinator.ctx, coordinator.cfg, RunPipelineWithLocalSourceOptions{
PipelineID: job.pipeline.ID,
SourceRoot: job.stagedRoot,
DryRun: job.request.DryRun,
Force: job.request.Force,
})
if err == nil {
coordinator.setStagedRoot(job.recordID, staged.Root)
var report RunReport
report, err = coordinator.run(coordinator.ctx, coordinator.cfg, RunPipelineWithLocalSourceOptions{
PipelineID: job.pipeline.ID,
SourceRoot: staged.Root,
DryRun: job.request.DryRun,
Force: job.request.Force,
})
coordinator.complete(job, &report, err)
return
coordinator.complete(job, &report, err)
}
func (coordinator *UploadCoordinator) releaseReservation() {
coordinator.mu.Lock()
defer coordinator.mu.Unlock()
coordinator.reservedCount--
}
func (coordinator *UploadCoordinator) queueFullLocked() bool {
return len(coordinator.pending)+coordinator.reservedCount >= coordinator.queueSize
}
func uploadMaxFileCount(value int) int {
if value > 0 {
return value
}
coordinator.complete(job, nil, err)
}
func (coordinator *UploadCoordinator) currentRecord(runID UploadRunID) UploadRunRecord {
coordinator.mu.Lock()
defer coordinator.mu.Unlock()
return coordinator.records[runID]
}
func (coordinator *UploadCoordinator) setStagedRoot(runID UploadRunID, root string) {
coordinator.mu.Lock()
defer coordinator.mu.Unlock()
record := coordinator.records[runID]
record.StagedRoot = root
coordinator.records[runID] = record
return DefaultUploadMaxFileCount
}
func (coordinator *UploadCoordinator) complete(job *uploadJob, report *RunReport, runErr error) {