yao/job/worker.go

package job

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"sync"
	"time"

	"github.com/google/uuid"
	jsoniter "github.com/json-iterator/go"
	"github.com/yaoapp/kun/log"
)

// WorkerManager manages job execution workers
type WorkerManager struct {
	maxWorkers    int
	activeWorkers map[string]*Worker
	workQueue     chan *WorkRequest
	workerPool    chan chan *WorkRequest
	quit          chan bool
	mu            sync.RWMutex
}

// Worker represents a single worker instance
type Worker struct {
	ID         string
	WorkerPool chan chan *WorkRequest
	JobChannel chan *WorkRequest
	Quit       chan bool
	Mode       ModeType
	ctx        context.Context
	cancel     context.CancelFunc
}

// WorkRequest represents a job execution request
type WorkRequest struct {
	Job       *Job
	Execution *Execution
	Handler   HandlerFunc
	Context   context.Context
}

// Global worker manager instance
var globalWorkerManager *WorkerManager
var workerManagerOnce sync.Once

// GetWorkerManager returns the global worker manager instance
func GetWorkerManager() *WorkerManager {
	workerManagerOnce.Do(func() {
		globalWorkerManager = NewWorkerManager(runtime.NumCPU() * 2) // Default to 2x CPU cores
	})
	return globalWorkerManager
}

// NewWorkerManagerForTest creates a new worker manager for testing (not singleton)
func NewWorkerManagerForTest(maxWorkers int) *WorkerManager {
	return NewWorkerManager(maxWorkers)
}

// NewWorkerManager creates a new worker manager
func NewWorkerManager(maxWorkers int) *WorkerManager {
	return &WorkerManager{
		maxWorkers:    maxWorkers,
		activeWorkers: make(map[string]*Worker),
		workQueue:     make(chan *WorkRequest, maxWorkers*2), // Buffer for queue
		workerPool:    make(chan chan *WorkRequest, maxWorkers),
		quit:          make(chan bool),
	}
}

// Start starts the worker manager
func (wm *WorkerManager) Start() {
	// Start workers
	for i := 0; i < wm.maxWorkers; i++ {
		worker := NewWorker(wm.workerPool, GOROUTINE)
		worker.Start()

		wm.mu.Lock()
		wm.activeWorkers[worker.ID] = worker
		wm.mu.Unlock()
	}

	// Start dispatcher
	go wm.dispatch()
	log.Info("Worker manager started with %d workers", wm.maxWorkers)
}

// Stop stops the worker manager
func (wm *WorkerManager) Stop() {
	log.Info("Stopping worker manager...")

	// Stop dispatcher first
	select {
	case <-wm.quit:
		// Already stopped
		return
	default:
		close(wm.quit)
	}

	// Stop all workers
	wm.mu.Lock()
	workers := make([]*Worker, 0, len(wm.activeWorkers))
	for _, worker := range wm.activeWorkers {
		workers = append(workers, worker)
	}
	wm.activeWorkers = make(map[string]*Worker)
	wm.mu.Unlock()

	// Stop workers and wait for them to finish
	for _, worker := range workers {
		worker.Stop()
	}

	// Give workers time to finish their current operations
	time.Sleep(200 * time.Millisecond)

	log.Info("Worker manager stopped")
}

// SubmitJob submits a job for execution
func (wm *WorkerManager) SubmitJob(job *Job, handler HandlerFunc) error {
	// Create execution record
	execution := &Execution{
		ExecutionID:     uuid.New().String(),
		JobID:           job.JobID,
		Status:          "queued",
		TriggerCategory: "manual", // Default trigger
		RetryAttempt:    0,
		Progress:        0,
		CreatedAt:       time.Now(),
		UpdatedAt:       time.Now(),
	}

	// Save execution to database
	if err := SaveExecution(execution); err != nil {
		return fmt.Errorf("failed to save execution: %w", err)
	}

	// Create work request
	workRequest := &WorkRequest{
		Job:       job,
		Execution: execution,
		Handler:   handler,
		Context:   context.Background(),
	}

	// Submit to work queue
	select {
	case wm.workQueue <- workRequest:
		log.Debug("Job %s submitted to work queue", job.JobID)
		return nil
	default:
		return fmt.Errorf("work queue is full")
	}
}

// dispatch dispatches work requests to available workers
func (wm *WorkerManager) dispatch() {
	for {
		select {
		case work := <-wm.workQueue:
			// Get an available worker
			select {
			case jobChannel := <-wm.workerPool:
				// Send work to worker
				jobChannel <- work
			case <-wm.quit:
				return
			}
		case <-wm.quit:
			return
		}
	}
}

// GetActiveWorkers returns the number of active workers
func (wm *WorkerManager) GetActiveWorkers() int {
	wm.mu.RLock()
	defer wm.mu.RUnlock()
	return len(wm.activeWorkers)
}

// NewWorker creates a new worker
func NewWorker(workerPool chan chan *WorkRequest, mode ModeType) *Worker {
	ctx, cancel := context.WithCancel(context.Background())

	return &Worker{
		ID:         uuid.New().String(),
		WorkerPool: workerPool,
		JobChannel: make(chan *WorkRequest),
		Quit:       make(chan bool),
		Mode:       mode,
		ctx:        ctx,
		cancel:     cancel,
	}
}

// Start starts the worker
func (w *Worker) Start() {
	go func() {
		for {
			// Register worker in the worker pool
			w.WorkerPool <- w.JobChannel

			select {
			case work := <-w.JobChannel:
				// Process the work
				w.processWork(work)
			case <-w.Quit:
				return
			}
		}
	}()
}

// Stop stops the worker
func (w *Worker) Stop() {
	w.cancel()
	select {
	case <-w.Quit:
		// Already stopped
		return
	default:
		close(w.Quit)
	}
}

// processWork processes a work request
func (w *Worker) processWork(work *WorkRequest) {
	log.Debug("Worker %s processing job %s", w.ID, work.Job.JobID)

	// Update execution status
	work.Execution.Status = "running"
	work.Execution.WorkerID = &w.ID
	work.Execution.StartedAt = &time.Time{}
	*work.Execution.StartedAt = time.Now()

	// Try to save execution, but don't fail if database is closed
	if err := SaveExecution(work.Execution); err != nil {
		log.Warn("Failed to save execution status (database may be closed): %v", err)
	}

	// Update job status
	work.Job.Status = "running"
	work.Job.CurrentExecutionID = &work.Execution.ExecutionID
	work.Job.LastRunAt = work.Execution.StartedAt

	// Try to save job, but don't fail if database is closed
	if err := SaveJob(work.Job); err != nil {
		log.Warn("Failed to save job status (database may be closed): %v", err)
	}

	// Create execution context with timeout
	ctx := work.Context
	if work.Job.DefaultTimeout != nil && *work.Job.DefaultTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(work.Context, time.Duration(*work.Job.DefaultTimeout)*time.Second)
		defer cancel()
	}

	// Set up progress tracking
	progress := &Progress{
		ExecutionID: work.Execution.ExecutionID,
		Progress:    0,
		Message:     "Starting execution",
	}

	// Update execution with progress manager
	work.Execution.Job = work.Job // Set job reference for progress updates

	var err error
	startTime := time.Now()

	// Execute based on mode
	switch w.Mode {
	case GOROUTINE:
		err = w.executeInGoroutine(ctx, work, progress)
	case PROCESS:
		err = w.executeInProcess(ctx, work, progress)
	default:
		err = fmt.Errorf("unsupported execution mode: %s", w.Mode)
	}

	// Calculate duration
	duration := int(time.Since(startTime).Milliseconds())
	endTime := time.Now()

	// Update execution with results
	work.Execution.EndedAt = &endTime
	work.Execution.Duration = &duration

	if err != nil {
		work.Execution.Status = "failed"
		errorInfo := map[string]interface{}{
			"error":  err.Error(),
			"time":   endTime,
			"worker": w.ID,
		}
		errorData, _ := jsoniter.Marshal(errorInfo)
		work.Execution.ErrorInfo = (*json.RawMessage)(&errorData)

		log.Error("Job %s execution failed: %v", work.Job.JobID, err)

		// Log error
		logEntry := &Log{
			JobID:       work.Job.JobID,
			Level:       "error",
			Message:     fmt.Sprintf("Execution failed: %v", err),
			ExecutionID: &work.Execution.ExecutionID,
			WorkerID:    &w.ID,
			Timestamp:   time.Now(),
			Sequence:    0,
		}
		if err := SaveLog(logEntry); err != nil {
			log.Warn("Failed to save error log (database may be closed): %v", err)
		}
	} else {
		work.Execution.Status = "completed"
		work.Execution.Progress = 100

		log.Info("Job %s execution completed successfully", work.Job.JobID)

		// Log completion
		logEntry := &Log{
			JobID:       work.Job.JobID,
			Level:       "info",
			Message:     "Execution completed successfully",
			ExecutionID: &work.Execution.ExecutionID,
			WorkerID:    &w.ID,
			Progress:    &work.Execution.Progress,
			Duration:    &duration,
			Timestamp:   time.Now(),
			Sequence:    1,
		}
		if err := SaveLog(logEntry); err != nil {
			log.Warn("Failed to save completion log (database may be closed): %v", err)
		}
	}

	// Update execution in database
	if err := SaveExecution(work.Execution); err != nil {
		log.Warn("Failed to save final execution status (database may be closed): %v", err)
	}

	// Update job status
	if work.Job.ScheduleType == string(ScheduleTypeOnce) {
		work.Job.Status = "completed"
	} else {
		work.Job.Status = "ready" // Ready for next execution
	}
	work.Job.CurrentExecutionID = nil
	if err := SaveJob(work.Job); err != nil {
		log.Warn("Failed to save final job status (database may be closed): %v", err)
	}

	log.Debug("Worker %s finished processing job %s", w.ID, work.Job.JobID)
}

// executeInGoroutine executes job in goroutine mode
func (w *Worker) executeInGoroutine(ctx context.Context, work *WorkRequest, progress *Progress) error {
	// Execute handler directly in current goroutine
	return work.Handler(ctx, work.Execution)
}

// executeInProcess executes job in process mode
func (w *Worker) executeInProcess(ctx context.Context, work *WorkRequest, progress *Progress) error {
	// For process mode, we would typically spawn a separate process
	// For now, we'll simulate this with a goroutine but with process isolation concepts

	// Set process ID (simulated)
	processID := fmt.Sprintf("proc_%s", uuid.New().String()[:8])
	work.Execution.ProcessID = &processID

	// Create a separate goroutine to simulate process isolation
	done := make(chan error, 1)

	go func() {
		defer func() {
			if r := recover(); r != nil {
				done <- fmt.Errorf("process panic: %v", r)
			}
		}()

		// Execute handler
		err := work.Handler(ctx, work.Execution)
		done <- err
	}()

	// Wait for completion or timeout
	select {
	case err := <-done:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// executeInRealProcess executes job in a real separate process (future implementation)
func (w *Worker) executeInRealProcess(ctx context.Context, work *WorkRequest, progress *Progress) error {
	// This would be used for true process isolation
	// For now, it's a placeholder for future implementation

	// Create command to execute job in separate process
	cmd := exec.CommandContext(ctx, os.Args[0], "job-execute", work.Execution.ExecutionID)

	// Set environment variables
	cmd.Env = append(os.Environ(),
		fmt.Sprintf("JOB_ID=%s", work.Job.JobID),
		fmt.Sprintf("EXECUTION_ID=%s", work.Execution.ExecutionID),
	)

	// Execute command
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("process execution failed: %v, output: %s", err, string(output))
	}

	return nil
}