backend/internal/grpc/runner/task_manager.go

package runner

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"with_you/proto/runner"

	"github.com/google/uuid"
	"go.uber.org/zap"
)

// TaskCallback 任务完成回调函数类型
type TaskCallback func(result *runner.TaskResult)

// PendingTask 表示一个等待处理的任务
type PendingTask struct {
	Task      *runner.Task
	Callback  TaskCallback
	CreatedAt time.Time
}

// TaskManager 管理任务的生命周期
type TaskManager struct {
	mu      sync.RWMutex
	pending map[string]*PendingTask
	hub     *RunnerHub
	timeout time.Duration
	logger  *zap.Logger
}

// NewTaskManager 创建新的任务管理器
func NewTaskManager(hub *RunnerHub, timeout time.Duration, logger *zap.Logger) *TaskManager {
	return &TaskManager{
		pending: make(map[string]*PendingTask),
		hub:     hub,
		timeout: timeout,
		logger:  logger,
	}
}

// SubmitTask 提交任务并等待结果
func (tm *TaskManager) SubmitTask(ctx context.Context, taskType runner.TaskType, payload any, callback TaskCallback) (*runner.Task, error) {
	payloadBytes, err := json.Marshal(payload)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal payload: %w", err)
	}

	task := &runner.Task{
		TaskId:         uuid.New().String(),
		Type:           taskType,
		CreatedAt:      time.Now().Unix(),
		TimeoutSeconds: int32(tm.timeout.Seconds()),
		Payload:        payloadBytes,
	}

	tm.mu.Lock()
	tm.pending[task.TaskId] = &PendingTask{
		Task:      task,
		Callback:  callback,
		CreatedAt: time.Now(),
	}
	tm.mu.Unlock()

	tm.logger.Debug("Task submitted",
		zap.String("task_id", task.TaskId),
		zap.String("type", taskType.String()))

	// 获取可用 Runner
	conn, err := tm.hub.GetAvailableRunner(taskType)
	if err != nil {
		tm.cleanupTask(task.TaskId)
		return nil, fmt.Errorf("no available runner: %w", err)
	}

	// 发送任务
	if err := tm.hub.SendTask(conn.ID, task); err != nil {
		tm.cleanupTask(task.TaskId)
		return nil, fmt.Errorf("failed to send task: %w", err)
	}

	// 启动超时检查
	go tm.waitForTimeout(task.TaskId)

	return task, nil
}

// SubmitTaskSync 同步提交任务并等待结果
func (tm *TaskManager) SubmitTaskSync(ctx context.Context, taskType runner.TaskType, payload any) (*runner.TaskResult, error) {
	resultChan := make(chan *runner.TaskResult, 1)

	task, err := tm.SubmitTask(ctx, taskType, payload, func(result *runner.TaskResult) {
		select {
		case resultChan <- result:
		default:
		}
	})
	if err != nil {
		return nil, err
	}

	select {
	case result := <-resultChan:
		return result, nil
	case <-ctx.Done():
		tm.cleanupTask(task.TaskId)
		return nil, ctx.Err()
	case <-time.After(tm.timeout):
		tm.cleanupTask(task.TaskId)
		return nil, fmt.Errorf("task timeout")
	}
}

// HandleResult 处理任务结果
func (tm *TaskManager) HandleResult(result *runner.TaskResult) {
	tm.mu.Lock()
	pending, ok := tm.pending[result.TaskId]
	if ok {
		delete(tm.pending, result.TaskId)
	}
	tm.mu.Unlock()

	if !ok {
		tm.logger.Warn("Received result for unknown task",
			zap.String("task_id", result.TaskId))
		return
	}

	tm.logger.Debug("Task completed",
		zap.String("task_id", result.TaskId),
		zap.String("status", result.Status.String()))

	if pending.Callback != nil {
		pending.Callback(result)
	}
}

// waitForTimeout 等待任务超时
func (tm *TaskManager) waitForTimeout(taskID string) {
	time.Sleep(tm.timeout)

	tm.mu.Lock()
	defer tm.mu.Unlock()

	if pending, ok := tm.pending[taskID]; ok {
		tm.logger.Warn("Task timeout",
			zap.String("task_id", taskID),
			zap.String("type", pending.Task.Type.String()))

		delete(tm.pending, taskID)

		// 调用回调，返回超时结果
		if pending.Callback != nil {
			pending.Callback(&runner.TaskResult{
				TaskId:       taskID,
				Status:       runner.TaskStatus_TASK_STATUS_TIMEOUT,
				ErrorMessage: "task timeout",
				CompletedAt:  time.Now().Unix(),
			})
		}
	}
}

// cleanupTask 清理任务
func (tm *TaskManager) cleanupTask(taskID string) {
	tm.mu.Lock()
	defer tm.mu.Unlock()
	delete(tm.pending, taskID)
}

// CancelTask 取消任务
func (tm *TaskManager) CancelTask(taskID string) error {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	if pending, ok := tm.pending[taskID]; ok {
		tm.logger.Debug("Task cancelled", zap.String("task_id", taskID))

		delete(tm.pending, taskID)

		// 调用回调，返回取消结果
		if pending.Callback != nil {
			pending.Callback(&runner.TaskResult{
				TaskId:       taskID,
				Status:       runner.TaskStatus_TASK_STATUS_CANCELLED,
				ErrorMessage: "task cancelled",
				CompletedAt:  time.Now().Unix(),
			})
		}
		return nil
	}

	return fmt.Errorf("task not found: %s", taskID)
}

// GetPendingTaskCount 获取等待处理的任务数量
func (tm *TaskManager) GetPendingTaskCount() int {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	return len(tm.pending)
}

// SetHub 设置 RunnerHub（用于解决循环依赖）
func (tm *TaskManager) SetHub(hub *RunnerHub) {
	tm.hub = hub
}

// CleanupStaleTasks 清理超时的任务
func (tm *TaskManager) CleanupStaleTasks() {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	now := time.Now()
	for taskID, pending := range tm.pending {
		if now.Sub(pending.CreatedAt) > tm.timeout {
			tm.logger.Warn("Cleaning up stale task",
				zap.String("task_id", taskID),
				zap.String("type", pending.Task.Type.String()))

			delete(tm.pending, taskID)

			if pending.Callback != nil {
				pending.Callback(&runner.TaskResult{
					TaskId:       taskID,
					Status:       runner.TaskStatus_TASK_STATUS_TIMEOUT,
					ErrorMessage: "task timeout",
					CompletedAt:  now.Unix(),
				})
			}
		}
	}
}