feat(grpc): add gRPC server infrastructure and schedule sync service
Add gRPC server support with configurable TLS, environment-based settings, and Wire injection. Implement schedule synchronization service with task runner integration for external course data fetching. - Add gRPC configuration with env var overrides (APP_GRPC_ENABLED, APP_GRPC_PORT, etc.) - Create gRPC server infrastructure with runner task management - Implement ScheduleSyncService for course data synchronization - Add sync endpoint to schedule handler for external system integration - Update repository with context-aware batch delete operation - Wire all dependencies including zap logger for structured logging
This commit is contained in:
299
internal/grpc/runner/hub.go
Normal file
299
internal/grpc/runner/hub.go
Normal file
@@ -0,0 +1,299 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"carrot_bbs/proto/runner"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// RunnerConnection 表示一个已连接的 Runner
|
||||
type RunnerConnection struct {
|
||||
ID string
|
||||
Version string
|
||||
Capabilities map[string]string
|
||||
Stream runner.RunnerHub_ConnectServer
|
||||
LastSeen time.Time
|
||||
sendChan chan *runner.StreamMessage
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// RunnerHub 管理 Runner 连接和任务分发
|
||||
type RunnerHub struct {
|
||||
runner.UnimplementedRunnerHubServer // 嵌入以实现 gRPC 接口
|
||||
|
||||
mu sync.RWMutex
|
||||
runners map[string]*RunnerConnection
|
||||
taskManager *TaskManager
|
||||
registerChan chan *RunnerConnection
|
||||
unregisterChan chan string
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewRunnerHub 创建新的 RunnerHub
|
||||
func NewRunnerHub(taskManager *TaskManager, logger *zap.Logger) *RunnerHub {
|
||||
return &RunnerHub{
|
||||
runners: make(map[string]*RunnerConnection),
|
||||
taskManager: taskManager,
|
||||
registerChan: make(chan *RunnerConnection, 100),
|
||||
unregisterChan: make(chan string, 100),
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// Connect 实现 RunnerHub 服务的 Connect 方法
|
||||
func (h *RunnerHub) Connect(stream runner.RunnerHub_ConnectServer) error {
|
||||
var conn *RunnerConnection
|
||||
ctx := stream.Context()
|
||||
|
||||
defer func() {
|
||||
if conn != nil {
|
||||
h.unregisterRunner(conn.ID)
|
||||
}
|
||||
}()
|
||||
|
||||
// 启动发送协程
|
||||
sendDone := make(chan struct{})
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-sendDone:
|
||||
return nil
|
||||
default:
|
||||
msg, err := stream.Recv()
|
||||
if err != nil {
|
||||
h.logger.Debug("Stream receive error", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
switch m := msg.Message.(type) {
|
||||
case *runner.StreamMessage_Register:
|
||||
conn = h.handleRegister(m.Register, stream)
|
||||
if conn != nil {
|
||||
// 启动发送协程
|
||||
go h.sendLoop(conn, sendDone)
|
||||
}
|
||||
case *runner.StreamMessage_Heartbeat:
|
||||
h.handleHeartbeat(conn, m.Heartbeat)
|
||||
case *runner.StreamMessage_Result:
|
||||
h.handleTaskResult(m.Result)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// sendLoop 处理向 Runner 发送消息
|
||||
func (h *RunnerHub) sendLoop(conn *RunnerConnection, done chan struct{}) {
|
||||
for {
|
||||
select {
|
||||
case msg, ok := <-conn.sendChan:
|
||||
if !ok {
|
||||
// Channel 已关闭
|
||||
return
|
||||
}
|
||||
if msg == nil {
|
||||
// 忽略 nil 消息
|
||||
continue
|
||||
}
|
||||
if err := conn.Stream.Send(msg); err != nil {
|
||||
h.logger.Error("Failed to send message to runner",
|
||||
zap.String("runner_id", conn.ID),
|
||||
zap.Error(err))
|
||||
return
|
||||
}
|
||||
case <-done:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *RunnerHub) handleRegister(req *runner.RegisterRequest, stream runner.RunnerHub_ConnectServer) *RunnerConnection {
|
||||
h.logger.Info("Runner registering",
|
||||
zap.String("runner_id", req.RunnerId),
|
||||
zap.String("version", req.Version),
|
||||
zap.Any("capabilities", req.Capabilities))
|
||||
|
||||
conn := &RunnerConnection{
|
||||
ID: req.RunnerId,
|
||||
Version: req.Version,
|
||||
Capabilities: req.Capabilities,
|
||||
Stream: stream,
|
||||
LastSeen: time.Now(),
|
||||
sendChan: make(chan *runner.StreamMessage, 100),
|
||||
logger: h.logger,
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
// 如果已存在旧连接，先关闭
|
||||
if oldConn, ok := h.runners[req.RunnerId]; ok {
|
||||
close(oldConn.sendChan)
|
||||
}
|
||||
h.runners[req.RunnerId] = conn
|
||||
h.mu.Unlock()
|
||||
|
||||
// 发送注册响应
|
||||
if err := conn.send(&runner.StreamMessage{
|
||||
Message: &runner.StreamMessage_RegisterResponse{
|
||||
RegisterResponse: &runner.RegisterResponse{
|
||||
Success: true,
|
||||
SessionId: conn.ID,
|
||||
Message: "Registered successfully",
|
||||
HeartbeatInterval: 30,
|
||||
},
|
||||
},
|
||||
}); err != nil {
|
||||
h.logger.Error("Failed to send register response", zap.Error(err))
|
||||
}
|
||||
|
||||
h.logger.Info("Runner registered successfully", zap.String("runner_id", req.RunnerId))
|
||||
return conn
|
||||
}
|
||||
|
||||
func (h *RunnerHub) handleHeartbeat(conn *RunnerConnection, hb *runner.Heartbeat) {
|
||||
if conn == nil {
|
||||
return
|
||||
}
|
||||
conn.LastSeen = time.Now()
|
||||
|
||||
if err := conn.send(&runner.StreamMessage{
|
||||
Message: &runner.StreamMessage_HeartbeatAck{
|
||||
HeartbeatAck: &runner.HeartbeatAck{
|
||||
Timestamp: time.Now().Unix(),
|
||||
},
|
||||
},
|
||||
}); err != nil {
|
||||
h.logger.Debug("Failed to send heartbeat ack",
|
||||
zap.String("runner_id", conn.ID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
func (h *RunnerHub) handleTaskResult(result *runner.TaskResult) {
|
||||
h.logger.Debug("Received task result",
|
||||
zap.String("task_id", result.TaskId),
|
||||
zap.String("status", result.Status.String()))
|
||||
|
||||
if h.taskManager != nil {
|
||||
h.taskManager.HandleResult(result)
|
||||
}
|
||||
}
|
||||
|
||||
// unregisterRunner 注销 Runner
|
||||
func (h *RunnerHub) unregisterRunner(runnerID string) {
|
||||
h.mu.Lock()
|
||||
if conn, ok := h.runners[runnerID]; ok {
|
||||
close(conn.sendChan)
|
||||
delete(h.runners, runnerID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
h.logger.Info("Runner unregistered", zap.String("runner_id", runnerID))
|
||||
}
|
||||
|
||||
// SendTask 向指定 Runner 发送任务
|
||||
func (h *RunnerHub) SendTask(runnerID string, task *runner.Task) error {
|
||||
h.mu.RLock()
|
||||
conn, ok := h.runners[runnerID]
|
||||
h.mu.RUnlock()
|
||||
|
||||
if !ok {
|
||||
return fmt.Errorf("runner not found: %s", runnerID)
|
||||
}
|
||||
|
||||
return conn.send(&runner.StreamMessage{
|
||||
Message: &runner.StreamMessage_Task{
|
||||
Task: task,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// GetAvailableRunner 获取支持指定任务类型的可用 Runner
|
||||
func (h *RunnerHub) GetAvailableRunner(taskType runner.TaskType) (*RunnerConnection, error) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
|
||||
taskTypeStr := taskType.String()
|
||||
for _, conn := range h.runners {
|
||||
for capKey, capValue := range conn.Capabilities {
|
||||
// 支持两种格式：
|
||||
// 1. key = "TASK_TYPE_GET_SCHEDULE", value = "true"
|
||||
// 2. key = "schedule", value = "TASK_TYPE_GET_SCHEDULE"
|
||||
if capKey == taskTypeStr && capValue == "true" {
|
||||
return conn, nil
|
||||
}
|
||||
if capValue == taskTypeStr {
|
||||
return conn, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("no available runner for task type: %s", taskType)
|
||||
}
|
||||
|
||||
// GetAllRunners 获取所有已连接的 Runner 信息
|
||||
func (h *RunnerHub) GetAllRunners() []map[string]interface{} {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
|
||||
result := make([]map[string]interface{}, 0, len(h.runners))
|
||||
for id, conn := range h.runners {
|
||||
result = append(result, map[string]interface{}{
|
||||
"id": id,
|
||||
"version": conn.Version,
|
||||
"capabilities": conn.Capabilities,
|
||||
"last_seen": conn.LastSeen,
|
||||
})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// SetTaskManager 设置任务管理器（用于解决循环依赖）
|
||||
func (h *RunnerHub) SetTaskManager(tm *TaskManager) {
|
||||
h.taskManager = tm
|
||||
}
|
||||
|
||||
// send 向 RunnerConnection 发送消息
|
||||
func (c *RunnerConnection) send(msg *runner.StreamMessage) error {
|
||||
select {
|
||||
case c.sendChan <- msg:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("send buffer full for runner: %s", c.ID)
|
||||
}
|
||||
}
|
||||
|
||||
// CheckStaleRunners 检查并清理超时的 Runner 连接
|
||||
func (h *RunnerHub) CheckStaleRunners(timeout time.Duration) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
for id, conn := range h.runners {
|
||||
if now.Sub(conn.LastSeen) > timeout {
|
||||
h.logger.Warn("Runner timeout, removing",
|
||||
zap.String("runner_id", id),
|
||||
zap.Time("last_seen", conn.LastSeen))
|
||||
close(conn.sendChan)
|
||||
delete(h.runners, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown 关闭所有 Runner 连接
|
||||
func (h *RunnerHub) Shutdown(ctx context.Context) error {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
for id, conn := range h.runners {
|
||||
h.logger.Info("Closing runner connection", zap.String("runner_id", id))
|
||||
close(conn.sendChan)
|
||||
delete(h.runners, id)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
145
internal/grpc/runner/server.go
Normal file
145
internal/grpc/runner/server.go
Normal file
@@ -0,0 +1,145 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"carrot_bbs/internal/config"
|
||||
"carrot_bbs/proto/runner"
|
||||
|
||||
"go.uber.org/zap"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/credentials"
|
||||
"google.golang.org/grpc/credentials/insecure"
|
||||
)
|
||||
|
||||
// Server gRPC 服务器
|
||||
type Server struct {
|
||||
cfg *config.GRPCConfig
|
||||
logger *zap.Logger
|
||||
hub *RunnerHub
|
||||
grpcSrv *grpc.Server
|
||||
}
|
||||
|
||||
// NewServer 创建新的 gRPC 服务器
|
||||
func NewServer(cfg *config.GRPCConfig, logger *zap.Logger) *Server {
|
||||
// 创建 TaskManager 和 RunnerHub，解决循环依赖
|
||||
taskManager := NewTaskManager(nil, 60*time.Second, logger)
|
||||
hub := NewRunnerHub(taskManager, logger)
|
||||
taskManager.SetHub(hub)
|
||||
|
||||
return &Server{
|
||||
cfg: cfg,
|
||||
logger: logger,
|
||||
hub: hub,
|
||||
}
|
||||
}
|
||||
|
||||
// Start 启动 gRPC 服务器
|
||||
func (s *Server) Start() error {
|
||||
if !s.cfg.Enabled {
|
||||
s.logger.Info("gRPC server is disabled")
|
||||
return nil
|
||||
}
|
||||
|
||||
lis, err := net.Listen("tcp", s.cfg.Address())
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to listen on %s: %w", s.cfg.Address(), err)
|
||||
}
|
||||
|
||||
// 配置 gRPC 服务器选项
|
||||
var opts []grpc.ServerOption
|
||||
|
||||
// 配置 TLS
|
||||
if s.cfg.TLSEnabled {
|
||||
if s.cfg.TLSCertFile == "" || s.cfg.TLSKeyFile == "" {
|
||||
return fmt.Errorf("TLS enabled but cert/key files not specified")
|
||||
}
|
||||
creds, err := credentials.NewServerTLSFromFile(s.cfg.TLSCertFile, s.cfg.TLSKeyFile)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load TLS credentials: %w", err)
|
||||
}
|
||||
opts = append(opts, grpc.Creds(creds))
|
||||
s.logger.Info("gRPC server TLS enabled")
|
||||
} else {
|
||||
s.logger.Warn("gRPC server running without TLS")
|
||||
}
|
||||
|
||||
s.grpcSrv = grpc.NewServer(opts...)
|
||||
runner.RegisterRunnerHubServer(s.grpcSrv, s.hub)
|
||||
|
||||
s.logger.Info("gRPC server starting", zap.String("address", s.cfg.Address()))
|
||||
|
||||
go func() {
|
||||
if err := s.grpcSrv.Serve(lis); err != nil {
|
||||
s.logger.Error("gRPC server error", zap.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
// 启动后台任务清理协程
|
||||
go s.backgroundCleanup()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// backgroundCleanup 后台清理任务
|
||||
func (s *Server) backgroundCleanup() {
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for range ticker.C {
|
||||
// 清理超时的 Runner 连接
|
||||
s.hub.CheckStaleRunners(90 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
// Stop 停止 gRPC 服务器
|
||||
func (s *Server) Stop() {
|
||||
if s.grpcSrv != nil {
|
||||
s.logger.Info("Stopping gRPC server")
|
||||
s.grpcSrv.GracefulStop()
|
||||
}
|
||||
}
|
||||
|
||||
// GetHub 获取 RunnerHub
|
||||
func (s *Server) GetHub() *RunnerHub {
|
||||
return s.hub
|
||||
}
|
||||
|
||||
// GetTaskManager 获取 TaskManager
|
||||
func (s *Server) GetTaskManager() *TaskManager {
|
||||
return s.hub.taskManager
|
||||
}
|
||||
|
||||
// Shutdown 优雅关闭
|
||||
func (s *Server) Shutdown(ctx context.Context) error {
|
||||
if s.grpcSrv == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
s.grpcSrv.GracefulStop()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
s.logger.Info("gRPC server stopped gracefully")
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
s.logger.Warn("gRPC server shutdown timeout, forcing stop")
|
||||
s.grpcSrv.Stop()
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// GetClientCredentials 获取客户端凭证配置
|
||||
func GetClientCredentials(cfg *config.GRPCConfig) (credentials.TransportCredentials, error) {
|
||||
if cfg.TLSEnabled {
|
||||
return credentials.NewClientTLSFromFile(cfg.TLSCertFile, "")
|
||||
}
|
||||
return insecure.NewCredentials(), nil
|
||||
}
|
||||
235
internal/grpc/runner/task_manager.go
Normal file
235
internal/grpc/runner/task_manager.go
Normal file
@@ -0,0 +1,235 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"carrot_bbs/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 interface{}, 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 interface{}) (*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(),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user