feat: queue extractions (#1227)

This commit is contained in:
Matteo Pietro Dazzi authored and GitHub committed 2026-01-05 15:53:19 +08:00
1 parent de0afb1481
commit 24db9bb58b
35 files changed
+1047 -159

No files matched your search

+118 -61
View File
@@ -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
+36 -27
View File
@@ -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)")
}
}
+199
View File
@@ -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
}
+545
View File
@@ -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")
}
}