feat: queue extractions (#1227)

pkg/download/extract_queue.go

@@ -0,0 +1,199 @@
+package download
+
+import (
+	"sync"
+)
+
+// ExtractionJob represents a single extraction job in the queue
+type ExtractionJob struct {
+	// ID is a unique identifier for this job (usually task ID or multi-part base name)
+	ID string
+	// Execute is the function that performs the actual extraction
+	Execute func()
+	// done channel is signaled when the job has been executed
+	done chan struct{}
+}
+
+// NewExtractionJob creates a new extraction job
+func NewExtractionJob(id string, execute func()) *ExtractionJob {
+	return &ExtractionJob{
+		ID:      id,
+		Execute: execute,
+		done:    make(chan struct{}),
+	}
+}
+
+// Wait blocks until the job has been executed
+func (j *ExtractionJob) Wait() {
+	<-j.done
+}
+
+// ExtractionQueue manages a queue of extraction jobs to prevent resource exhaustion
+// by ensuring only one extraction (or one multi-part archive extraction) runs at a time
+type ExtractionQueue struct {
+	mu       sync.Mutex
+	cond     *sync.Cond
+	jobs     []*ExtractionJob
+	running  bool
+	shutdown bool
+	wg       sync.WaitGroup
+}
+
+// NewExtractionQueue creates a new extraction queue
+func NewExtractionQueue() *ExtractionQueue {
+	q := &ExtractionQueue{
+		jobs: make([]*ExtractionJob, 0),
+	}
+	q.cond = sync.NewCond(&q.mu)
+	return q
+}
+
+// Start starts the queue worker that processes jobs sequentially
+func (q *ExtractionQueue) Start() {
+	q.mu.Lock()
+	if q.running {
+		q.mu.Unlock()
+		return
+	}
+	q.running = true
+	q.shutdown = false
+	q.mu.Unlock()
+
+	q.wg.Add(1)
+	go q.worker()
+}
+
+// Stop stops the queue worker gracefully
+// It waits for the current job to complete but discards pending jobs
+func (q *ExtractionQueue) Stop() {
+	q.mu.Lock()
+	if !q.running {
+		q.mu.Unlock()
+		return
+	}
+	q.shutdown = true
+	q.cond.Signal()
+	q.mu.Unlock()
+
+	q.wg.Wait()
+}
+
+// Enqueue adds a new extraction job to the queue
+// The job will be executed when its turn comes (FIFO order)
+// Returns the job so the caller can wait for completion if needed
+func (q *ExtractionQueue) Enqueue(job *ExtractionJob) *ExtractionJob {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	if q.shutdown {
+		// Queue is shutting down, signal done immediately without executing
+		close(job.done)
+		return job
+	}
+
+	q.jobs = append(q.jobs, job)
+	q.cond.Signal()
+	return job
+}
+
+// EnqueueAndWait adds a new extraction job to the queue and waits for it to complete
+func (q *ExtractionQueue) EnqueueAndWait(job *ExtractionJob) {
+	q.Enqueue(job)
+	job.Wait()
+}
+
+// QueueLength returns the current number of pending jobs in the queue
+func (q *ExtractionQueue) QueueLength() int {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	return len(q.jobs)
+}
+
+// IsRunning returns true if the queue worker is running
+func (q *ExtractionQueue) IsRunning() bool {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	return q.running
+}
+
+// HasPendingJob checks if there's a pending job with the given ID
+func (q *ExtractionQueue) HasPendingJob(id string) bool {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	for _, job := range q.jobs {
+		if job.ID == id {
+			return true
+		}
+	}
+	return false
+}
+
+// RemovePendingJob removes a pending job with the given ID from the queue
+// Returns true if a job was removed, false if not found
+// Note: This cannot remove a job that is currently being executed
+func (q *ExtractionQueue) RemovePendingJob(id string) bool {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	for i, job := range q.jobs {
+		if job.ID == id {
+			// Close the done channel to unblock any waiters
+			close(job.done)
+			// Remove from queue
+			q.jobs = append(q.jobs[:i], q.jobs[i+1:]...)
+			return true
+		}
+	}
+	return false
+}
+
+// worker is the main loop that processes jobs sequentially
+func (q *ExtractionQueue) worker() {
+	defer q.wg.Done()
+
+	for {
+		q.mu.Lock()
+		// Wait for a job or shutdown signal
+		for len(q.jobs) == 0 && !q.shutdown {
+			q.cond.Wait()
+		}
+
+		if q.shutdown {
+			// Close done channels for all remaining jobs
+			for _, job := range q.jobs {
+				close(job.done)
+			}
+			q.jobs = nil
+			q.running = false
+			q.mu.Unlock()
+			return
+		}
+
+		// Dequeue the first job
+		job := q.jobs[0]
+		q.jobs = q.jobs[1:]
+		q.mu.Unlock()
+
+		// Execute the job outside the lock
+		if job.Execute != nil {
+			job.Execute()
+		}
+
+		// Signal that the job is done
+		close(job.done)
+	}
+}
+
+// Global extraction queue instance
+var globalExtractionQueue *ExtractionQueue
+var extractionQueueOnce sync.Once
+
+// GetExtractionQueue returns the global extraction queue instance
+func GetExtractionQueue() *ExtractionQueue {
+	extractionQueueOnce.Do(func() {
+		globalExtractionQueue = NewExtractionQueue()
+		globalExtractionQueue.Start()
+	})
+	return globalExtractionQueue
+}