Add async upload coordination
This commit is contained in:
382
internal/app/upload_coordinator.go
Normal file
382
internal/app/upload_coordinator.go
Normal file
@@ -0,0 +1,382 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/distributor/internal/config"
|
||||
"gitea.maximumdirect.net/eric/distributor/internal/ingest"
|
||||
)
|
||||
|
||||
const DefaultUploadMaxFileCount = 4096
|
||||
|
||||
type UploadRunID string
|
||||
|
||||
type UploadStatus string
|
||||
|
||||
const (
|
||||
UploadStatusAccepted UploadStatus = "accepted"
|
||||
UploadStatusQueued UploadStatus = "queued"
|
||||
UploadStatusRunning UploadStatus = "running"
|
||||
UploadStatusSucceeded UploadStatus = "succeeded"
|
||||
UploadStatusFailed UploadStatus = "failed"
|
||||
UploadStatusExpired UploadStatus = "expired"
|
||||
)
|
||||
|
||||
type UploadRunRecord struct {
|
||||
ID UploadRunID `json:"run_id"`
|
||||
PipelineID string `json:"pipeline_id"`
|
||||
Status UploadStatus `json:"status"`
|
||||
AcceptedAt time.Time `json:"accepted_at"`
|
||||
StartedAt *time.Time `json:"started_at,omitempty"`
|
||||
FinishedAt *time.Time `json:"finished_at,omitempty"`
|
||||
Report *RunReport `json:"report,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
StagedRoot string `json:"-"`
|
||||
}
|
||||
|
||||
type UploadRequest struct {
|
||||
PipelineID string
|
||||
ContentType string
|
||||
Body io.Reader
|
||||
DryRun bool
|
||||
Force bool
|
||||
MaxFileCount int
|
||||
}
|
||||
|
||||
type UploadQueueFullError struct {
|
||||
QueueSize int
|
||||
}
|
||||
|
||||
func (err UploadQueueFullError) Error() string {
|
||||
return fmt.Sprintf("upload queue is full with capacity %d", err.QueueSize)
|
||||
}
|
||||
|
||||
func IsUploadQueueFull(err error) bool {
|
||||
var full UploadQueueFullError
|
||||
return errors.As(err, &full)
|
||||
}
|
||||
|
||||
type UploadCoordinator struct {
|
||||
ctx context.Context
|
||||
cfg config.Config
|
||||
stage uploadStageFunc
|
||||
run uploadRunFunc
|
||||
now func() time.Time
|
||||
randomSuffix func() (string, error)
|
||||
retention time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
signal chan struct{}
|
||||
queueSize int
|
||||
maxConcurrency int
|
||||
runningCount int
|
||||
activePipeline map[string]bool
|
||||
pending []*uploadJob
|
||||
records map[UploadRunID]UploadRunRecord
|
||||
}
|
||||
|
||||
type uploadStageFunc func(context.Context, ingest.StageOptions) (ingest.StagedBundle, error)
|
||||
|
||||
type uploadRunFunc func(context.Context, config.Config, RunPipelineWithLocalSourceOptions) (RunReport, error)
|
||||
|
||||
type uploadJob struct {
|
||||
recordID UploadRunID
|
||||
request UploadRequest
|
||||
pipeline config.Pipeline
|
||||
}
|
||||
|
||||
type uploadCoordinatorHooks struct {
|
||||
stage uploadStageFunc
|
||||
run uploadRunFunc
|
||||
now func() time.Time
|
||||
randomSuffix func() (string, error)
|
||||
}
|
||||
|
||||
func NewUploadCoordinator(ctx context.Context, cfg config.Config) *UploadCoordinator {
|
||||
return newUploadCoordinator(ctx, cfg, uploadCoordinatorHooks{})
|
||||
}
|
||||
|
||||
func newUploadCoordinator(ctx context.Context, cfg config.Config, hooks uploadCoordinatorHooks) *UploadCoordinator {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
config.ApplyDefaults(&cfg)
|
||||
stage := hooks.stage
|
||||
if stage == nil {
|
||||
stage = ingest.StageArchive
|
||||
}
|
||||
run := hooks.run
|
||||
if run == nil {
|
||||
run = runPipelineConfigWithLocalSource
|
||||
}
|
||||
now := hooks.now
|
||||
if now == nil {
|
||||
now = time.Now
|
||||
}
|
||||
randomSuffix := hooks.randomSuffix
|
||||
if randomSuffix == nil {
|
||||
randomSuffix = randomRunIDSuffix
|
||||
}
|
||||
coordinator := &UploadCoordinator{
|
||||
ctx: ctx,
|
||||
cfg: cfg,
|
||||
stage: stage,
|
||||
run: run,
|
||||
now: now,
|
||||
randomSuffix: randomSuffix,
|
||||
retention: cfg.Server.HTTP.Retention.AsDuration(),
|
||||
signal: make(chan struct{}, 1),
|
||||
queueSize: cfg.Server.HTTP.QueueSize,
|
||||
maxConcurrency: cfg.Server.HTTP.MaxConcurrency,
|
||||
activePipeline: map[string]bool{},
|
||||
records: map[UploadRunID]UploadRunRecord{},
|
||||
}
|
||||
go coordinator.dispatchLoop()
|
||||
return coordinator
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) Submit(ctx context.Context, request UploadRequest) (UploadRunRecord, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return UploadRunRecord{}, err
|
||||
}
|
||||
if request.Body == nil {
|
||||
return UploadRunRecord{}, fmt.Errorf("upload body is required")
|
||||
}
|
||||
pipeline, ok := findPipeline(coordinator.cfg, request.PipelineID)
|
||||
if !ok {
|
||||
return UploadRunRecord{}, PipelineNotFoundError{ID: request.PipelineID}
|
||||
}
|
||||
if pipeline.Source.Backend != config.BackendHTTPUpload {
|
||||
return UploadRunRecord{}, fmt.Errorf("pipeline %s source backend %s is not configured for uploads", pipeline.ID, pipeline.Source.Backend)
|
||||
}
|
||||
runID, err := coordinator.newRunID(pipeline.ID)
|
||||
if err != nil {
|
||||
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}
|
||||
}
|
||||
record := UploadRunRecord{
|
||||
ID: runID,
|
||||
PipelineID: pipeline.ID,
|
||||
Status: UploadStatusAccepted,
|
||||
AcceptedAt: coordinator.now().UTC(),
|
||||
}
|
||||
coordinator.records[runID] = record
|
||||
coordinator.pending = append(coordinator.pending, &uploadJob{
|
||||
recordID: runID,
|
||||
request: request,
|
||||
pipeline: pipeline,
|
||||
})
|
||||
coordinator.notify()
|
||||
return record, nil
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) Status(runID UploadRunID) (UploadRunRecord, bool) {
|
||||
coordinator.mu.Lock()
|
||||
defer coordinator.mu.Unlock()
|
||||
coordinator.expireLocked(coordinator.now().UTC())
|
||||
record, ok := coordinator.records[runID]
|
||||
return record, ok
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) Expire() []UploadRunRecord {
|
||||
coordinator.mu.Lock()
|
||||
defer coordinator.mu.Unlock()
|
||||
return coordinator.expireLocked(coordinator.now().UTC())
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) QueueDepth() int {
|
||||
coordinator.mu.Lock()
|
||||
defer coordinator.mu.Unlock()
|
||||
return len(coordinator.pending)
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) RunningCount() int {
|
||||
coordinator.mu.Lock()
|
||||
defer coordinator.mu.Unlock()
|
||||
return coordinator.runningCount
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) newRunID(pipelineID string) (UploadRunID, error) {
|
||||
suffix, err := coordinator.randomSuffix()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
timestamp := coordinator.now().UTC().Format("20060102T150405Z")
|
||||
return UploadRunID(pipelineID + "." + timestamp + "." + suffix), nil
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) dispatchLoop() {
|
||||
for {
|
||||
select {
|
||||
case <-coordinator.ctx.Done():
|
||||
return
|
||||
case <-coordinator.signal:
|
||||
for coordinator.startNext() {
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) startNext() bool {
|
||||
coordinator.mu.Lock()
|
||||
defer coordinator.mu.Unlock()
|
||||
if coordinator.runningCount >= coordinator.maxConcurrency {
|
||||
coordinator.markPendingQueuedLocked()
|
||||
return false
|
||||
}
|
||||
index := -1
|
||||
for candidateIndex, job := range coordinator.pending {
|
||||
if coordinator.activePipeline[job.pipeline.ID] {
|
||||
record := coordinator.records[job.recordID]
|
||||
if record.Status == UploadStatusAccepted {
|
||||
record.Status = UploadStatusQueued
|
||||
coordinator.records[job.recordID] = record
|
||||
}
|
||||
continue
|
||||
}
|
||||
index = candidateIndex
|
||||
break
|
||||
}
|
||||
if index < 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
job := coordinator.pending[index]
|
||||
coordinator.pending = append(coordinator.pending[:index], coordinator.pending[index+1:]...)
|
||||
now := coordinator.now().UTC()
|
||||
record := coordinator.records[job.recordID]
|
||||
record.Status = UploadStatusRunning
|
||||
record.StartedAt = &now
|
||||
coordinator.records[job.recordID] = record
|
||||
coordinator.runningCount++
|
||||
coordinator.activePipeline[job.pipeline.ID] = true
|
||||
|
||||
go coordinator.runJob(job)
|
||||
return true
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) markPendingQueuedLocked() {
|
||||
for _, job := range coordinator.pending {
|
||||
record := coordinator.records[job.recordID]
|
||||
if record.Status == UploadStatusAccepted {
|
||||
record.Status = UploadStatusQueued
|
||||
coordinator.records[job.recordID] = record
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
})
|
||||
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, 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
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) complete(job *uploadJob, report *RunReport, runErr error) {
|
||||
coordinator.mu.Lock()
|
||||
defer coordinator.mu.Unlock()
|
||||
record := coordinator.records[job.recordID]
|
||||
finishedAt := coordinator.now().UTC()
|
||||
record.FinishedAt = &finishedAt
|
||||
record.Report = report
|
||||
if runErr != nil {
|
||||
record.Status = UploadStatusFailed
|
||||
record.Error = runErr.Error()
|
||||
} else {
|
||||
record.Status = UploadStatusSucceeded
|
||||
}
|
||||
coordinator.records[job.recordID] = record
|
||||
coordinator.runningCount--
|
||||
delete(coordinator.activePipeline, job.pipeline.ID)
|
||||
coordinator.notify()
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) expireLocked(now time.Time) []UploadRunRecord {
|
||||
var expired []UploadRunRecord
|
||||
for runID, record := range coordinator.records {
|
||||
if record.FinishedAt == nil || record.Status == UploadStatusExpired {
|
||||
continue
|
||||
}
|
||||
if now.Before(record.FinishedAt.Add(coordinator.retention)) {
|
||||
continue
|
||||
}
|
||||
if record.StagedRoot != "" {
|
||||
_ = os.RemoveAll(record.StagedRoot)
|
||||
}
|
||||
record.Status = UploadStatusExpired
|
||||
record.Report = nil
|
||||
record.Error = ""
|
||||
expired = append(expired, record)
|
||||
delete(coordinator.records, runID)
|
||||
}
|
||||
return expired
|
||||
}
|
||||
|
||||
func (coordinator *UploadCoordinator) notify() {
|
||||
select {
|
||||
case coordinator.signal <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func randomRunIDSuffix() (string, error) {
|
||||
var data [4]byte
|
||||
if _, err := rand.Read(data[:]); err != nil {
|
||||
return "", fmt.Errorf("generate run id suffix: %w", err)
|
||||
}
|
||||
return hex.EncodeToString(data[:]), nil
|
||||
}
|
||||
Reference in New Issue
Block a user