refactor: implement worker pool pattern, add persistent task queue, and decouple target loading logic
This commit is contained in:
@@ -0,0 +1,306 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/hireme/hireme/internal/analyzer"
|
||||
"github.com/hireme/hireme/internal/crawler"
|
||||
"github.com/hireme/hireme/internal/db"
|
||||
"github.com/hireme/hireme/internal/filter"
|
||||
"github.com/hireme/hireme/pkg/crawler/google"
|
||||
)
|
||||
|
||||
type WorkerPool struct {
|
||||
db *db.DB
|
||||
analyzer analyzer.LLMAnalyzer
|
||||
filterPipe *filter.FilterPipeline
|
||||
isRunning atomic.Bool
|
||||
workerCount int
|
||||
wg sync.WaitGroup
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
}
|
||||
|
||||
func NewWorkerPool(database *db.DB, analyzer analyzer.LLMAnalyzer, numWorkers int) *WorkerPool {
|
||||
// Initialize default filters
|
||||
hardKeywords := []string{"cleared", "clearance", "polygraph", "ts/sci", "secret", "citizen", "citizenship", "phd", "ph.d"}
|
||||
scoringWeights := []filter.KeywordWeight{
|
||||
{Keyword: "golang", Weight: 3.0}, {Keyword: "go", Weight: 2.0},
|
||||
{Keyword: "python", Weight: 2.0}, {Keyword: "kubernetes", Weight: 3.0},
|
||||
{Keyword: "k8s", Weight: 3.0}, {Keyword: "docker", Weight: 2.0},
|
||||
{Keyword: "aws", Weight: 1.0}, {Keyword: "gcp", Weight: 1.0},
|
||||
{Keyword: "distributed systems", Weight: 4.0},
|
||||
{Keyword: "backend", Weight: 3.0}, {Keyword: "ai", Weight: -5.0},
|
||||
{Keyword: "machine learning", Weight: -5.0},
|
||||
}
|
||||
|
||||
filterPipe := filter.NewPipeline(
|
||||
filter.NewHardKeywordFilter(hardKeywords),
|
||||
filter.NewScoringFilter(scoringWeights),
|
||||
)
|
||||
|
||||
return &WorkerPool{
|
||||
db: database,
|
||||
analyzer: analyzer,
|
||||
filterPipe: filterPipe,
|
||||
workerCount: numWorkers,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *WorkerPool) IsRunning() bool {
|
||||
return p.isRunning.Load()
|
||||
}
|
||||
|
||||
func (p *WorkerPool) Toggle() bool {
|
||||
if p.IsRunning() {
|
||||
p.Stop()
|
||||
return false
|
||||
}
|
||||
p.Start()
|
||||
return true
|
||||
}
|
||||
|
||||
func (p *WorkerPool) Start() {
|
||||
if p.isRunning.Load() {
|
||||
return
|
||||
}
|
||||
p.ctx, p.cancel = context.WithCancel(context.Background())
|
||||
p.isRunning.Store(true)
|
||||
|
||||
for i := 0; i < p.workerCount; i++ {
|
||||
p.wg.Add(1)
|
||||
go p.workerLoop(i)
|
||||
}
|
||||
log.Printf("[WorkerPool] Started %d workers", p.workerCount)
|
||||
}
|
||||
|
||||
func (p *WorkerPool) Stop() {
|
||||
if !p.isRunning.Load() {
|
||||
return
|
||||
}
|
||||
p.cancel()
|
||||
p.isRunning.Store(false)
|
||||
p.wg.Wait()
|
||||
log.Println("[WorkerPool] All workers stopped")
|
||||
}
|
||||
|
||||
func (p *WorkerPool) workerLoop(workerID int) {
|
||||
defer p.wg.Done()
|
||||
log.Printf("[Worker %d] Starting loop", workerID)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-p.ctx.Done():
|
||||
log.Printf("[Worker %d] Context cancelled, exiting", workerID)
|
||||
return
|
||||
default:
|
||||
task, err := p.db.ClaimNextTask()
|
||||
if err != nil {
|
||||
log.Printf("[Worker %d] Error claiming task: %v", workerID, err)
|
||||
time.Sleep(2 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
if task == nil {
|
||||
// No tasks available, sleep and poll again
|
||||
time.Sleep(1 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
log.Printf("[Worker %d] Picked up task %d (%s)", workerID, task.ID, task.TaskType)
|
||||
|
||||
// Process task based on type
|
||||
var processErr error
|
||||
switch task.TaskType {
|
||||
case "DISCOVER":
|
||||
processErr = p.handleDiscover(p.ctx, task)
|
||||
case "FETCH":
|
||||
processErr = p.handleFetch(p.ctx, task)
|
||||
case "PROCESS":
|
||||
processErr = p.handleProcess(p.ctx, task)
|
||||
case "ANALYZE":
|
||||
processErr = p.handleAnalyze(p.ctx, task)
|
||||
default:
|
||||
processErr = fmt.Errorf("unknown task type: %s", task.TaskType)
|
||||
}
|
||||
|
||||
if processErr != nil {
|
||||
log.Printf("[Worker %d] Task %d FAILED: %v", workerID, task.ID, processErr)
|
||||
_ = p.db.FailTask(task.ID, processErr.Error())
|
||||
} else {
|
||||
log.Printf("[Worker %d] Task %d COMPLETED", workerID, task.ID)
|
||||
_ = p.db.CompleteTask(task.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------
|
||||
// Task Payloads
|
||||
// ---------------------------------------------------------
|
||||
|
||||
type DiscoverPayload struct {
|
||||
TargetID int64 `json:"target_id"`
|
||||
TargetURL string `json:"target_url"`
|
||||
Platform string `json:"platform"`
|
||||
}
|
||||
|
||||
type FetchPayload struct {
|
||||
TargetID int64 `json:"target_id"`
|
||||
Platform string `json:"platform"`
|
||||
ApplyURL string `json:"apply_url"`
|
||||
CanonicalID string `json:"canonical_id"`
|
||||
}
|
||||
|
||||
type ProcessPayload struct {
|
||||
RawJobID int64 `json:"raw_job_id"`
|
||||
}
|
||||
|
||||
type AnalyzePayload struct {
|
||||
RawJobID int64 `json:"raw_job_id"`
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------
|
||||
// Task Handlers
|
||||
// ---------------------------------------------------------
|
||||
|
||||
func (p *WorkerPool) handleDiscover(ctx context.Context, task *db.Task) error {
|
||||
var payload DiscoverPayload
|
||||
if err := json.Unmarshal([]byte(task.Payload), &payload); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
c := p.getCrawlerForPlatform(payload.Platform)
|
||||
discovered, err := c.DiscoverJobs(ctx, payload.TargetURL)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, disc := range discovered {
|
||||
// Check deduplication (already fetched/inserted jobs)
|
||||
exists, err := p.db.CanonicalIDExists(disc.CanonicalID)
|
||||
if err == nil && exists {
|
||||
continue // Skip existing job
|
||||
}
|
||||
|
||||
// Check if it's already in the task queue to be fetched
|
||||
pending, err := p.db.HasPendingFetchTask(disc.CanonicalID)
|
||||
if err == nil && pending {
|
||||
continue // Skip duplicate pending fetch
|
||||
}
|
||||
|
||||
// Enqueue FETCH task
|
||||
fetchPayload := FetchPayload{
|
||||
TargetID: payload.TargetID,
|
||||
Platform: payload.Platform,
|
||||
ApplyURL: disc.ApplyURL,
|
||||
CanonicalID: disc.CanonicalID,
|
||||
}
|
||||
_ = p.db.EnqueueTask("FETCH", fetchPayload)
|
||||
}
|
||||
|
||||
// Update the target's last crawled time
|
||||
_ = p.db.UpdateTargetCrawledAt(payload.TargetID)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *WorkerPool) handleFetch(ctx context.Context, task *db.Task) error {
|
||||
var payload FetchPayload
|
||||
if err := json.Unmarshal([]byte(task.Payload), &payload); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
c := p.getCrawlerForPlatform(payload.Platform)
|
||||
extracted, err := c.FetchJobDetail(ctx, payload.ApplyURL)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
tID := payload.TargetID
|
||||
rawJob, inserted, err := p.db.InsertRawJob(extracted.CanonicalID, &tID, extracted.Title, extracted.Company, extracted.Location, extracted.ApplyURL, extracted.RawContent)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if inserted && rawJob != nil {
|
||||
processPayload := ProcessPayload{RawJobID: rawJob.ID}
|
||||
_ = p.db.EnqueueTask("PROCESS", processPayload)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *WorkerPool) handleProcess(ctx context.Context, task *db.Task) error {
|
||||
var payload ProcessPayload
|
||||
if err := json.Unmarshal([]byte(task.Payload), &payload); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
job, err := p.db.GetJobByID(payload.RawJobID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
keep, score, reason := p.filterPipe.Execute(ctx, job.Title, job.RawContent)
|
||||
if !keep {
|
||||
// Create a rejected JobAnalysis record. Store reason in missing_highlights JSON array
|
||||
reasonJSON, _ := json.Marshal([]string{reason})
|
||||
_ = p.db.CreateJobAnalysisFull(payload.RawJobID, score, false, "[]", string(reasonJSON), "[]", "PARSED")
|
||||
return nil
|
||||
}
|
||||
|
||||
// It's qualified for analysis. Create the base record and enqueue it.
|
||||
_ = p.db.CreateJobAnalysis(payload.RawJobID, score, true, "PENDING")
|
||||
|
||||
analyzePayload := AnalyzePayload(payload)
|
||||
_ = p.db.EnqueueTask("ANALYZE", analyzePayload)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *WorkerPool) handleAnalyze(ctx context.Context, task *db.Task) error {
|
||||
var payload AnalyzePayload
|
||||
if err := json.Unmarshal([]byte(task.Payload), &payload); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
job, err := p.db.GetJobByID(payload.RawJobID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Use empty matched/missing keywords for now, or pass from filter if stored
|
||||
result, err := p.analyzer.AnalyzeJD(ctx, job.Title, job.RawContent)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
matchedStr, _ := json.Marshal(result.MatchedHighlights)
|
||||
missingStr, _ := json.Marshal(result.MissingHighlights)
|
||||
tailoringStr, _ := json.Marshal(result.TailoringSuggestions)
|
||||
|
||||
err = p.db.UpdateJobAnalysisLLM(
|
||||
payload.RawJobID,
|
||||
string(matchedStr),
|
||||
string(missingStr),
|
||||
string(tailoringStr),
|
||||
)
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
func (p *WorkerPool) getCrawlerForPlatform(platform string) crawler.SiteCrawler {
|
||||
plat := strings.ToLower(platform)
|
||||
if strings.Contains(plat, "google") {
|
||||
return google.NewGoogleCrawler()
|
||||
}
|
||||
return crawler.NewGenericCrawler()
|
||||
}
|
||||
Reference in New Issue
Block a user