mirror of
https://github.com/GopeedLab/gopeed.git
synced 2026-11-07 17:10:29 +00:00
feat: queue extractions (#1227)
This commit is contained in:
1 parent
de0afb1481
commit
24db9bb58b
35 files changed
+1047
-159
No files matched your search
+118
-61
@@ -52,6 +52,8 @@ type ExtractStatus string
|
||||
const (
|
||||
// ExtractStatusNone indicates extraction has not started
|
||||
ExtractStatusNone ExtractStatus = ""
|
||||
// ExtractStatusQueued indicates extraction is waiting in the queue
|
||||
ExtractStatusQueued ExtractStatus = "queued"
|
||||
// ExtractStatusWaitingParts indicates waiting for other multi-part archive parts to complete
|
||||
ExtractStatusWaitingParts ExtractStatus = "waitingParts"
|
||||
// ExtractStatusExtracting indicates extraction is in progress
|
||||
@@ -102,6 +104,10 @@ type Downloader struct {
|
||||
checkDuplicateLock *sync.Mutex
|
||||
closed atomic.Bool
|
||||
|
||||
// claimedExtractions tracks which multi-part archives have been claimed for extraction
|
||||
// Key: fullBaseName (e.g., "/path/archive.7z"), Value: taskID that claimed it
|
||||
claimedExtractions sync.Map
|
||||
|
||||
extensions []*Extension
|
||||
}
|
||||
|
||||
@@ -736,6 +742,9 @@ func (d *Downloader) emit(eventKey EventKey, task *Task, errs ...error) {
|
||||
}
|
||||
|
||||
func (d *Downloader) GetTask(id string) *Task {
|
||||
d.lock.Lock()
|
||||
defer d.lock.Unlock()
|
||||
|
||||
for _, task := range d.tasks {
|
||||
if task.ID == id {
|
||||
return task
|
||||
@@ -745,12 +754,18 @@ func (d *Downloader) GetTask(id string) *Task {
|
||||
}
|
||||
|
||||
func (d *Downloader) GetTasks() []*Task {
|
||||
d.lock.Lock()
|
||||
defer d.lock.Unlock()
|
||||
|
||||
return d.tasks
|
||||
}
|
||||
|
||||
// GetTasksByFilter get tasks by filter, if filter is nil, return all tasks
|
||||
// return tasks and if match all tasks
|
||||
func (d *Downloader) GetTasksByFilter(filter *TaskFilter) []*Task {
|
||||
d.lock.Lock()
|
||||
defer d.lock.Unlock()
|
||||
|
||||
if filter == nil || filter.IsEmpty() {
|
||||
return d.tasks
|
||||
}
|
||||
@@ -905,9 +920,10 @@ func (d *Downloader) watch(task *Task) {
|
||||
}()
|
||||
}
|
||||
|
||||
// Auto-extract archive files
|
||||
// Auto-extract archive files using the extraction queue
|
||||
// This ensures only one extraction runs at a time to prevent resource exhaustion
|
||||
if e.AutoExtract && isArchiveFile(downloadFilePath) {
|
||||
go d.handleAutoExtract(task, downloadFilePath, e)
|
||||
d.enqueueExtraction(task, downloadFilePath, e)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1200,55 +1216,93 @@ func (d *Downloader) buildFetcher(url string) (fetcher.Fetcher, error) {
|
||||
return fetcher, nil
|
||||
}
|
||||
|
||||
// handleAutoExtract handles automatic extraction of archive files, including multi-part archives
|
||||
func (d *Downloader) handleAutoExtract(task *Task, downloadFilePath string, opts *http.OptsExtra) {
|
||||
// Check if this is a multi-part archive
|
||||
// enqueueExtraction adds an extraction job to the global extraction queue
|
||||
// This ensures only one extraction (or one multi-part archive extraction) runs at a time
|
||||
// to prevent resource exhaustion
|
||||
func (d *Downloader) enqueueExtraction(task *Task, downloadFilePath string, opts *http.OptsExtra) {
|
||||
partInfo := getArchivePartInfo(downloadFilePath)
|
||||
|
||||
if partInfo.IsMultiPart {
|
||||
// Set multi-part info on the task
|
||||
task.Progress.MultiPartBaseName = partInfo.BaseName
|
||||
task.Progress.MultiPartNumber = partInfo.PartNumber
|
||||
task.Progress.MultiPartIsFirst = isFirstPart(downloadFilePath)
|
||||
d.storage.Put(bucketTask, task.ID, task.clone())
|
||||
// For multi-part archives, handle specially
|
||||
d.enqueueMultiPartExtraction(task, downloadFilePath, partInfo, opts)
|
||||
} else {
|
||||
// For single archives, queue immediately
|
||||
d.enqueueSingleExtraction(task, downloadFilePath, opts)
|
||||
}
|
||||
}
|
||||
|
||||
// Check if all parts are downloaded
|
||||
destDir := task.Meta.Opts.Path
|
||||
allPartsReady, missingParts := d.checkMultiPartArchiveReady(downloadFilePath, destDir, partInfo)
|
||||
// enqueueSingleExtraction queues extraction for a single (non-multi-part) archive
|
||||
func (d *Downloader) enqueueSingleExtraction(task *Task, downloadFilePath string, opts *http.OptsExtra) {
|
||||
jobID := "single:" + task.ID
|
||||
|
||||
if !allPartsReady {
|
||||
// Not all parts are ready yet, set status to waiting
|
||||
task.Progress.ExtractStatus = ExtractStatusWaitingParts
|
||||
d.emit(EventKeyProgress, task)
|
||||
d.storage.Put(bucketTask, task.ID, task.clone())
|
||||
d.Logger.Info().Msgf("multi-part archive waiting for other parts, task id: %s, missing: %v", task.ID, missingParts)
|
||||
return
|
||||
}
|
||||
// Set extraction status to queued
|
||||
task.Progress.ExtractStatus = ExtractStatusQueued
|
||||
d.emit(EventKeyProgress, task)
|
||||
d.storage.Put(bucketTask, task.ID, task.clone())
|
||||
d.Logger.Info().Msgf("extraction queued, task id: %s, job id: %s", task.ID, jobID)
|
||||
|
||||
// All parts are ready!
|
||||
// Only one task should perform extraction - we'll use the one that completes last
|
||||
// Check if extraction has already been started by another part
|
||||
if d.isMultiPartExtractionInProgress(partInfo.BaseName) {
|
||||
// Another part already started extraction, mark this as done
|
||||
task.Progress.ExtractStatus = ExtractStatusDone
|
||||
task.Progress.ExtractProgress = 100
|
||||
d.emit(EventKeyProgress, task)
|
||||
d.storage.Put(bucketTask, task.ID, task.clone())
|
||||
d.Logger.Info().Msgf("multi-part archive part completed, extraction handled by another part, task id: %s", task.ID)
|
||||
return
|
||||
}
|
||||
// Create and enqueue the extraction job
|
||||
job := NewExtractionJob(jobID, func() {
|
||||
d.performExtraction(task, downloadFilePath, task.Meta.Opts.Path, opts)
|
||||
})
|
||||
|
||||
// This task will handle extraction - mark as extracting first to prevent others
|
||||
task.Progress.ExtractStatus = ExtractStatusExtracting
|
||||
go func() {
|
||||
GetExtractionQueue().Enqueue(job)
|
||||
}()
|
||||
}
|
||||
|
||||
// enqueueMultiPartExtraction handles queueing for multi-part archives
|
||||
// It ensures only ONE extraction job is queued when ALL parts are ready
|
||||
func (d *Downloader) enqueueMultiPartExtraction(task *Task, downloadFilePath string, partInfo ArchivePartInfo, opts *http.OptsExtra) {
|
||||
// Set multi-part info on the task
|
||||
task.Progress.MultiPartBaseName = partInfo.BaseName
|
||||
task.Progress.MultiPartNumber = partInfo.PartNumber
|
||||
task.Progress.MultiPartIsFirst = isFirstPart(downloadFilePath)
|
||||
|
||||
// Check if all parts are downloaded
|
||||
destDir := task.Meta.Opts.Path
|
||||
allPartsReady, missingParts := d.checkMultiPartArchiveReady(downloadFilePath, destDir, partInfo)
|
||||
|
||||
if !allPartsReady {
|
||||
// Not all parts are ready yet - just set status to waiting, don't queue anything
|
||||
task.Progress.ExtractStatus = ExtractStatusWaitingParts
|
||||
d.emit(EventKeyProgress, task)
|
||||
d.storage.Put(bucketTask, task.ID, task.clone())
|
||||
|
||||
// Extract the multi-part archive
|
||||
d.performMultiPartExtraction(task, partInfo.FirstPartPath, destDir, opts)
|
||||
} else {
|
||||
// Regular (non-multi-part) archive
|
||||
d.performExtraction(task, downloadFilePath, task.Meta.Opts.Path, opts)
|
||||
d.Logger.Info().Msgf("multi-part archive waiting for other parts, task id: %s, missing: %v", task.ID, missingParts)
|
||||
return
|
||||
}
|
||||
|
||||
// All parts are ready! Atomically check if extraction has already been started/queued
|
||||
// and if not, mark this task as the one that will handle it
|
||||
// Use GetMultiPartArchiveBaseName to get the full path for comparison
|
||||
fullBaseName := GetMultiPartArchiveBaseName(downloadFilePath)
|
||||
shouldQueue := d.tryClaimMultiPartExtraction(task, fullBaseName)
|
||||
|
||||
if !shouldQueue {
|
||||
// Another part already started/queued extraction, mark this task as done
|
||||
task.Progress.ExtractStatus = ExtractStatusDone
|
||||
task.Progress.ExtractProgress = 100
|
||||
d.emit(EventKeyProgress, task)
|
||||
d.storage.Put(bucketTask, task.ID, task.clone())
|
||||
d.Logger.Info().Msgf("multi-part archive extraction already handled by another part, task id: %s", task.ID)
|
||||
return
|
||||
}
|
||||
|
||||
// This task claimed the extraction - status already set to queued in tryClaimMultiPartExtraction
|
||||
d.emit(EventKeyProgress, task)
|
||||
d.storage.Put(bucketTask, task.ID, task.clone())
|
||||
|
||||
jobID := "multipart:" + fullBaseName
|
||||
d.Logger.Info().Msgf("multi-part extraction queued, task id: %s, job id: %s", task.ID, jobID)
|
||||
|
||||
// Create and enqueue the extraction job
|
||||
job := NewExtractionJob(jobID, func() {
|
||||
d.performMultiPartExtraction(task, partInfo.FirstPartPath, destDir, opts)
|
||||
})
|
||||
|
||||
go func() {
|
||||
GetExtractionQueue().Enqueue(job)
|
||||
}()
|
||||
}
|
||||
|
||||
// checkMultiPartArchiveReady checks if all parts of a multi-part archive are downloaded
|
||||
@@ -1298,28 +1352,25 @@ func (d *Downloader) checkAllMultiPartTasksDone(baseName string) (bool, []string
|
||||
return len(notDoneParts) == 0, notDoneParts
|
||||
}
|
||||
|
||||
// isMultiPartExtractionInProgress checks if any task for this multi-part archive
|
||||
// has already started extraction (to prevent duplicate extractions)
|
||||
func (d *Downloader) isMultiPartExtractionInProgress(baseName string) bool {
|
||||
d.lock.Lock()
|
||||
defer d.lock.Unlock()
|
||||
|
||||
for _, task := range d.tasks {
|
||||
taskBaseName := ""
|
||||
if task.Meta != nil && task.Meta.Res != nil && len(task.Meta.Res.Files) > 0 {
|
||||
taskBaseName = GetMultiPartArchiveBaseName(task.Meta.SingleFilepath())
|
||||
}
|
||||
if taskBaseName == baseName {
|
||||
// Check if this task is already extracting or has completed extraction
|
||||
if task.Progress.ExtractStatus == ExtractStatusExtracting ||
|
||||
task.Progress.ExtractStatus == ExtractStatusDone ||
|
||||
task.Progress.ExtractStatus == ExtractStatusError {
|
||||
return true
|
||||
}
|
||||
}
|
||||
// tryClaimMultiPartExtraction atomically checks if extraction can be claimed for a multi-part archive
|
||||
// and if so, marks the task as queued. Returns true if this task should proceed with queueing.
|
||||
// This uses sync.Map.LoadOrStore for atomic claim to prevent race conditions.
|
||||
func (d *Downloader) tryClaimMultiPartExtraction(task *Task, baseName string) bool {
|
||||
// Use LoadOrStore for atomic claim - if another goroutine already stored a value, we get that value back
|
||||
_, alreadyClaimed := d.claimedExtractions.LoadOrStore(baseName, task.ID)
|
||||
if alreadyClaimed {
|
||||
return false // Another task already claimed it
|
||||
}
|
||||
|
||||
return false
|
||||
// This task successfully claimed it
|
||||
task.Progress.ExtractStatus = ExtractStatusQueued
|
||||
return true
|
||||
}
|
||||
|
||||
// releaseMultiPartExtractionClaim releases the extraction claim for a multi-part archive
|
||||
// This is primarily used for testing purposes
|
||||
func (d *Downloader) releaseMultiPartExtractionClaim(baseName string) {
|
||||
d.claimedExtractions.Delete(baseName)
|
||||
}
|
||||
|
||||
// performExtraction performs extraction for a regular (non-multi-part) archive
|
||||
@@ -1341,6 +1392,9 @@ func (d *Downloader) performExtraction(task *Task, archivePath string, destDir s
|
||||
|
||||
// performMultiPartExtraction performs extraction for a multi-part archive
|
||||
func (d *Downloader) performMultiPartExtraction(task *Task, firstPartPath string, destDir string, opts *http.OptsExtra) {
|
||||
// Get the baseName for releasing the claim later
|
||||
fullBaseName := GetMultiPartArchiveBaseName(firstPartPath)
|
||||
|
||||
// Set extraction status to extracting
|
||||
task.Progress.ExtractStatus = ExtractStatusExtracting
|
||||
task.Progress.ExtractProgress = 0
|
||||
@@ -1362,6 +1416,9 @@ func (d *Downloader) performMultiPartExtraction(task *Task, firstPartPath string
|
||||
|
||||
// Update status for all related multi-part tasks
|
||||
d.updateMultiPartTasksStatus(task, extractErr)
|
||||
|
||||
// Release the claim so future downloads of the same archive can be extracted
|
||||
d.releaseMultiPartExtractionClaim(fullBaseName)
|
||||
}
|
||||
|
||||
// collectMultiPartFiles collects all files belonging to a multi-part archive
|
||||
|
||||
@@ -1420,6 +1420,7 @@ func TestExtractStatus(t *testing.T) {
|
||||
expected string
|
||||
}{
|
||||
{ExtractStatusNone, ""},
|
||||
{ExtractStatusQueued, "queued"},
|
||||
{ExtractStatusExtracting, "extracting"},
|
||||
{ExtractStatusDone, "done"},
|
||||
{ExtractStatusError, "error"},
|
||||
@@ -2083,7 +2084,7 @@ func TestDownloader_CheckAllMultiPartTasksDone_NoRelatedTasks(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloader_IsMultiPartExtractionInProgress(t *testing.T) {
|
||||
func TestDownloader_TryClaimMultiPartExtraction(t *testing.T) {
|
||||
downloader := NewDownloader(nil)
|
||||
if err := downloader.Setup(); err != nil {
|
||||
t.Fatal(err)
|
||||
@@ -2097,50 +2098,58 @@ func TestDownloader_IsMultiPartExtractionInProgress(t *testing.T) {
|
||||
defer os.RemoveAll(tempDir)
|
||||
|
||||
baseName := "archive.7z"
|
||||
basePath := filepath.Join(tempDir, baseName)
|
||||
// GetMultiPartArchiveBaseName returns filepath.Join(dir, baseName)
|
||||
fullBaseName := filepath.Join(tempDir, baseName)
|
||||
|
||||
// Create task with no extraction status
|
||||
// Create tasks
|
||||
task1 := &Task{
|
||||
ID: "task1",
|
||||
Status: base.DownloadStatusDone,
|
||||
Meta: &fetcher.FetcherMeta{
|
||||
Opts: &base.Options{Path: tempDir},
|
||||
Res: &base.Resource{
|
||||
Files: []*base.FileInfo{{Name: baseName + ".001"}},
|
||||
Files: []*base.FileInfo{{Name: baseName + ".001", Path: ""}},
|
||||
},
|
||||
},
|
||||
Progress: &Progress{ExtractStatus: ExtractStatusNone},
|
||||
}
|
||||
initTask(task1)
|
||||
downloader.tasks = []*Task{task1}
|
||||
|
||||
// No extraction in progress
|
||||
if downloader.isMultiPartExtractionInProgress(basePath) {
|
||||
t.Error("isMultiPartExtractionInProgress() = true, want false")
|
||||
task2 := &Task{
|
||||
ID: "task2",
|
||||
Status: base.DownloadStatusDone,
|
||||
Meta: &fetcher.FetcherMeta{
|
||||
Opts: &base.Options{Path: tempDir},
|
||||
Res: &base.Resource{
|
||||
Files: []*base.FileInfo{{Name: baseName + ".002", Path: ""}},
|
||||
},
|
||||
},
|
||||
Progress: &Progress{ExtractStatus: ExtractStatusNone},
|
||||
}
|
||||
initTask(task2)
|
||||
|
||||
downloader.tasks = []*Task{task1, task2}
|
||||
|
||||
// Task1 should be able to claim extraction (no one has claimed yet)
|
||||
if !downloader.tryClaimMultiPartExtraction(task1, fullBaseName) {
|
||||
t.Error("tryClaimMultiPartExtraction() = false, want true (first claim)")
|
||||
}
|
||||
// task1's status should now be Queued
|
||||
if task1.Progress.ExtractStatus != ExtractStatusQueued {
|
||||
t.Errorf("task1.ExtractStatus = %v, want %v", task1.Progress.ExtractStatus, ExtractStatusQueued)
|
||||
}
|
||||
|
||||
// Set to extracting
|
||||
task1.Progress.ExtractStatus = ExtractStatusExtracting
|
||||
if !downloader.isMultiPartExtractionInProgress(basePath) {
|
||||
t.Error("isMultiPartExtractionInProgress() = false, want true (extracting)")
|
||||
// Task2 should NOT be able to claim (task1 already claimed via sync.Map)
|
||||
if downloader.tryClaimMultiPartExtraction(task2, fullBaseName) {
|
||||
t.Error("tryClaimMultiPartExtraction() = true, want false (already claimed)")
|
||||
}
|
||||
|
||||
// Set to done
|
||||
task1.Progress.ExtractStatus = ExtractStatusDone
|
||||
if !downloader.isMultiPartExtractionInProgress(basePath) {
|
||||
t.Error("isMultiPartExtractionInProgress() = false, want true (done)")
|
||||
}
|
||||
// Release the claim
|
||||
downloader.releaseMultiPartExtractionClaim(fullBaseName)
|
||||
|
||||
// Set to error
|
||||
task1.Progress.ExtractStatus = ExtractStatusError
|
||||
if !downloader.isMultiPartExtractionInProgress(basePath) {
|
||||
t.Error("isMultiPartExtractionInProgress() = false, want true (error)")
|
||||
}
|
||||
|
||||
// Set to waiting parts (not in progress)
|
||||
task1.Progress.ExtractStatus = ExtractStatusWaitingParts
|
||||
if downloader.isMultiPartExtractionInProgress(basePath) {
|
||||
t.Error("isMultiPartExtractionInProgress() = true, want false (waitingParts)")
|
||||
// Now task2 CAN claim (claim was released)
|
||||
if !downloader.tryClaimMultiPartExtraction(task2, fullBaseName) {
|
||||
t.Error("tryClaimMultiPartExtraction() = false, want true (claim was released)")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,545 @@
|
||||
package download
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestNewExtractionQueue(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
if q == nil {
|
||||
t.Fatal("NewExtractionQueue returned nil")
|
||||
}
|
||||
if q.jobs == nil {
|
||||
t.Fatal("jobs slice should be initialized")
|
||||
}
|
||||
if len(q.jobs) != 0 {
|
||||
t.Fatalf("expected empty jobs slice, got %d", len(q.jobs))
|
||||
}
|
||||
if q.running {
|
||||
t.Fatal("queue should not be running initially")
|
||||
}
|
||||
if q.cond == nil {
|
||||
t.Fatal("condition variable should be initialized")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_StartStop(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
|
||||
// Test Start
|
||||
q.Start()
|
||||
if !q.IsRunning() {
|
||||
t.Fatal("queue should be running after Start")
|
||||
}
|
||||
|
||||
// Test double Start (should be idempotent)
|
||||
q.Start()
|
||||
if !q.IsRunning() {
|
||||
t.Fatal("queue should still be running after second Start")
|
||||
}
|
||||
|
||||
// Test Stop
|
||||
q.Stop()
|
||||
if q.IsRunning() {
|
||||
t.Fatal("queue should not be running after Stop")
|
||||
}
|
||||
|
||||
// Test double Stop (should be idempotent)
|
||||
q.Stop()
|
||||
if q.IsRunning() {
|
||||
t.Fatal("queue should still not be running after second Stop")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_EnqueueSingleJob(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
executed := false
|
||||
job := NewExtractionJob("test-1", func() {
|
||||
executed = true
|
||||
})
|
||||
|
||||
q.Enqueue(job)
|
||||
job.Wait()
|
||||
|
||||
if !executed {
|
||||
t.Fatal("job should have been executed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_EnqueueAndWait(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
executed := false
|
||||
job := NewExtractionJob("test-1", func() {
|
||||
executed = true
|
||||
})
|
||||
|
||||
q.EnqueueAndWait(job)
|
||||
|
||||
if !executed {
|
||||
t.Fatal("job should have been executed after EnqueueAndWait returns")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_FIFOOrder(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
// Don't start yet - we want to queue multiple jobs first
|
||||
|
||||
var executionOrder []string
|
||||
var mu sync.Mutex
|
||||
|
||||
addToOrder := func(id string) {
|
||||
mu.Lock()
|
||||
executionOrder = append(executionOrder, id)
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
job1 := NewExtractionJob("job-1", func() {
|
||||
addToOrder("job-1")
|
||||
})
|
||||
job2 := NewExtractionJob("job-2", func() {
|
||||
addToOrder("job-2")
|
||||
})
|
||||
job3 := NewExtractionJob("job-3", func() {
|
||||
addToOrder("job-3")
|
||||
})
|
||||
|
||||
// Enqueue jobs before starting (they'll wait in queue)
|
||||
q.mu.Lock()
|
||||
q.jobs = append(q.jobs, job1, job2, job3)
|
||||
q.mu.Unlock()
|
||||
|
||||
// Now start processing
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
// Wait for all jobs
|
||||
job1.Wait()
|
||||
job2.Wait()
|
||||
job3.Wait()
|
||||
|
||||
// Verify FIFO order
|
||||
if len(executionOrder) != 3 {
|
||||
t.Fatalf("expected 3 executions, got %d", len(executionOrder))
|
||||
}
|
||||
if executionOrder[0] != "job-1" || executionOrder[1] != "job-2" || executionOrder[2] != "job-3" {
|
||||
t.Fatalf("expected FIFO order [job-1, job-2, job-3], got %v", executionOrder)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_SequentialExecution(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
var activeJobs int32
|
||||
var maxConcurrent int32
|
||||
var mu sync.Mutex
|
||||
|
||||
// Create jobs that take some time to execute
|
||||
createJob := func(id string) *ExtractionJob {
|
||||
return NewExtractionJob(id, func() {
|
||||
current := atomic.AddInt32(&activeJobs, 1)
|
||||
mu.Lock()
|
||||
if current > maxConcurrent {
|
||||
maxConcurrent = current
|
||||
}
|
||||
mu.Unlock()
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
atomic.AddInt32(&activeJobs, -1)
|
||||
})
|
||||
}
|
||||
|
||||
jobs := make([]*ExtractionJob, 5)
|
||||
for i := range jobs {
|
||||
jobs[i] = q.Enqueue(createJob(string(rune('A' + i))))
|
||||
}
|
||||
|
||||
// Wait for all jobs
|
||||
for _, job := range jobs {
|
||||
job.Wait()
|
||||
}
|
||||
|
||||
// Verify only one job ran at a time
|
||||
if maxConcurrent > 1 {
|
||||
t.Fatalf("expected max 1 concurrent job, got %d", maxConcurrent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_QueueLength(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
|
||||
if q.QueueLength() != 0 {
|
||||
t.Fatalf("expected queue length 0, got %d", q.QueueLength())
|
||||
}
|
||||
|
||||
// Add jobs without starting (they'll stay in queue)
|
||||
blockChan := make(chan struct{})
|
||||
blockingJob := NewExtractionJob("blocking", func() {
|
||||
<-blockChan // Block until signaled
|
||||
})
|
||||
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
q.Enqueue(blockingJob)
|
||||
|
||||
// Give worker time to pick up the blocking job
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
|
||||
// Queue more jobs while blocking job is running
|
||||
job2 := q.Enqueue(NewExtractionJob("job-2", func() {}))
|
||||
job3 := q.Enqueue(NewExtractionJob("job-3", func() {}))
|
||||
|
||||
// Should have 2 jobs waiting
|
||||
if q.QueueLength() != 2 {
|
||||
t.Fatalf("expected queue length 2, got %d", q.QueueLength())
|
||||
}
|
||||
|
||||
// Unblock and wait
|
||||
close(blockChan)
|
||||
blockingJob.Wait()
|
||||
job2.Wait()
|
||||
job3.Wait()
|
||||
|
||||
if q.QueueLength() != 0 {
|
||||
t.Fatalf("expected queue length 0 after completion, got %d", q.QueueLength())
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_HasPendingJob(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
|
||||
blockChan := make(chan struct{})
|
||||
blockingJob := NewExtractionJob("blocking", func() {
|
||||
<-blockChan
|
||||
})
|
||||
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
q.Enqueue(blockingJob)
|
||||
time.Sleep(20 * time.Millisecond) // Let blocking job start
|
||||
|
||||
// Queue more jobs
|
||||
job2 := q.Enqueue(NewExtractionJob("job-2", func() {}))
|
||||
|
||||
if !q.HasPendingJob("job-2") {
|
||||
t.Fatal("expected HasPendingJob to return true for job-2")
|
||||
}
|
||||
|
||||
if q.HasPendingJob("job-99") {
|
||||
t.Fatal("expected HasPendingJob to return false for non-existent job")
|
||||
}
|
||||
|
||||
// Unblock and wait
|
||||
close(blockChan)
|
||||
blockingJob.Wait()
|
||||
job2.Wait()
|
||||
|
||||
if q.HasPendingJob("job-2") {
|
||||
t.Fatal("expected HasPendingJob to return false after job completion")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_RemovePendingJob(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
|
||||
blockChan := make(chan struct{})
|
||||
blockingJob := NewExtractionJob("blocking", func() {
|
||||
<-blockChan
|
||||
})
|
||||
|
||||
executed := false
|
||||
job2 := NewExtractionJob("job-2", func() {
|
||||
executed = true
|
||||
})
|
||||
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
q.Enqueue(blockingJob)
|
||||
time.Sleep(20 * time.Millisecond) // Let blocking job start
|
||||
q.Enqueue(job2)
|
||||
|
||||
// Remove job-2 before it executes
|
||||
removed := q.RemovePendingJob("job-2")
|
||||
if !removed {
|
||||
t.Fatal("expected RemovePendingJob to return true")
|
||||
}
|
||||
|
||||
// Try to remove again (should return false)
|
||||
removed = q.RemovePendingJob("job-2")
|
||||
if removed {
|
||||
t.Fatal("expected second RemovePendingJob to return false")
|
||||
}
|
||||
|
||||
// Unblock
|
||||
close(blockChan)
|
||||
blockingJob.Wait()
|
||||
|
||||
// Wait a bit for any potential execution
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
if executed {
|
||||
t.Fatal("removed job should not have been executed")
|
||||
}
|
||||
|
||||
// The removed job's done channel should be closed
|
||||
select {
|
||||
case <-job2.done:
|
||||
// Expected
|
||||
default:
|
||||
t.Fatal("removed job's done channel should be closed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_StopDiscardsPendingJobs(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
|
||||
blockChan := make(chan struct{})
|
||||
executed1 := false
|
||||
executed2 := false
|
||||
|
||||
blockingJob := NewExtractionJob("blocking", func() {
|
||||
<-blockChan
|
||||
executed1 = true
|
||||
})
|
||||
pendingJob := NewExtractionJob("pending", func() {
|
||||
executed2 = true
|
||||
})
|
||||
|
||||
q.Start()
|
||||
|
||||
q.Enqueue(blockingJob)
|
||||
time.Sleep(20 * time.Millisecond) // Let blocking job start
|
||||
q.Enqueue(pendingJob)
|
||||
|
||||
// Stop without unblocking
|
||||
go func() {
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
close(blockChan)
|
||||
}()
|
||||
q.Stop()
|
||||
|
||||
// Blocking job should have completed, but pending job was discarded
|
||||
if !executed1 {
|
||||
t.Log("Note: blocking job may not complete if Stop() won races")
|
||||
}
|
||||
if executed2 {
|
||||
t.Fatal("pending job should not have been executed after Stop")
|
||||
}
|
||||
|
||||
// Pending job's done channel should be closed
|
||||
select {
|
||||
case <-pendingJob.done:
|
||||
// Expected
|
||||
default:
|
||||
t.Fatal("pending job's done channel should be closed after Stop")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_EnqueueAfterShutdown(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
q.Start()
|
||||
q.Stop()
|
||||
|
||||
executed := false
|
||||
job := NewExtractionJob("after-shutdown", func() {
|
||||
executed = true
|
||||
})
|
||||
|
||||
q.Enqueue(job)
|
||||
|
||||
// Job should be immediately done (without execution)
|
||||
select {
|
||||
case <-job.done:
|
||||
// Expected
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("job should be immediately done after shutdown")
|
||||
}
|
||||
|
||||
if executed {
|
||||
t.Fatal("job should not be executed after shutdown")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewExtractionJob(t *testing.T) {
|
||||
job := NewExtractionJob("test-id", func() {})
|
||||
|
||||
if job.ID != "test-id" {
|
||||
t.Fatalf("expected ID 'test-id', got '%s'", job.ID)
|
||||
}
|
||||
if job.Execute == nil {
|
||||
t.Fatal("Execute should not be nil")
|
||||
}
|
||||
if job.done == nil {
|
||||
t.Fatal("done channel should be initialized")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionJob_Wait(t *testing.T) {
|
||||
job := NewExtractionJob("test", func() {})
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
job.Wait()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
// Wait should block
|
||||
select {
|
||||
case <-done:
|
||||
t.Fatal("Wait should block until job is done")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
// Expected
|
||||
}
|
||||
|
||||
// Close done channel
|
||||
close(job.done)
|
||||
|
||||
// Now Wait should return
|
||||
select {
|
||||
case <-done:
|
||||
// Expected
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("Wait should return after job is done")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_NilExecuteFunction(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
// Job with nil Execute should not panic
|
||||
job := &ExtractionJob{
|
||||
ID: "nil-execute",
|
||||
Execute: nil,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
|
||||
q.Enqueue(job)
|
||||
job.Wait()
|
||||
// No panic means success
|
||||
}
|
||||
|
||||
func TestExtractionQueue_ConcurrentEnqueue(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
var counter int32
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 10
|
||||
jobsPerGoroutine := 10
|
||||
|
||||
wg.Add(numGoroutines)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func(goroutineID int) {
|
||||
defer wg.Done()
|
||||
for j := 0; j < jobsPerGoroutine; j++ {
|
||||
job := NewExtractionJob("", func() {
|
||||
atomic.AddInt32(&counter, 1)
|
||||
})
|
||||
q.EnqueueAndWait(job)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
expected := int32(numGoroutines * jobsPerGoroutine)
|
||||
if counter != expected {
|
||||
t.Fatalf("expected counter %d, got %d", expected, counter)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_Restart(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
|
||||
// First run
|
||||
q.Start()
|
||||
executed1 := false
|
||||
job1 := NewExtractionJob("job-1", func() {
|
||||
executed1 = true
|
||||
})
|
||||
q.EnqueueAndWait(job1)
|
||||
q.Stop()
|
||||
|
||||
if !executed1 {
|
||||
t.Fatal("first job should have been executed")
|
||||
}
|
||||
|
||||
// Restart
|
||||
q.Start()
|
||||
executed2 := false
|
||||
job2 := NewExtractionJob("job-2", func() {
|
||||
executed2 = true
|
||||
})
|
||||
q.EnqueueAndWait(job2)
|
||||
q.Stop()
|
||||
|
||||
if !executed2 {
|
||||
t.Fatal("second job should have been executed after restart")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetExtractionQueue(t *testing.T) {
|
||||
// Note: This test modifies global state, so it should ideally be run in isolation
|
||||
// However, the implementation ensures the queue is created only once
|
||||
|
||||
queue := GetExtractionQueue()
|
||||
if queue == nil {
|
||||
t.Fatal("GetExtractionQueue should not return nil")
|
||||
}
|
||||
if !queue.IsRunning() {
|
||||
t.Fatal("global queue should be running")
|
||||
}
|
||||
|
||||
// Calling again should return the same instance
|
||||
queue2 := GetExtractionQueue()
|
||||
if queue != queue2 {
|
||||
t.Fatal("GetExtractionQueue should return the same instance")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionQueue_LongRunningJob(t *testing.T) {
|
||||
q := NewExtractionQueue()
|
||||
q.Start()
|
||||
defer q.Stop()
|
||||
|
||||
startTime := time.Now()
|
||||
longJobDuration := 100 * time.Millisecond
|
||||
|
||||
longJob := NewExtractionJob("long-job", func() {
|
||||
time.Sleep(longJobDuration)
|
||||
})
|
||||
|
||||
shortJobExecuted := false
|
||||
shortJob := NewExtractionJob("short-job", func() {
|
||||
shortJobExecuted = true
|
||||
})
|
||||
|
||||
q.Enqueue(longJob)
|
||||
q.Enqueue(shortJob)
|
||||
|
||||
shortJob.Wait()
|
||||
|
||||
elapsed := time.Since(startTime)
|
||||
if elapsed < longJobDuration {
|
||||
t.Fatalf("short job should wait for long job, elapsed: %v", elapsed)
|
||||
}
|
||||
|
||||
if !shortJobExecuted {
|
||||
t.Fatal("short job should have been executed")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user