400 lines
11 KiB
Go
400 lines
11 KiB
Go
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
|
|
reservedCount 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
|
|
stagedRoot string
|
|
}
|
|
|
|
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
|
|
}
|
|
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.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,
|
|
stagedRoot: staged.Root,
|
|
})
|
|
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) CanAccept() bool {
|
|
coordinator.mu.Lock()
|
|
defer coordinator.mu.Unlock()
|
|
coordinator.expireLocked(coordinator.now().UTC())
|
|
return !coordinator.queueFullLocked()
|
|
}
|
|
|
|
func (coordinator *UploadCoordinator) QueueDepth() int {
|
|
coordinator.mu.Lock()
|
|
defer coordinator.mu.Unlock()
|
|
return len(coordinator.pending) + coordinator.reservedCount
|
|
}
|
|
|
|
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) {
|
|
report, err := coordinator.run(coordinator.ctx, coordinator.cfg, RunPipelineWithLocalSourceOptions{
|
|
PipelineID: job.pipeline.ID,
|
|
SourceRoot: job.stagedRoot,
|
|
DryRun: job.request.DryRun,
|
|
Force: job.request.Force,
|
|
})
|
|
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
|
|
}
|
|
return DefaultUploadMaxFileCount
|
|
}
|
|
|
|
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
|
|
}
|