schedule_converter/grpc/handler.go

package grpc

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

	pb "schedule_converter/proto/runner"
)

type TaskHandler struct {
	cancelledTasks sync.Map
}

func NewTaskHandler() *TaskHandler {
	return &TaskHandler{}
}

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 h.IsCancelled(task.TaskId) {
		result.Status = pb.TaskStatus_TASK_STATUS_CANCELLED
		result.ErrorMessage = "任务已取消"
		log.Info("任务已取消")
		return result
	}

	ctx := context.Background()
	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
}

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

func (h *TaskHandler) dispatch(ctx context.Context, task *pb.Task) ([]byte, error) {
	switch task.Type {
	case pb.TaskType_TASK_TYPE_GET_SCHEDULE:
		return h.getSchedule(ctx, task.Payload)
	case pb.TaskType_TASK_TYPE_GET_GRADES:
		return h.getGrades(ctx, task.Payload)
	case pb.TaskType_TASK_TYPE_GET_EXAMS:
		return h.getExams(ctx, task.Payload)
	case pb.TaskType_TASK_TYPE_LOGIN:
		return h.login(ctx, task.Payload)
	case pb.TaskType_TASK_TYPE_GET_EMPTY_CLASSROOM:
		return h.getEmptyClassrooms(ctx, task.Payload)
	default:
		return nil, fmt.Errorf("%w: %v", errUnknownTaskType, task.Type)
	}
}

func (h *TaskHandler) CancelTask(taskID string) {
	h.cancelledTasks.Store(taskID, true)
	slog.Info("任务已标记为取消", "task_id", taskID)
}

func (h *TaskHandler) IsCancelled(taskID string) bool {
	v, ok := h.cancelledTasks.Load(taskID)
	return ok && v.(bool)
}