307 lines
7.8 KiB
Go
307 lines
7.8 KiB
Go
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()
|
|
}
|