package grpc import ( "context" "encoding/json" "errors" "fmt" "log/slog" "sync" "time" "schedule_converter/service" pb "schedule_converter/proto/runner" ) type TaskHandler struct { cancelledTasks sync.Map service *service.Service } func NewTaskHandler() *TaskHandler { return &TaskHandler{service: service.NewService()} } func supportedTaskTypes() []pb.TaskType { return []pb.TaskType{ pb.TaskType_TASK_TYPE_LOGIN, pb.TaskType_TASK_TYPE_GET_SCHEDULE, pb.TaskType_TASK_TYPE_GET_GRADES, pb.TaskType_TASK_TYPE_GET_EXAMS, pb.TaskType_TASK_TYPE_GET_EMPTY_CLASSROOM, pb.TaskType_TASK_TYPE_GET_CET, } } func unmarshalPayload(payload []byte, v any) error { if len(payload) == 0 { return fmt.Errorf("任务 payload 为空") } if err := json.Unmarshal(payload, v); err != nil { return fmt.Errorf("解析任务 payload 失败: %w", err) } return nil } 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) case pb.TaskType_TASK_TYPE_GET_CET: return h.getCET(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) }