mirror of
https://github.com/crawlab-team/crawlab.git
synced 2026-01-25 17:42:25 +01:00
feat: implement task management service operations, stream manager, and worker pool
- Added service_operations.go for task management including run, cancel, and execution logic. - Introduced stream_manager.go to handle task streams and manage cancellation signals. - Created worker_pool.go to manage a bounded pool of workers for executing tasks concurrently. - Implemented graceful shutdown and cleanup mechanisms for task runners and streams. - Enhanced error handling and logging throughout the task management process.
This commit is contained in:
Marvin Zhang
@@ -2,9 +2,7 @@ package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"runtime"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -55,137 +53,6 @@ type Service struct {
|
||||
interfaces.Logger
|
||||
}
|
||||
|
||||
// StreamManager manages task streams without goroutine leaks
|
||||
type StreamManager struct {
|
||||
streams sync.Map // map[primitive.ObjectID]*TaskStream
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
service *Service
|
||||
messageQueue chan *StreamMessage
|
||||
maxStreams int
|
||||
}
|
||||
|
||||
// TaskStream represents a single task's stream
|
||||
type TaskStream struct {
|
||||
taskId primitive.ObjectID
|
||||
stream grpc.TaskService_SubscribeClient
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
lastActive time.Time
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// StreamMessage represents a message from a stream
|
||||
type StreamMessage struct {
|
||||
taskId primitive.ObjectID
|
||||
msg *grpc.TaskServiceSubscribeResponse
|
||||
err error
|
||||
}
|
||||
|
||||
// taskRequest represents a task execution request
|
||||
type taskRequest struct {
|
||||
taskId primitive.ObjectID
|
||||
}
|
||||
|
||||
// TaskWorkerPool manages a bounded pool of workers for task execution
|
||||
type TaskWorkerPool struct {
|
||||
workers int
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
taskQueue chan taskRequest
|
||||
service *Service
|
||||
}
|
||||
|
||||
func NewTaskWorkerPool(workers int, service *Service) *TaskWorkerPool {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
// Use a more generous queue size to handle task bursts
|
||||
// Queue size is workers * 5 to allow for better buffering
|
||||
queueSize := workers * 5
|
||||
if queueSize < 50 {
|
||||
queueSize = 50 // Minimum queue size
|
||||
}
|
||||
|
||||
return &TaskWorkerPool{
|
||||
workers: workers,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
taskQueue: make(chan taskRequest, queueSize),
|
||||
service: service,
|
||||
}
|
||||
}
|
||||
|
||||
func (pool *TaskWorkerPool) Start() {
|
||||
for i := 0; i < pool.workers; i++ {
|
||||
pool.wg.Add(1)
|
||||
go pool.worker(i)
|
||||
}
|
||||
}
|
||||
|
||||
func (pool *TaskWorkerPool) Stop() {
|
||||
pool.cancel()
|
||||
close(pool.taskQueue)
|
||||
pool.wg.Wait()
|
||||
}
|
||||
|
||||
func (pool *TaskWorkerPool) SubmitTask(taskId primitive.ObjectID) error {
|
||||
req := taskRequest{
|
||||
taskId: taskId,
|
||||
}
|
||||
|
||||
select {
|
||||
case pool.taskQueue <- req:
|
||||
pool.service.Debugf("task[%s] queued for parallel execution, queue usage: %d/%d",
|
||||
taskId.Hex(), len(pool.taskQueue), cap(pool.taskQueue))
|
||||
return nil // Return immediately - task will execute in parallel
|
||||
case <-pool.ctx.Done():
|
||||
return fmt.Errorf("worker pool is shutting down")
|
||||
default:
|
||||
queueLen := len(pool.taskQueue)
|
||||
queueCap := cap(pool.taskQueue)
|
||||
pool.service.Warnf("task queue is full (%d/%d), consider increasing task.workers configuration",
|
||||
queueLen, queueCap)
|
||||
return fmt.Errorf("task queue is full (%d/%d), consider increasing task.workers configuration",
|
||||
queueLen, queueCap)
|
||||
}
|
||||
}
|
||||
|
||||
func (pool *TaskWorkerPool) worker(workerID int) {
|
||||
defer pool.wg.Done()
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
pool.service.Errorf("worker %d panic recovered: %v", workerID, r)
|
||||
}
|
||||
}()
|
||||
|
||||
pool.service.Debugf("worker %d started", workerID)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-pool.ctx.Done():
|
||||
pool.service.Debugf("worker %d shutting down", workerID)
|
||||
return
|
||||
case req, ok := <-pool.taskQueue:
|
||||
if !ok {
|
||||
pool.service.Debugf("worker %d: task queue closed", workerID)
|
||||
return
|
||||
}
|
||||
|
||||
// Execute task asynchronously - each worker handles one task at a time
|
||||
// but multiple workers can process different tasks simultaneously
|
||||
pool.service.Debugf("worker %d processing task[%s]", workerID, req.taskId.Hex())
|
||||
err := pool.service.executeTask(req.taskId)
|
||||
if err != nil {
|
||||
pool.service.Errorf("worker %d failed to execute task[%s]: %v",
|
||||
workerID, req.taskId.Hex(), err)
|
||||
} else {
|
||||
pool.service.Debugf("worker %d completed task[%s]", workerID, req.taskId.Hex())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (svc *Service) Start() {
|
||||
// Initialize context for graceful shutdown
|
||||
svc.ctx, svc.cancel = context.WithCancel(context.Background())
|
||||
@@ -308,14 +175,6 @@ func (svc *Service) startGoroutineMonitoring() {
|
||||
}()
|
||||
}
|
||||
|
||||
func (svc *Service) Run(taskId primitive.ObjectID) (err error) {
|
||||
return svc.runTask(taskId)
|
||||
}
|
||||
|
||||
func (svc *Service) Cancel(taskId primitive.ObjectID, force bool) (err error) {
|
||||
return svc.cancelTask(taskId, force)
|
||||
}
|
||||
|
||||
func (svc *Service) fetchAndRunTasks() {
|
||||
defer svc.wg.Done()
|
||||
defer func() {
|
||||
@@ -507,35 +366,6 @@ func (svc *Service) getRunnerCount() (count int) {
|
||||
return count
|
||||
}
|
||||
|
||||
func (svc *Service) getRunner(taskId primitive.ObjectID) (r interfaces.TaskRunner, err error) {
|
||||
svc.Debugf("get runner: taskId[%v]", taskId)
|
||||
v, ok := svc.runners.Load(taskId)
|
||||
if !ok {
|
||||
err = fmt.Errorf("task[%s] not exists", taskId.Hex())
|
||||
svc.Errorf("get runner error: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
switch v := v.(type) {
|
||||
case interfaces.TaskRunner:
|
||||
r = v
|
||||
default:
|
||||
err = fmt.Errorf("invalid type: %T", v)
|
||||
svc.Errorf("get runner error: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func (svc *Service) addRunner(taskId primitive.ObjectID, r interfaces.TaskRunner) {
|
||||
svc.Debugf("add runner: taskId[%s]", taskId.Hex())
|
||||
svc.runners.Store(taskId, r)
|
||||
}
|
||||
|
||||
func (svc *Service) deleteRunner(taskId primitive.ObjectID) {
|
||||
svc.Debugf("delete runner: taskId[%v]", taskId)
|
||||
svc.runners.Delete(taskId)
|
||||
}
|
||||
|
||||
func (svc *Service) updateNodeStatus() (err error) {
|
||||
// current node
|
||||
n, err := svc.GetCurrentNode()
|
||||
@@ -581,411 +411,6 @@ func (svc *Service) fetchTask() (tid primitive.ObjectID, err error) {
|
||||
return tid, nil
|
||||
}
|
||||
|
||||
func (svc *Service) runTask(taskId primitive.ObjectID) (err error) {
|
||||
// attempt to get runner from pool
|
||||
_, ok := svc.runners.Load(taskId)
|
||||
if ok {
|
||||
err = fmt.Errorf("task[%s] already exists", taskId.Hex())
|
||||
svc.Errorf("run task error: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// Use worker pool for bounded task execution
|
||||
return svc.workerPool.SubmitTask(taskId)
|
||||
}
|
||||
|
||||
// executeTask is the actual task execution logic called by worker pool
|
||||
func (svc *Service) executeTask(taskId primitive.ObjectID) (err error) {
|
||||
// attempt to get runner from pool
|
||||
_, ok := svc.runners.Load(taskId)
|
||||
if ok {
|
||||
err = fmt.Errorf("task[%s] already exists", taskId.Hex())
|
||||
svc.Errorf("execute task error: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// create a new task runner
|
||||
r, err := newTaskRunner(taskId, svc)
|
||||
if err != nil {
|
||||
err = fmt.Errorf("failed to create task runner: %v", err)
|
||||
svc.Errorf("execute task error: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// add runner to pool
|
||||
svc.addRunner(taskId, r)
|
||||
|
||||
// Ensure cleanup always happens
|
||||
defer func() {
|
||||
if rec := recover(); rec != nil {
|
||||
svc.Errorf("task[%s] panic recovered: %v", taskId.Hex(), rec)
|
||||
}
|
||||
// Always cleanup runner from pool and stream
|
||||
svc.deleteRunner(taskId)
|
||||
svc.streamManager.RemoveTaskStream(taskId)
|
||||
}()
|
||||
|
||||
// Add task to stream manager for cancellation support
|
||||
if err := svc.streamManager.AddTaskStream(r.GetTaskId()); err != nil {
|
||||
svc.Warnf("failed to add task[%s] to stream manager: %v", r.GetTaskId().Hex(), err)
|
||||
svc.Warnf("task[%s] will not be able to receive cancellation messages", r.GetTaskId().Hex())
|
||||
} else {
|
||||
svc.Debugf("task[%s] added to stream manager for cancellation support", r.GetTaskId().Hex())
|
||||
}
|
||||
|
||||
// run task process (blocking) error or finish after task runner ends
|
||||
if err := r.Run(); err != nil {
|
||||
switch {
|
||||
case errors.Is(err, constants.ErrTaskError):
|
||||
svc.Errorf("task[%s] finished with error: %v", r.GetTaskId().Hex(), err)
|
||||
case errors.Is(err, constants.ErrTaskCancelled):
|
||||
svc.Infof("task[%s] cancelled", r.GetTaskId().Hex())
|
||||
default:
|
||||
svc.Errorf("task[%s] finished with unknown error: %v", r.GetTaskId().Hex(), err)
|
||||
}
|
||||
} else {
|
||||
svc.Infof("task[%s] finished successfully", r.GetTaskId().Hex())
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// subscribeTask attempts to subscribe to task stream
|
||||
func (svc *Service) subscribeTask(taskId primitive.ObjectID) (stream grpc.TaskService_SubscribeClient, err error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
req := &grpc.TaskServiceSubscribeRequest{
|
||||
TaskId: taskId.Hex(),
|
||||
}
|
||||
taskClient, err := svc.c.GetTaskClient()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get task client: %v", err)
|
||||
}
|
||||
stream, err = taskClient.Subscribe(ctx, req)
|
||||
if err != nil {
|
||||
svc.Errorf("failed to subscribe task[%s]: %v", taskId.Hex(), err)
|
||||
return nil, err
|
||||
}
|
||||
return stream, nil
|
||||
}
|
||||
|
||||
// subscribeTaskWithRetry attempts to subscribe to task stream with retry logic
|
||||
func (svc *Service) subscribeTaskWithRetry(ctx context.Context, taskId primitive.ObjectID, maxRetries int) (stream grpc.TaskService_SubscribeClient, err error) {
|
||||
for i := 0; i < maxRetries; i++ {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
stream, err = svc.subscribeTask(taskId)
|
||||
if err == nil {
|
||||
return stream, nil
|
||||
}
|
||||
|
||||
svc.Warnf("failed to subscribe task[%s] (attempt %d/%d): %v", taskId.Hex(), i+1, maxRetries, err)
|
||||
|
||||
if i < maxRetries-1 {
|
||||
// Wait before retry with exponential backoff
|
||||
backoff := time.Duration(i+1) * time.Second
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(backoff):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("failed to subscribe after %d retries: %w", maxRetries, err)
|
||||
}
|
||||
|
||||
func (svc *Service) handleStreamMessagesSync(ctx context.Context, taskId primitive.ObjectID, stream grpc.TaskService_SubscribeClient) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
svc.Errorf("handleStreamMessagesSync[%s] panic recovered: %v", taskId.Hex(), r)
|
||||
}
|
||||
// Ensure stream is properly closed
|
||||
if stream != nil {
|
||||
if err := stream.CloseSend(); err != nil {
|
||||
svc.Debugf("task[%s] failed to close stream: %v", taskId.Hex(), err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
svc.Debugf("task[%s] starting synchronous stream message handling", taskId.Hex())
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
svc.Debugf("task[%s] stream handler stopped by context", taskId.Hex())
|
||||
return
|
||||
case <-svc.ctx.Done():
|
||||
svc.Debugf("task[%s] stream handler stopped by service context", taskId.Hex())
|
||||
return
|
||||
default:
|
||||
// Set a reasonable timeout for stream receive
|
||||
recvCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
|
||||
// Create a done channel to handle the recv operation
|
||||
done := make(chan struct {
|
||||
msg *grpc.TaskServiceSubscribeResponse
|
||||
err error
|
||||
}, 1)
|
||||
|
||||
// Use a goroutine only for the blocking recv call, but ensure it's cleaned up
|
||||
go func() {
|
||||
defer cancel()
|
||||
msg, err := stream.Recv()
|
||||
select {
|
||||
case done <- struct {
|
||||
msg *grpc.TaskServiceSubscribeResponse
|
||||
err error
|
||||
}{msg, err}:
|
||||
case <-recvCtx.Done():
|
||||
// Timeout occurred, abandon this receive
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case result := <-done:
|
||||
cancel() // Clean up the context
|
||||
if result.err != nil {
|
||||
if errors.Is(result.err, io.EOF) {
|
||||
svc.Infof("task[%s] received EOF, stream closed", taskId.Hex())
|
||||
return
|
||||
}
|
||||
svc.Errorf("task[%s] stream error: %v", taskId.Hex(), result.err)
|
||||
return
|
||||
}
|
||||
svc.processStreamMessage(taskId, result.msg)
|
||||
case <-recvCtx.Done():
|
||||
cancel()
|
||||
// Timeout on receive - continue to next iteration
|
||||
svc.Debugf("task[%s] stream receive timeout", taskId.Hex())
|
||||
case <-ctx.Done():
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (svc *Service) handleStreamMessages(taskId primitive.ObjectID, stream grpc.TaskService_SubscribeClient, stopCh chan struct{}) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
svc.Errorf("handleStreamMessages[%s] panic recovered: %v", taskId.Hex(), r)
|
||||
}
|
||||
// Ensure stream is properly closed
|
||||
if stream != nil {
|
||||
if err := stream.CloseSend(); err != nil {
|
||||
svc.Debugf("task[%s] failed to close stream: %v", taskId.Hex(), err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Create timeout for stream operations
|
||||
streamTimeout := 30 * time.Second
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-stopCh:
|
||||
svc.Debugf("task[%s] stream handler received stop signal", taskId.Hex())
|
||||
return
|
||||
case <-svc.ctx.Done():
|
||||
svc.Debugf("task[%s] stream handler stopped by service context", taskId.Hex())
|
||||
return
|
||||
default:
|
||||
// Set deadline for receive operation
|
||||
ctx, cancel := context.WithTimeout(context.Background(), streamTimeout)
|
||||
|
||||
// Create a buffered channel to receive the result
|
||||
result := make(chan struct {
|
||||
msg *grpc.TaskServiceSubscribeResponse
|
||||
err error
|
||||
}, 1)
|
||||
|
||||
// Use a single goroutine to handle the blocking Recv call
|
||||
go func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
svc.Errorf("task[%s] stream recv goroutine panic: %v", taskId.Hex(), r)
|
||||
}
|
||||
}()
|
||||
|
||||
msg, err := stream.Recv()
|
||||
select {
|
||||
case result <- struct {
|
||||
msg *grpc.TaskServiceSubscribeResponse
|
||||
err error
|
||||
}{msg, err}:
|
||||
case <-ctx.Done():
|
||||
// Context cancelled, don't send result
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case res := <-result:
|
||||
cancel()
|
||||
if res.err != nil {
|
||||
if errors.Is(res.err, io.EOF) {
|
||||
svc.Infof("task[%s] received EOF, stream closed", taskId.Hex())
|
||||
return
|
||||
}
|
||||
svc.Errorf("task[%s] stream error: %v", taskId.Hex(), res.err)
|
||||
return
|
||||
}
|
||||
svc.processStreamMessage(taskId, res.msg)
|
||||
case <-ctx.Done():
|
||||
cancel()
|
||||
svc.Warnf("task[%s] stream receive timeout", taskId.Hex())
|
||||
// Continue loop to try again
|
||||
case <-stopCh:
|
||||
cancel()
|
||||
return
|
||||
case <-svc.ctx.Done():
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (svc *Service) processStreamMessage(taskId primitive.ObjectID, msg *grpc.TaskServiceSubscribeResponse) {
|
||||
switch msg.Code {
|
||||
case grpc.TaskServiceSubscribeCode_CANCEL:
|
||||
svc.Infof("task[%s] received cancel signal", taskId.Hex())
|
||||
// Handle cancel synchronously to avoid goroutine accumulation
|
||||
svc.handleCancel(msg, taskId)
|
||||
default:
|
||||
svc.Debugf("task[%s] received unknown stream message code: %v", taskId.Hex(), msg.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func (svc *Service) handleCancel(msg *grpc.TaskServiceSubscribeResponse, taskId primitive.ObjectID) {
|
||||
// validate task id
|
||||
if msg.TaskId != taskId.Hex() {
|
||||
svc.Errorf("task[%s] received cancel signal for another task[%s]", taskId.Hex(), msg.TaskId)
|
||||
return
|
||||
}
|
||||
|
||||
// cancel task
|
||||
err := svc.cancelTask(taskId, msg.Force)
|
||||
if err != nil {
|
||||
svc.Errorf("task[%s] failed to cancel: %v", taskId.Hex(), err)
|
||||
return
|
||||
}
|
||||
svc.Infof("task[%s] cancelled", taskId.Hex())
|
||||
|
||||
// set task status as "cancelled"
|
||||
t, err := svc.GetTaskById(taskId)
|
||||
if err != nil {
|
||||
svc.Errorf("task[%s] failed to get task: %v", taskId.Hex(), err)
|
||||
return
|
||||
}
|
||||
t.Status = constants.TaskStatusCancelled
|
||||
err = svc.UpdateTask(t)
|
||||
if err != nil {
|
||||
svc.Errorf("task[%s] failed to update task: %v", taskId.Hex(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func (svc *Service) cancelTask(taskId primitive.ObjectID, force bool) (err error) {
|
||||
r, err := svc.getRunner(taskId)
|
||||
if err != nil {
|
||||
// Runner not found, task might already be finished
|
||||
svc.Warnf("runner not found for task[%s]: %v", taskId.Hex(), err)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Attempt cancellation with timeout
|
||||
cancelCtx, cancelFunc := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancelFunc()
|
||||
|
||||
cancelDone := make(chan error, 1)
|
||||
go func() {
|
||||
cancelDone <- r.Cancel(force)
|
||||
}()
|
||||
|
||||
select {
|
||||
case err = <-cancelDone:
|
||||
if err != nil {
|
||||
svc.Errorf("failed to cancel task[%s]: %v", taskId.Hex(), err)
|
||||
// If cancellation failed and force is not set, try force cancellation
|
||||
if !force {
|
||||
svc.Warnf("escalating to force cancellation for task[%s]", taskId.Hex())
|
||||
return svc.cancelTask(taskId, true)
|
||||
}
|
||||
return err
|
||||
}
|
||||
svc.Infof("task[%s] cancelled successfully", taskId.Hex())
|
||||
case <-cancelCtx.Done():
|
||||
svc.Errorf("timeout cancelling task[%s], removing runner from pool", taskId.Hex())
|
||||
// Remove runner from pool to prevent further issues
|
||||
svc.runners.Delete(taskId)
|
||||
return fmt.Errorf("task cancellation timeout")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// stopAllRunners gracefully stops all running tasks
|
||||
func (svc *Service) stopAllRunners() {
|
||||
svc.Infof("Stopping all running tasks...")
|
||||
|
||||
var runnerIds []primitive.ObjectID
|
||||
|
||||
// Collect all runner IDs
|
||||
svc.runners.Range(func(key, value interface{}) bool {
|
||||
if taskId, ok := key.(primitive.ObjectID); ok {
|
||||
runnerIds = append(runnerIds, taskId)
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
// Cancel all runners with bounded concurrency to prevent goroutine explosion
|
||||
const maxConcurrentCancellations = 10
|
||||
var wg sync.WaitGroup
|
||||
semaphore := make(chan struct{}, maxConcurrentCancellations)
|
||||
|
||||
for _, taskId := range runnerIds {
|
||||
wg.Add(1)
|
||||
|
||||
// Acquire semaphore to limit concurrent cancellations
|
||||
semaphore <- struct{}{}
|
||||
|
||||
go func(tid primitive.ObjectID) {
|
||||
defer func() {
|
||||
<-semaphore // Release semaphore
|
||||
wg.Done()
|
||||
if r := recover(); r != nil {
|
||||
svc.Errorf("stopAllRunners panic for task[%s]: %v", tid.Hex(), r)
|
||||
}
|
||||
}()
|
||||
|
||||
if err := svc.cancelTask(tid, false); err != nil {
|
||||
svc.Errorf("failed to cancel task[%s]: %v", tid.Hex(), err)
|
||||
// Force cancel after timeout
|
||||
time.Sleep(5 * time.Second)
|
||||
_ = svc.cancelTask(tid, true)
|
||||
}
|
||||
}(taskId)
|
||||
}
|
||||
|
||||
// Wait for all cancellations with timeout
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
svc.Infof("All tasks stopped gracefully")
|
||||
case <-time.After(30 * time.Second):
|
||||
svc.Warnf("Some tasks did not stop within timeout")
|
||||
}
|
||||
}
|
||||
|
||||
func (svc *Service) startStuckTaskCleanup() {
|
||||
svc.wg.Add(1) // Track this goroutine in the WaitGroup
|
||||
go func() {
|
||||
@@ -1067,222 +492,6 @@ func (svc *Service) checkAndCleanupStuckTasks() {
|
||||
}
|
||||
}
|
||||
|
||||
func NewStreamManager(service *Service) *StreamManager {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return &StreamManager{
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
service: service,
|
||||
messageQueue: make(chan *StreamMessage, 100), // Buffered channel for messages
|
||||
maxStreams: 50, // Limit concurrent streams
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *StreamManager) Start() {
|
||||
sm.wg.Add(2)
|
||||
go sm.messageProcessor()
|
||||
go sm.streamCleaner()
|
||||
}
|
||||
|
||||
func (sm *StreamManager) Stop() {
|
||||
sm.cancel()
|
||||
close(sm.messageQueue)
|
||||
|
||||
// Close all active streams
|
||||
sm.streams.Range(func(key, value interface{}) bool {
|
||||
if ts, ok := value.(*TaskStream); ok {
|
||||
ts.Close()
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
sm.wg.Wait()
|
||||
}
|
||||
|
||||
func (sm *StreamManager) AddTaskStream(taskId primitive.ObjectID) error {
|
||||
// Check if we're at the stream limit
|
||||
streamCount := 0
|
||||
sm.streams.Range(func(key, value interface{}) bool {
|
||||
streamCount++
|
||||
return true
|
||||
})
|
||||
|
||||
if streamCount >= sm.maxStreams {
|
||||
return fmt.Errorf("stream limit reached (%d)", sm.maxStreams)
|
||||
}
|
||||
|
||||
// Create new stream
|
||||
stream, err := sm.service.subscribeTask(taskId)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to subscribe to task stream: %v", err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(sm.ctx)
|
||||
taskStream := &TaskStream{
|
||||
taskId: taskId,
|
||||
stream: stream,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
lastActive: time.Now(),
|
||||
}
|
||||
|
||||
sm.streams.Store(taskId, taskStream)
|
||||
|
||||
// Start listening for messages in a single goroutine per stream
|
||||
sm.wg.Add(1)
|
||||
go sm.streamListener(taskStream)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (sm *StreamManager) RemoveTaskStream(taskId primitive.ObjectID) {
|
||||
if value, ok := sm.streams.LoadAndDelete(taskId); ok {
|
||||
if ts, ok := value.(*TaskStream); ok {
|
||||
ts.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *StreamManager) streamListener(ts *TaskStream) {
|
||||
defer sm.wg.Done()
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
sm.service.Errorf("stream listener panic for task[%s]: %v", ts.taskId.Hex(), r)
|
||||
}
|
||||
ts.Close()
|
||||
sm.streams.Delete(ts.taskId)
|
||||
}()
|
||||
|
||||
sm.service.Debugf("stream listener started for task[%s]", ts.taskId.Hex())
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ts.ctx.Done():
|
||||
sm.service.Debugf("stream listener stopped for task[%s]", ts.taskId.Hex())
|
||||
return
|
||||
case <-sm.ctx.Done():
|
||||
return
|
||||
default:
|
||||
msg, err := ts.stream.Recv()
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) {
|
||||
sm.service.Debugf("stream EOF for task[%s]", ts.taskId.Hex())
|
||||
return
|
||||
}
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return
|
||||
}
|
||||
sm.service.Errorf("stream error for task[%s]: %v", ts.taskId.Hex(), err)
|
||||
return
|
||||
}
|
||||
|
||||
// Update last active time
|
||||
ts.mu.Lock()
|
||||
ts.lastActive = time.Now()
|
||||
ts.mu.Unlock()
|
||||
|
||||
// Send message to processor
|
||||
select {
|
||||
case sm.messageQueue <- &StreamMessage{
|
||||
taskId: ts.taskId,
|
||||
msg: msg,
|
||||
err: nil,
|
||||
}:
|
||||
case <-ts.ctx.Done():
|
||||
return
|
||||
case <-sm.ctx.Done():
|
||||
return
|
||||
default:
|
||||
sm.service.Warnf("message queue full, dropping message for task[%s]", ts.taskId.Hex())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *StreamManager) messageProcessor() {
|
||||
defer sm.wg.Done()
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
sm.service.Errorf("message processor panic: %v", r)
|
||||
}
|
||||
}()
|
||||
|
||||
sm.service.Debugf("stream message processor started")
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-sm.ctx.Done():
|
||||
sm.service.Debugf("stream message processor shutting down")
|
||||
return
|
||||
case msg, ok := <-sm.messageQueue:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
sm.processMessage(msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *StreamManager) processMessage(streamMsg *StreamMessage) {
|
||||
if streamMsg.err != nil {
|
||||
sm.service.Errorf("stream message error for task[%s]: %v", streamMsg.taskId.Hex(), streamMsg.err)
|
||||
return
|
||||
}
|
||||
|
||||
// Process the actual message
|
||||
sm.service.processStreamMessage(streamMsg.taskId, streamMsg.msg)
|
||||
}
|
||||
|
||||
func (sm *StreamManager) streamCleaner() {
|
||||
defer sm.wg.Done()
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
sm.service.Errorf("stream cleaner panic: %v", r)
|
||||
}
|
||||
}()
|
||||
|
||||
ticker := time.NewTicker(2 * time.Minute)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-sm.ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
sm.cleanupInactiveStreams()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sm *StreamManager) cleanupInactiveStreams() {
|
||||
now := time.Now()
|
||||
inactiveThreshold := 10 * time.Minute
|
||||
|
||||
sm.streams.Range(func(key, value interface{}) bool {
|
||||
taskId := key.(primitive.ObjectID)
|
||||
ts := value.(*TaskStream)
|
||||
|
||||
ts.mu.RLock()
|
||||
lastActive := ts.lastActive
|
||||
ts.mu.RUnlock()
|
||||
|
||||
if now.Sub(lastActive) > inactiveThreshold {
|
||||
sm.service.Debugf("cleaning up inactive stream for task[%s]", taskId.Hex())
|
||||
sm.RemoveTaskStream(taskId)
|
||||
}
|
||||
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
func (ts *TaskStream) Close() {
|
||||
ts.cancel()
|
||||
if ts.stream != nil {
|
||||
_ = ts.stream.CloseSend()
|
||||
}
|
||||
}
|
||||
|
||||
func newTaskHandlerService() *Service {
|
||||
// service
|
||||
svc := &Service{
|
||||
|
||||
Reference in New Issue
Block a user