115 lines
2.8 KiB
Go
115 lines
2.8 KiB
Go
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)
|
|
}
|