schedule_converter/grpc/handler.go

package grpc

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"schedule_converter/service"

	pb "schedule_converter/proto/runner"
)

type TaskHandler struct {
	service      *service.Service
	cancels      sync.Map // taskID -> context.CancelFunc（运行中的任务）
	preCancelled sync.Map // taskID -> struct{}（任务开始前就已取消）
}

func NewTaskHandler() *TaskHandler {
	return &TaskHandler{
		service: service.NewService(),
	}
}

func (h *TaskHandler) Execute(task *pb.Task) *pb.TaskResult {
	log := slog.With("task_id", task.TaskId, "task_type", task.Type.String())
	log.Info("开始执行任务")

	result := &pb.TaskResult{
		TaskId:      task.TaskId,
		CompletedAt: time.Now().UnixMilli(),
	}

	// 任务开始前检查是否已被预取消。
	if _, cancelled := h.preCancelled.LoadAndDelete(task.TaskId); cancelled {
		result.Status = pb.TaskStatus_TASK_STATUS_CANCELLED
		result.ErrorMessage = "任务已取消"
		log.Info("任务已取消")
		return result
	}

	ctx, cancel := h.buildTaskContext(task)
	defer cancel()

	data, err := h.dispatch(ctx, task)

	if err != nil {
		if errors.Is(err, errUnknownTaskType) {
			result.Status = pb.TaskStatus_TASK_STATUS_FAILED
			result.ErrorMessage = err.Error()
			log.Warn("未知的任务类型")
			return result
		}
		result.Status = pb.TaskStatus_TASK_STATUS_FAILED
		result.ErrorMessage = err.Error()
		log.Error("任务执行失败", "error", err)
		return result
	}

	result.Status = pb.TaskStatus_TASK_STATUS_SUCCESS
	result.Data = data
	log.Info("任务执行成功")
	return result
}

// buildTaskContext 构造任务执行 context：根据 TimeoutSeconds 设置截止时间，
// 并注册 cancel 函数，使 CancelTask 能够中断正在执行的任务。
func (h *TaskHandler) buildTaskContext(task *pb.Task) (context.Context, context.CancelFunc) {
	ctx := context.Background()
	if task.TimeoutSeconds > 0 {
		ctx, cancel := context.WithTimeout(ctx, time.Duration(task.TimeoutSeconds)*time.Second)
		h.cancels.Store(task.TaskId, context.CancelFunc(cancel))
		return ctx, func() {
			h.cancels.Delete(task.TaskId)
			cancel()
		}
	}
	ctx, cancel := context.WithCancel(ctx)
	h.cancels.Store(task.TaskId, context.CancelFunc(cancel))
	return ctx, func() {
		h.cancels.Delete(task.TaskId)
		cancel()
	}
}

var errUnknownTaskType = fmt.Errorf("unknown task type")

// taskHandlers 将任务类型映射到处理方法。
var taskHandlers = map[pb.TaskType]func(h *TaskHandler, ctx context.Context, payload []byte) ([]byte, error){
	pb.TaskType_TASK_TYPE_GET_SCHEDULE:        (*TaskHandler).getSchedule,
	pb.TaskType_TASK_TYPE_GET_GRADES:          (*TaskHandler).getGrades,
	pb.TaskType_TASK_TYPE_GET_EXAMS:           (*TaskHandler).getExams,
	pb.TaskType_TASK_TYPE_LOGIN:               (*TaskHandler).login,
	pb.TaskType_TASK_TYPE_GET_EMPTY_CLASSROOM: (*TaskHandler).getEmptyClassrooms,
}

// supportedTaskTypes 返回当前支持的任务类型列表，供注册 capabilities 使用。
func supportedTaskTypes() []pb.TaskType {
	types := make([]pb.TaskType, 0, len(taskHandlers))
	for t := range taskHandlers {
		types = append(types, t)
	}
	return types
}

func (h *TaskHandler) dispatch(ctx context.Context, task *pb.Task) ([]byte, error) {
	handler, ok := taskHandlers[task.Type]
	if !ok {
		return nil, fmt.Errorf("%w: %v", errUnknownTaskType, task.Type)
	}
	return handler(h, ctx, task.Payload)
}

// CancelTask 取消任务：若任务正在执行则中断其 context；若尚未开始则记为预取消。
func (h *TaskHandler) CancelTask(taskID string) {
	if v, ok := h.cancels.LoadAndDelete(taskID); ok {
		if cancel, ok := v.(context.CancelFunc); ok {
			cancel()
		}
		return
	}
	// 任务尚未开始执行，标记为预取消，待 Execute 时直接返回取消。
	h.preCancelled.Store(taskID, struct{}{})
	slog.Info("任务已标记为取消", "task_id", taskID)
}

// unmarshalPayload 解码任务载荷。
func unmarshalPayload(payload []byte, v any) error {
	return json.Unmarshal(payload, v)
}