Skip to content

TM is fast enough in some cases that waiting on the main worker to finalize the results takes as long as the search. #564

Description

@bHimes

Background Finalization Thread for Template Matching to Unblock Workers

Problem Statement

The match_template program's master process currently blocks workers during result finalization. When a worker completes template matching for an image, the master process performs extensive I/O and computation before workers can proceed to the next image:

  • Writing 8+ MRC files per image (mip, psi, theta, phi, defocus, pixel_size, scaled_mip, sums, square_sums)
  • Writing histogram text files
  • Peak finding and projection generation
  • Sending results to GUI via socket

Impact: Workers sit idle during master's bookkeeping, wasting compute resources.

Location: src/programs/match_template/match_template.cpp:1692 (MasterHandleProgramDefinedResult)

Proposed Solution: Sequential Async Finalization

Implement a background finalization thread that processes completed results sequentially, allowing the master to immediately return to distributing work while finalization happens asynchronously.

┌─────────────────────────────┐
│ MasterHandleProgramDefined  │
│  1. Aggregate partial data  │ ← FAST (in-memory)
│  2. Enqueue if complete     │ ← FAST (queue push)
│  3. Return immediately      │ ← Workers continue!
└─────────────────────────────┘
            ↓
   ┌────────────────────┐
   │ Thread-Safe Queue  │
   └────────────────────┘
            ↓
┌─────────────────────────────┐
│ Background Thread           │
│  - Dequeue (FIFO order)     │
│  - Write MRC files          │ ← SLOW (disk I/O)
│  - Peak finding             │
│  - Send to GUI (w/ mutex)   │
│  - Repeat                   │
└─────────────────────────────┘

Implementation Details

1. Core Data Structures

// Container for one finalization job
struct FinalizationJob {
    AggregatedTemplateResult result;
    int image_number;

    // Copy of job parameters needed for finalization
    wxString input_reconstruction_filename;
    wxString mip_filename;
    wxString psi_filename;
    // ... all output filenames and parameters

    bool use_gpu;
    float min_peak_radius;
    float search_pixel_size;
    // ... other parameters
};

2. Thread-Safe Queue

class FinalizationQueue {
private:
    std::deque<FinalizationJob> jobs;
    std::mutex queue_mutex;
    std::condition_variable cv;
    bool shutdown_requested = false;

public:
    void Enqueue(FinalizationJob job);
    bool Dequeue(FinalizationJob& job);  // Blocks until work available
    void Shutdown();
    size_t Size() const;
};

3. Modified MatchTemplateApp Class

class MatchTemplateApp : public MyApp {
private:
    // NEW: Background finalization infrastructure
    FinalizationQueue finalization_queue;
    std::thread finalization_thread;
    std::mutex socket_mutex;  // CRITICAL: Protect GUI socket writes
    bool finalization_thread_running = false;

public:
    void StartFinalizationThread();
    void StopFinalizationThread();  // Waits for queue to drain
};

4. Modified Result Handling Flow

Current flow:

void MasterHandleProgramDefinedResult(...) {
    // 1. Aggregate partial results
    // 2. If complete: Write files + peak find + send to GUI  ← BLOCKS HERE
}

New flow:

void MasterHandleProgramDefinedResult(...) {
    // 1. Aggregate partial results (unchanged)

    // 2. If complete:
    if (all_parts_received) {
        FinalizationJob job;
        job.result = std::move(aggregated_results[i]);
        job.image_number = result_number;
        // Copy all needed parameters from current_job_package
        job.mip_filename = current_job_package.jobs[...].arguments[21]...;
        // ...

        finalization_queue.Enqueue(std::move(job));
        aggregated_results.RemoveAt(i);

        return;  // RETURN IMMEDIATELY - workers continue!
    }
}

5. Background Thread Worker

void FinalizationThreadWorker(FinalizationQueue* queue,
                              wxSocketBase* controller_socket,
                              std::mutex* socket_mutex) {
    FinalizationJob job;

    while (queue->Dequeue(job)) {  // Blocks until work available
        // Extract current finalization code (lines 1732-2138)
        // - Allocate images
        // - Write all MRC files
        // - Peak finding
        // - Generate projections

        // Send to GUI with mutex protection
        {
            std::lock_guard<std::mutex> lock(*socket_mutex);
            SendTemplateMatchingResultToSocket(
                controller_socket,
                job.image_number,
                expected_threshold,
                all_peak_infos,
                blank_changes
            );
        }
    }
}

Benefits

  1. Workers never wait: Master returns immediately after queueing finalization
  2. Preserves ordering: Sequential processing maintains expected template_match_id sequence
  3. Simple reasoning: FIFO queue ensures deterministic behavior
  4. Natural backpressure: Growing queue size indicates finalization bottleneck
  5. Clean shutdown: Thread drains all pending work before exit
  6. No data races: Sequential processing + socket mutex = thread-safe

Risks and Mitigations

Risk Mitigation
Socket corruption from concurrent writes Add std::mutex around all SendTemplateMatchingResultToSocket() calls
Memory growth if finalization can't keep up Queue size bounded by total images (finite)
Background thread errors not visible Wrap in try/catch, set error flag, check in main thread
Deadlock on shutdown Use timeout in condition_variable wait, proper shutdown signaling

Alternatives Considered

Option 1: Out-of-Order Finalization (Parallel)

  • Multiple background threads process any available result
  • Rejected: Creates non-sequential template_match_id values causing cosmetic issues with filenames and database display order

Option 2: Pre-Assign IDs + Parallel Finalization

  • Reserve block of template_match_id values upfront, allows parallel finalization
  • Rejected: More complex, requires schema changes, harder to reason about

Option 3: Sequential Finalization (Selected)

  • Single background thread processes in FIFO order
  • Selected: Simplest implementation, preserves ordering, achieves main goal (unblock workers)

Files to Modify

src/programs/match_template/match_template.cpp
  - Add FinalizationQueue class
  - Add FinalizationJob struct
  - Modify MatchTemplateApp to manage thread lifecycle
  - Split MasterHandleProgramDefinedResult into aggregate + finalize
  - Add FinalizationThreadWorker function
  - Add socket_mutex for thread safety

Testing Considerations

  1. Functional testing: Verify results identical to sequential version
  2. Performance testing: Measure worker idle time before/after
  3. Stress testing: Large image sets to test queue behavior
  4. Error testing: Simulate disk I/O failures, socket disconnects
  5. Shutdown testing: Verify clean exit with pending work

Success Metrics

  • Worker idle time reduced by ~XX% (measure master finalization time currently)
  • No functional regressions (results bit-identical to current implementation)
  • Clean shutdown (no lost results, no memory leaks)
  • No data races (ThreadSanitizer clean)

Related Code Patterns

This pattern already exists in other cisTEM workflows:

  • FindCTFPanel::ProcessAllJobsFinished() - similar result aggregation
  • Job controller socket communication - shows socket thread safety patterns

References

  • Database schema: src/core/database/database_schema.h:61 (TEMPLATE_MATCH_LIST)
  • Current bottleneck: match_template.cpp:1692-2138
  • GUI result handling: src/gui/MatchTemplatePanel.cpp:502 (HandleSocketTemplateMatchResultReady)

Metadata

Metadata

Assignees

No one assigned

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions