zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-15 05:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: use NATS as the message queue (#15327)

Browse Source 
### What problem does this PR solve?

```
RAGFlow(admin)> mq publish 'msg2';
SUCCESS
RAGFlow(admin)> mq publish 'msg3';
SUCCESS
RAGFlow(admin)> mq list;
+---------+---------------+
| message | subject       |
+---------+---------------+
| msg1    | tasks.RAGFLOW |
| msg2    | tasks.RAGFLOW |
| msg3    | tasks.RAGFLOW |
+---------+---------------+
RAGFlow(admin)> mq pull 2;
+---------+---------------+
| message | subject       |
+---------+---------------+
| msg1    | tasks.RAGFLOW |
| msg2    | tasks.RAGFLOW |
+---------+---------------+
RAGFlow(admin)> mq pull noack;
+---------+---------------+
| message | subject       |
+---------+---------------+
| abc     | tasks.RAGFLOW |
+---------+---------------+
RAGFlow(admin)> mq show
+-------------------+----------------+--------+---------------+---------------+-------------------+---------------+
| ack_pending_count | consumer_count | memory | message_count | pending_count | redelivered_count | waiting_count |
+-------------------+----------------+--------+---------------+---------------+-------------------+---------------+
| 2                 | 1              | 0      | 2             | 0             | 1                 | 0             |
+-------------------+----------------+--------+---------------+---------------+-------------------+---------------+

RAGFlow(admin)> list ingestors;
+--------------+-------------------------------------------+--------+
| host         | name                                      | status |
+--------------+-------------------------------------------+--------+
| 192.168.1.38 | ingestor-8f0e4bd5650a4ac58b0151969fbf6935 | alive  |
+--------------+-------------------------------------------+--------+

RAGFlow(admin)> list ingestion tasks;
+----------------------------------+----------------------------------+-----------+------+-------------+----------------------------------+
| document_id                      | id                               | status    | step | user        | user_id                          |
+----------------------------------+----------------------------------+-----------+------+-------------+----------------------------------+
| ffe64fae423411f1a2d938a74640adcc | 90d3d0f6528941c1ac8eb0360effccc4 | COMPLETED | 5    | aaa@aaa.com | 2ba4881420fa11f19e9c38a74640adcc |
+----------------------------------+----------------------------------+-----------+------+-------------+----------------------------------+

RAGFlow(admin)> remove ingestion tasks '90d3d0f6528941c1ac8eb0360effccc4';
+---------+----------------------------------+
| delete  | task_id                          |
+---------+----------------------------------+
| success | 90d3d0f6528941c1ac8eb0360effccc4 |
+---------+----------------------------------+

RAGFlow(admin)> stop ingestion tasks 'e89e20d9a25848a1b79bd9345ddbfe1d';
+----------+----------------------------------+
| status   | task_id                          |
+----------+----------------------------------+
| STOPPING | e89e20d9a25848a1b79bd9345ddbfe1d |
+----------+----------------------------------+

# Publish a message
RAGFlow(admin)> mq publish 'cdd';
SUCCESS

# List current tasks in the message queue
RAGFlow(admin)> mq list
+----------------------------------+---------------+
| message                          | subject       |
+----------------------------------+---------------+
| 7ce392a3c1624cd2be4b5276e8825059 | tasks.RAGFLOW |
+----------------------------------+---------------+

# Consume a task from the message queue
RAGFlow(admin)> mq pull
+------+-----+----------------+
| ack  | id  | type           |
+------+-----+----------------+
| true | cdd | ingestion_test |
+------+-----+----------------+

# User mode
# List ingestion tasks, followed by dataset id
RAGFlow(user)> list ingestion tasks from '0abe79f9423311f1ad8d38a74640adcc';
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| create_date               | create_time   | dataset_id                       | document_id                      | id                               | schema | status    | update_date               | update_time   | user_id                          |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| 2026-05-30T20:21:06+08:00 | 1780143666289 | 0abe79f9423311f1ad8d38a74640adcc | ffe64fae423411f1a2d938a74640adcc | 8d758cd14a8b4ba8ab505003fb52017d |        | COMPLETED | 2026-05-30T20:21:26+08:00 | 1780143686431 | 2ba4881420fa11f19e9c38a74640adcc |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+

RAGFlow(user)> list ingestion tasks;
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| create_date               | create_time   | dataset_id                       | document_id                      | id                               | schema | status    | update_date               | update_time   | user_id                          |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| 2026-06-02T19:02:31+08:00 | 1780398151417 | 0abe79f9423311f1ad8d38a74640adcc | ffe64fae423411f1a2d938a74640adcc | e89e20d9a25848a1b79bd9345ddbfe1d |        | COMPLETED | 2026-06-02T19:02:52+08:00 | 1780398172208 | 2ba4881420fa11f19e9c38a74640adcc |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+

# Create an ingestion task
# First argument is document id, second argument is dataset id
RAGFlow(user)> start ingestion 'ffe64fae423411f1a2d938a74640adcc' from '0abe79f9423311f1ad8d38a74640adcc';
+----------------------------------+-------------------------------------------+
| document_id                      | result                                    |
+----------------------------------+-------------------------------------------+
| ffe64fae423411f1a2d938a74640adcc | task_id: 8d758cd14a8b4ba8ab505003fb52017d |
+----------------------------------+-------------------------------------------+

# Pause an ingestion task, first argument is ingestion id
RAGFlow(user)> stop ingestion '8d758cd14a8b4ba8ab505003fb52017d';
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| create_date               | create_time   | dataset_id                       | document_id                      | id                               | schema | status    | update_date               | update_time   | user_id                          |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+
| 2026-05-30T20:21:06+08:00 | 1780143666289 | 0abe79f9423311f1ad8d38a74640adcc | ffe64fae423411f1a2d938a74640adcc | 8d758cd14a8b4ba8ab505003fb52017d |        | COMPLETED | 2026-05-30T20:21:26+08:00 | 1780143686431 | 2ba4881420fa11f19e9c38a74640adcc |
+---------------------------+---------------+----------------------------------+----------------------------------+----------------------------------+--------+-----------+---------------------------+---------------+----------------------------------+

# Delete an ingestion task
RAGFlow(api/default)> remove ingestion tasks 'f366450a27d54677aec1c7090add30f0';
+---------+----------------------------------+
| remove  | task_id                          |
+---------+----------------------------------+
| success | f366450a27d54677aec1c7090add30f0 |
+---------+----------------------------------+

```

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-06-12 14:56:44 +08:00
committed by GitHub
parent 30724140d2
commit e96bc37d06
45 changed files with 2911 additions and 2337 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
internal/admin/handler.go
View File

@@ -17,12 +17,14 @@
package admin
import (
"encoding/json"
"errors"
"fmt"
"net/http"
"ragflow/internal/cache"
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/engine"
"ragflow/internal/server"
"ragflow/internal/service"
"ragflow/internal/utility"
@@ -203,15 +205,6 @@ func (h *Handler) AuthCheck(c *gin.Context) {
successNoData(c, "Admin is authorized")
}
// ListTasks handle list tasks
func (h *Handler) ListTasks(c *gin.Context) {
tasks, err := h.service.ListTasks()
if err != nil {
errorResponse(c, err.Error(), 500)
}
success(c, tasks, "Get all tasks")
}
// ListUsers handle list users
func (h *Handler) ListUsers(c *gin.Context) {
users, err := h.service.ListUsers()
@@ -261,7 +254,7 @@ func (h *Handler) GetUser(c *gin.Context) {
userDetails, err := h.service.GetUserDetails(username)
if err != nil {
if errors.Is(err, ErrUserNotFound) {
if errors.Is(err, common.ErrUserNotFound) {
errorResponse(c, "User not found", 404)
return
}
@@ -1256,57 +1249,206 @@ func (h *Handler) SetLogLevel(c *gin.Context) {
success(c, gin.H{"level": req.Level}, "Log level updated successfully")
}
type StartIngestionTaskRequest struct {
FileURI string `json:"uri" binding:"required"`
From string `json:"from" binding:"required"`
func (h *Handler) ListMessagesFromQueue(c *gin.Context) {
msgQueueEngine := engine.GetMessageQueueEngine()
messages, err := msgQueueEngine.ListMessages("ingestion", false)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
var result []map[string]string
for _, message := range messages {
var taskMessage common.TaskMessage
err = json.Unmarshal([]byte(message["message"]), &taskMessage)
if err != nil {
return
}
result = append(result, map[string]string{
"subject": message["subject"],
"id": taskMessage.TaskID,
"type": taskMessage.TaskType,
})
}
success(c, result, "List messages from queue successfully")
}
func (h *Handler) StartIngestionTask(c *gin.Context) {
var req StartIngestionTaskRequest
type PublishMessageToQueueRequest struct {
Message string `json:"message" binding:"required"`
}
func (h *Handler) PublishMessageToQueue(c *gin.Context) {
var req PublishMessageToQueueRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResponse(c, "file uri and from is required", 400)
errorResponse(c, "message is required", 400)
return
}
taskID := common.GenerateUUID()
ingestionManager.SubmitTask(&common.TaskAssignment{
TaskId: taskID,
TaskType: "start_ingestion_task",
Config: req.FileURI,
ComeFrom: req.From,
})
taskMessage := common.TaskMessage{
TaskID: req.Message,
TaskType: common.TaskTypeIngestionTest,
}
success(c, gin.H{"task_id": taskID}, "Send task for ingestion successfully")
// convert task
taskMessageStr, err := json.Marshal(taskMessage)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
msgQueueEngine := engine.GetMessageQueueEngine()
err = msgQueueEngine.PublishTask("tasks.RAGFLOW", taskMessageStr)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, nil, "Publish message successfully")
}
type PullMessageFromQueueRequest struct {
MessageCount int `json:"message_count" binding:"required"`
AckPolicy string `json:"ack_policy" binding:"required"`
}
func (h *Handler) PullMessageFromQueue(c *gin.Context) {
var req PullMessageFromQueueRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResponse(c, fmt.Sprintf("message count and ack_policy are required, error: %s", err.Error()), 400)
return
}
msgQueueEngine := engine.GetMessageQueueEngine()
err := msgQueueEngine.InitConsumer("tasks.RAGFLOW")
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
messages, err := msgQueueEngine.GetMessages(req.MessageCount)
var result []map[string]string
if req.AckPolicy == "ACK" {
for _, message := range messages {
taskMessage := message.GetMessage()
resultMessage := map[string]string{
"id": taskMessage.TaskID,
"type": taskMessage.TaskType,
}
err = message.Ack()
if err == nil {
resultMessage["ack"] = "true"
} else {
resultMessage["ack"] = "false"
}
result = append(result, resultMessage)
}
} else {
for _, message := range messages {
taskMessage := message.GetMessage()
resultMessage := map[string]string{
"id": taskMessage.TaskID,
"type": taskMessage.TaskType,
}
if err == nil {
resultMessage["nack"] = "true"
} else {
resultMessage["nack"] = "false"
}
result = append(result, resultMessage)
}
}
success(c, result, "Pull messages from queue successfully")
}
func (h *Handler) ShowMessageQueue(c *gin.Context) {
msgQueueEngine := engine.GetMessageQueueEngine()
result, err := msgQueueEngine.ShowMessageQueue()
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, result, "show message queue successfully")
}
type RemoveIngestionTaskRequest struct {
Tasks []string `json:"tasks" binding:"required"`
}
func (h *Handler) RemoveIngestionTasks(c *gin.Context) {
var req RemoveIngestionTaskRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResponse(c, "task id is required", 400)
return
}
tasks, err := h.service.RemoveIngestionTasks(req.Tasks)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, tasks, "Remove tasks successfully")
}
type StopIngestionTaskRequest struct {
TaskID string `json:"task_id" binding:"required"`
From string `json:"from" binding:"required"`
Tasks []string `json:"tasks" binding:"required"`
}
func (h *Handler) StopIngestionTask(c *gin.Context) {
func (h *Handler) StopIngestionTasks(c *gin.Context) {
var req StopIngestionTaskRequest
if err := c.ShouldBindJSON(&req); err != nil {
errorResponse(c, "task id and from is required", 400)
return
}
ingestionManager.SubmitTask(&common.TaskAssignment{
TaskId: req.TaskID,
TaskType: "cancel_ingestion_task",
ComeFrom: req.From,
})
tasks, err := h.service.StopIngestionTasks(req.Tasks)
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, gin.H{"task_id": req.TaskID}, "Cancel task successfully")
var result []map[string]string
for _, task := range tasks {
result = append(result, map[string]string{
"task_id": task.ID,
"status": task.Status,
})
}
success(c, result, "Stop tasks successfully")
}
func (h *Handler) ListIngestors(c *gin.Context) {
ingestionMgr := GetIngestionManager()
ingestors, err := ingestionMgr.ListIngestors()
// ListIngestionTasks
func (h *Handler) ListIngestionTasks(c *gin.Context) {
tasks, err := h.service.ListIngestionTasks()
if err != nil {
errorResponse(c, err.Error(), 500)
}
success(c, ingestors, "Get all tasks")
success(c, tasks, "Get all tasks")
}
func (h *Handler) ListIngestors(c *gin.Context) {
serverList := GlobalServerStore.ListInfos()
var ingestorResults []map[string]string
now := time.Now()
for _, ingestorServer := range serverList {
if ingestorServer.ServerType == common.ServerTypeIngestion {
ingestorResult := map[string]string{}
ingestorResult["name"] = ingestorServer.ServerName
ingestorResult["host"] = ingestorServer.Host
ingestorResult["status"] = ingestorServer.Version
if now.Sub(ingestorServer.Timestamp) < 30*time.Second {
ingestorResult["status"] = "alive"
} else {
ingestorResult["status"] = "timeout"
}
ingestorResults = append(ingestorResults, ingestorResult)
}
}
success(c, ingestorResults, "Get all tasks")
}
type ShutdownIngestorRequest struct {
@@ -1321,11 +1463,11 @@ func (h *Handler) ShutdownIngestor(c *gin.Context) {
}
taskID := common.GenerateUUID()
ingestionManager.SubmitTask(&common.TaskAssignment{
TaskId: taskID,
TaskType: "shutdown_ingestor",
AssignedTo: req.IngestorID,
})
//ingestionManager.SubmitTask(&common.TaskAssignment{
// TaskId: taskID,
// TaskType: "SHUTDOWN",
// AssignedTo: req.IngestorID,
//})
success(c, gin.H{"task_id": taskID, "ingestor_id": req.IngestorID}, "Shutdown ingestor")
}
@@ -1364,14 +1506,3 @@ func (h *Handler) Reports(c *gin.Context) {
responseWithCode(c, message, http.StatusOK, errCode)
}
// ListIngestionTasks
func (h *Handler) ListIngestionTasks(c *gin.Context) {
tasks, err := h.service.ListIngestionTasks()
if err != nil {
errorResponse(c, err.Error(), 400)
return
}
success(c, tasks, "")
}

587
internal/admin/ingestion_manager.go
View File

@@ -1,587 +0,0 @@
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package admin
import (
"context"
"fmt"
"net"
"sync"
"time"
"ragflow/internal/common"
"google.golang.org/grpc"
"google.golang.org/grpc/peer"
)
const heartbeatTimeout = 30 * time.Second
type IngestionManager struct {
common.UnimplementedIngestionManagerServer
mu sync.RWMutex
// Registered ingestion servers
ingestionServers map[string]*IngestorState // ingestor id -> ingestor id
taskStates map[string]*TaskState // task_id -> task state
// In-memory task queue
taskQueue chan *pendingTask
// Notifies that an ingestor slot may have freed up
slotFreed chan struct{}
grpcServer *grpc.Server // gRPC server instance for graceful shutdown via Stop()
ctx context.Context
cancel context.CancelFunc
}
type TaskState struct {
taskID string // same as task_id in database
status string // created, assigned, processing, completed, failed
comeFrom string // api server id
assignTo string // ingestor id
lastUpdate time.Time
startTime *time.Time
estimatedRemainingTime time.Duration // estimated cost in seconds to complete the task
errorMessage string
}
type IngestorState struct {
ID string
Info *common.RegisterInfo
LastHeartbeat time.Time
Stream common.IngestionManager_ActionServer
Status string // active, draining
Address string
ProcessID int64
cpuUsage float64
vmsUsage float64
rssUsage float64
}
type pendingTask struct {
Task *common.TaskAssignment
CreatedAt time.Time
}
var ingestionManager *IngestionManager
func GetIngestionManager() *IngestionManager {
return ingestionManager
}
func NewAdminServer() *IngestionManager {
ctx, cancel := context.WithCancel(context.Background())
ingestionManager = &IngestionManager{
taskStates: make(map[string]*TaskState),
ingestionServers: make(map[string]*IngestorState),
taskQueue: make(chan *pendingTask, 10000),
slotFreed: make(chan struct{}, 100),
ctx: ctx,
cancel: cancel,
}
go ingestionManager.dispatchLoop()
//go ingestionManager.heartbeatCheckLoop() no need to check heartbeat timeout
return ingestionManager
}
// Action handles the bidirectional streaming RPC from ingestion servers
func (s *IngestionManager) Action(stream common.IngestionManager_ActionServer) error {
var ingestionServerID string
var state *IngestorState
common.Info("New ingestion_server connection")
// Start receive goroutine
receiveErrorCH := make(chan error, 1)
go func() {
for {
msg, err := stream.Recv()
if err != nil {
receiveErrorCH <- err
return
}
s.handleMessage(stream, msg, &ingestionServerID, &state)
}
}()
// Start send goroutine: send tasks immediately when assigned to this ingestion_server
sendDone := make(chan struct{})
go func() {
defer close(sendDone)
for {
select {
case <-stream.Context().Done():
return
case <-s.ctx.Done():
return
}
}
}()
select {
case err := <-receiveErrorCH:
// Connection dropped, clean up
s.cleanupIngestionServer(ingestionServerID)
return err
case <-sendDone:
// Stream context canceled (client disconnect or server shutdown)
s.cleanupIngestionServer(ingestionServerID)
return nil
}
}
func (s *IngestionManager) handleMessage(
stream common.IngestionManager_ActionServer,
msg *common.IngestionMessage,
ingestionServerID *string,
state **IngestorState,
) {
switch msg.MessageType {
case "REGISTER":
s.handleRegister(stream, msg, ingestionServerID, state)
case "HEARTBEAT":
s.handleHeartbeat(msg, *ingestionServerID, *state)
case "TASK_RESULT":
s.handleTaskResult(msg, *ingestionServerID, *state)
case "TASK_PROGRESS":
s.handleTaskProgress(msg, *ingestionServerID, *state)
default:
common.Info(fmt.Sprintf("Unknown message type: %s", msg.MessageType))
err := stream.Send(&common.AdminMessage{
MessageType: "ERROR",
ErrorMessage: "unknown message type",
})
if err != nil {
common.Error("Fail to send unknown message", err)
return
}
}
}
func (s *IngestionManager) handleRegister(
stream common.IngestionManager_ActionServer,
msg *common.IngestionMessage,
ingestionServerID *string,
state **IngestorState,
) {
if msg.RegisterInfo == nil {
err := stream.Send(&common.AdminMessage{
MessageType: "ERROR",
ErrorMessage: "missing register info",
})
if err != nil {
common.Error("Fail to send missing register info", err)
return
}
return
}
peerHost, ok := peer.FromContext(stream.Context())
if !ok {
err := stream.Send(&common.AdminMessage{
MessageType: "ERROR",
ErrorMessage: "peer not found in context",
})
if err != nil {
common.Error("Fail to send 'peer not found' message", err)
return
}
return
}
clientAddr := peerHost.Addr.String()
*ingestionServerID = msg.IngestorId
*state = &IngestorState{
ID: msg.IngestorId,
Info: msg.RegisterInfo,
LastHeartbeat: time.Now().Truncate(time.Second),
Stream: stream,
Status: "active",
Address: clientAddr,
}
s.mu.Lock()
s.ingestionServers[*ingestionServerID] = *state
s.mu.Unlock()
err := stream.Send(&common.AdminMessage{
MessageType: "ACK",
AckInfo: &common.AckInfo{
TaskId: "",
Success: true,
Message: "registered successfully",
},
})
if err != nil {
common.Error("Fail to send ACK message", err)
return
}
common.Info(fmt.Sprintf("Ingestor %s registered, max_concurrency=%d, supported_types=%v",
*ingestionServerID, msg.RegisterInfo.MaxConcurrency, msg.RegisterInfo.SupportedDocTypes))
}
func (s *IngestionManager) handleHeartbeat(msg *common.IngestionMessage, ingestorID string, state *IngestorState) {
if state == nil {
return
}
state.LastHeartbeat = time.Now().Truncate(time.Second)
if msg.HeartbeatInfo != nil {
lastUpdateTime := time.Now().Truncate(time.Second)
s.mu.Lock()
ingestorState := s.ingestionServers[msg.IngestorId]
ingestorState.LastHeartbeat = lastUpdateTime
if ingestorState.Status == "timeout" {
ingestorState.Status = "active"
common.Info(fmt.Sprintf("Ingestor %s recovered from timeout, status set to active", msg.IngestorId))
}
ingestorState.ProcessID = msg.HeartbeatInfo.ProcessId
ingestorState.cpuUsage = float64(msg.HeartbeatInfo.CpuUsage)
ingestorState.vmsUsage = float64(msg.HeartbeatInfo.VmsUsage) / 1024 / 1024 // in MB
ingestorState.rssUsage = float64(msg.HeartbeatInfo.RssUsage) / 1024 / 1024 // in MB
// Delete expired terminal tasks from currentTasks
for _, taskID := range msg.HeartbeatInfo.DeleteTaskIds {
delete(s.taskStates, taskID)
}
for _, ingestorTaskState := range msg.HeartbeatInfo.TaskStates {
localTaskState := s.taskStates[ingestorTaskState.TaskId]
if localTaskState == nil {
startTime := time.Unix(0, ingestorTaskState.StartTime)
localTaskState = &TaskState{
taskID: ingestorTaskState.TaskId,
comeFrom: ingestorTaskState.ComeFrom,
startTime: &startTime,
}
}
localTaskState.estimatedRemainingTime = time.Duration(ingestorTaskState.EstimatedRemainingTimeSeconds)
localTaskState.lastUpdate = lastUpdateTime
localTaskState.status = ingestorTaskState.Status
localTaskState.errorMessage = ingestorTaskState.ErrorMessage
localTaskState.assignTo = msg.IngestorId
}
s.mu.Unlock()
common.Debug(fmt.Sprintf("Heartbeat from %s", ingestorID))
}
}
func (s *IngestionManager) handleTaskResult(msg *common.IngestionMessage, ingestorID string, state *IngestorState) {
if msg.TaskResult == nil {
return
}
result := msg.TaskResult
common.Info(fmt.Sprintf("Task result from %s: task=%s, status=%s, message=%s", ingestorID, result.TaskId, result.Status, result.ErrorMessage))
// Signal that a slot may have freed up for pending tasks
select {
case s.slotFreed <- struct{}{}:
default:
}
}
func (s *IngestionManager) handleTaskProgress(msg *common.IngestionMessage, ingestorID string, state *IngestorState) {
if msg.TaskProgress == nil {
return
}
progress := msg.TaskProgress
common.Info(fmt.Sprintf("Task progress from %s: task=%s, progress=%d%%, detail=%s",
ingestorID, progress.TaskId, progress.Progress, progress.Info))
}
// SubmitTask is for API Server to call (non-gRPC, for testing only)
func (s *IngestionManager) SubmitTask(task *common.TaskAssignment) {
s.taskQueue <- &pendingTask{
Task: task,
CreatedAt: time.Now().Truncate(time.Second),
}
common.Info(fmt.Sprintf("Task %s submitted to queue", task.TaskId))
// Wake up dispatchLoop if it's blocked waiting for a slot
select {
case s.slotFreed <- struct{}{}:
default:
}
}
// dispatchLoop pulls tasks from the queue and assigns them to available ingestors.
// Runs in a background goroutine.
func (s *IngestionManager) dispatchLoop() {
for {
select {
case <-s.ctx.Done():
return
case pending := <-s.taskQueue:
go s.tryAssign(pending.Task)
}
}
}
// heartbeatCheckLoop periodically checks all registered ingestors for heartbeat timeout.
// If an ingestor's LastHeartbeat is older than heartbeatTimeout, its status is set to "timeout".
func (s *IngestionManager) heartbeatCheckLoop() {
ticker := time.NewTicker(heartbeatTimeout / 3)
defer ticker.Stop()
for {
select {
case <-s.ctx.Done():
return
case <-ticker.C:
s.checkHeartbeats()
}
}
}
func (s *IngestionManager) checkHeartbeats() {
s.mu.Lock()
defer s.mu.Unlock()
now := time.Now().Truncate(time.Second)
for id, state := range s.ingestionServers {
if now.Sub(state.LastHeartbeat) > heartbeatTimeout {
if state.Status != "timeout" {
state.Status = "timeout"
common.Info(fmt.Sprintf("Ingestor %s heartbeat timeout, marked as timeout", id))
}
}
}
}
func (s *IngestionManager) SelectIngestorForTask(task *common.TaskAssignment) *IngestorState {
s.mu.Lock()
defer s.mu.Unlock()
switch task.TaskType {
case "start_ingestion_task":
for _, ingestor := range s.ingestionServers {
if ingestor.Status == "active" {
s.taskStates[task.TaskId] = &TaskState{
taskID: task.TaskId,
status: "DISPATCHED",
comeFrom: "CLI",
startTime: nil,
lastUpdate: time.Now().Truncate(time.Second),
assignTo: ingestor.ID,
}
return ingestor
}
}
case "cancel_ingestion_task":
taskState := s.taskStates[task.TaskId]
if taskState != nil {
switch taskState.status {
case "COMPLETED":
return nil
case "DISPATCHED":
{
taskState.status = "CANCELING"
return s.ingestionServers[taskState.assignTo]
}
default:
return s.ingestionServers[taskState.assignTo]
}
}
case "shutdown_ingestor":
return s.ingestionServers[task.AssignedTo]
}
return nil
}
// tryAssign repeatedly tries to find an available ingestor and assign the task.
// Blocks until either the task is assigned or the context is canceled.
func (s *IngestionManager) tryAssign(task *common.TaskAssignment) {
for {
target := s.SelectIngestorForTask(task)
if target != nil {
task.AssignedTo = target.ID
s.assignToIngestor(task, target)
return
}
if task.TaskType == "start_ingestion_task" {
// Receives a start ingestion task, save and change the states
s.mu.Lock()
s.taskStates[task.TaskId] = &TaskState{
taskID: task.TaskId,
status: "pending",
comeFrom: task.ComeFrom,
lastUpdate: time.Now().Truncate(time.Second),
startTime: nil,
}
s.mu.Unlock()
} else {
// shutdown ingestor or cancel task
common.Info("Task is completed, canceled, or ingestor is shutdown")
return
}
// No ingestor available, wait for a slot to free up
select {
case <-s.ctx.Done():
return
case <-s.slotFreed:
// A slot might be free, retry
case <-time.After(2 * time.Second):
// Periodic retry as fallback
}
}
}
func (s *IngestionManager) assignToIngestor(task *common.TaskAssignment, state *IngestorState) {
err := state.Stream.Send(&common.AdminMessage{
MessageType: "TASK_ASSIGNMENT",
TaskAssignment: task,
})
if err != nil {
common.Info(fmt.Sprintf("Failed to assign task %s to ingestor %s: %v", task.TaskId, state.ID, err))
// Re-queue the task
s.taskQueue <- &pendingTask{Task: task, CreatedAt: time.Now().Truncate(time.Second)}
return
}
common.Info(fmt.Sprintf("Assigned task %s to ingestion_server %s", task.TaskId, state.ID))
}
func (s *IngestionManager) cleanupIngestionServer(ingestorID string) {
s.mu.Lock()
defer s.mu.Unlock()
if ingestorID == "" {
// Client disconnected before REGISTER completed — nothing to clean up
common.Info("Unregistered ingestion server disconnected")
return
}
if _, exists := s.ingestionServers[ingestorID]; exists {
delete(s.ingestionServers, ingestorID)
common.Info(fmt.Sprintf("Ingestor %s cleaned up", ingestorID))
// Clean the tasks handled by this ingestor
var tasksToDelete []string
for _, taskState := range s.taskStates {
if taskState.assignTo == ingestorID {
tasksToDelete = append(tasksToDelete, taskState.taskID)
}
}
for _, taskID := range tasksToDelete {
delete(s.taskStates, taskID)
}
}
}
func (s *IngestionManager) ListIngestors() ([]map[string]interface{}, error) {
s.mu.Lock()
defer s.mu.Unlock()
var result []map[string]interface{}
for ingestorID, state := range s.ingestionServers {
var taskCount int64
for _, task := range s.taskStates {
if task.assignTo == ingestorID {
taskCount++
}
}
result = append(result, map[string]interface{}{
"id": ingestorID,
"name": state.Info.Name,
"address": state.Address,
"last_heartbeat": state.LastHeartbeat,
"task_count": taskCount,
"status": state.Status,
"cpu_usage": state.cpuUsage,
"rss_usage": state.rssUsage,
"vms_usage": state.vmsUsage,
"process_id": state.ProcessID,
})
}
return result, nil
}
func (s *IngestionManager) ListIngestionTasks() ([]map[string]interface{}, error) {
var result []map[string]interface{}
s.mu.Lock()
defer s.mu.Unlock()
for index, taskState := range s.taskStates {
common.Info(fmt.Sprintf("Task %s: %s", index, taskState.taskID))
result = append(result, map[string]interface{}{
"id": taskState.taskID,
"status": taskState.status,
"from": taskState.comeFrom,
"assign_to": taskState.assignTo,
"last_update": taskState.lastUpdate,
"start_time": taskState.startTime,
"ETA": taskState.estimatedRemainingTime,
"error": taskState.errorMessage,
})
}
return result, nil
}
// Start starts the admin service
func (s *IngestionManager) Start(port string) error {
lis, err := net.Listen("tcp", port)
if err != nil {
return err
}
s.grpcServer = grpc.NewServer()
common.RegisterIngestionManagerServer(s.grpcServer, s)
return s.grpcServer.Serve(lis)
}
// Stop gracefully shuts down the admin service
func (s *IngestionManager) Stop() {
common.Info("Stopping RAGFlow ingestion manager...")
// Notify all goroutines to exit
s.cancel()
// Gracefully stop gRPC server (stop accepting new connections, wait for in-flight requests)
if s.grpcServer != nil {
s.grpcServer.GracefulStop()
}
// Close the task queue
s.mu.Lock()
close(s.taskQueue)
s.mu.Unlock()
common.Info("RAGFlow ingestion manager stopped")
}

20
internal/admin/router.go
View File

@@ -46,6 +46,10 @@ func (r *Router) Setup(engine *gin.Engine) {
admin.POST("/reports", r.handler.Reports)
//admin.POST("/ingestion/tasks", r.handler.StartIngestionTask)
//admin.DELETE("/ingestion", r.handler.CancelIngestionTask) // cancel ingestion
//admin.GET("/ingestion/tasks", r.handler.ListIngestionTasks)
// Protected routes
protected := admin.Group("")
protected.Use(r.handler.AuthMiddleware())
@@ -55,9 +59,6 @@ func (r *Router) Setup(engine *gin.Engine) {
// Auth
protected.GET("/auth", r.handler.AuthCheck)
// Tasks
protected.GET("/tasks", r.handler.ListTasks)
// User management
protected.GET("/users", r.handler.ListUsers)
protected.POST("/users", r.handler.CreateUser)
@@ -137,12 +138,19 @@ func (r *Router) Setup(engine *gin.Engine) {
provider.GET("/:provider_name/models/:model_name", r.handler.ShowModel)
}
queue := protected.Group("/queue")
{
queue.GET("/", r.handler.ShowMessageQueue)
queue.POST("/messages", r.handler.PublishMessageToQueue)
queue.GET("/messages", r.handler.ListMessagesFromQueue)
queue.PUT("/messages", r.handler.PullMessageFromQueue)
}
protected.GET("/ingestors", r.handler.ListIngestors)
protected.DELETE("/ingestors", r.handler.ShutdownIngestor)
protected.POST("/ingestion", r.handler.StartIngestionTask) // start ingestion
protected.DELETE("/ingestion", r.handler.StopIngestionTask) // stop ingestion
protected.DELETE("/ingestion/tasks", r.handler.RemoveIngestionTasks)
protected.PUT("/ingestion/tasks", r.handler.StopIngestionTasks)
protected.GET("/ingestion/tasks", r.handler.ListIngestionTasks)
}
}

201
internal/admin/service.go
View File

@@ -29,6 +29,7 @@ import (
"ragflow/internal/cache"
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/engine"
"ragflow/internal/engine/elasticsearch"
"ragflow/internal/entity"
"ragflow/internal/server"
@@ -43,45 +44,49 @@ import (
// Service admin service layer
type Service struct {
userDAO *dao.UserDAO
licenseDAO *dao.LicenseDAO
timeRecordDAO *dao.TimeRecordDAO
systemSettingsDAO *dao.SystemSettingsDAO
tenantDAO *dao.TenantDAO
userTenantDAO *dao.UserTenantDAO
tenantLLMDAO *dao.TenantLLMDAO
fileDAO *dao.FileDAO
documentDAO *dao.DocumentDAO
taskDAO *dao.TaskDAO
kbDAO *dao.KnowledgebaseDAO
canvasDAO *dao.UserCanvasDAO
chatDAO *dao.ChatDAO
chatSessionDAO *dao.ChatSessionDAO
apiTokenDAO *dao.APITokenDAO
api4ConvDAO *dao.API4ConversationDAO
llmDAO *dao.LLMDAO
userDAO *dao.UserDAO
licenseDAO *dao.LicenseDAO
timeRecordDAO *dao.TimeRecordDAO
systemSettingsDAO *dao.SystemSettingsDAO
tenantDAO *dao.TenantDAO
userTenantDAO *dao.UserTenantDAO
tenantLLMDAO *dao.TenantLLMDAO
fileDAO *dao.FileDAO
documentDAO *dao.DocumentDAO
taskDAO *dao.TaskDAO
kbDAO *dao.KnowledgebaseDAO
canvasDAO *dao.UserCanvasDAO
chatDAO *dao.ChatDAO
chatSessionDAO *dao.ChatSessionDAO
apiTokenDAO *dao.APITokenDAO
api4ConvDAO *dao.API4ConversationDAO
llmDAO *dao.LLMDAO
ingestionTaskDAO *dao.IngestionTaskDAO
ingestionTaskLogDao *dao.IngestionTaskLogDAO
}
// NewService create admin service
func NewService() *Service {
return &Service{
userDAO: dao.NewUserDAO(),
licenseDAO: dao.NewLicenseDAO(),
timeRecordDAO: dao.NewTimeRecordDAO(),
systemSettingsDAO: dao.NewSystemSettingsDAO(),
tenantDAO: dao.NewTenantDAO(),
userTenantDAO: dao.NewUserTenantDAO(),
tenantLLMDAO: dao.NewTenantLLMDAO(),
fileDAO: dao.NewFileDAO(),
documentDAO: dao.NewDocumentDAO(),
taskDAO: dao.NewTaskDAO(),
kbDAO: dao.NewKnowledgebaseDAO(),
canvasDAO: dao.NewUserCanvasDAO(),
chatDAO: dao.NewChatDAO(),
chatSessionDAO: dao.NewChatSessionDAO(),
apiTokenDAO: dao.NewAPITokenDAO(),
api4ConvDAO: dao.NewAPI4ConversationDAO(),
llmDAO: dao.NewLLMDAO(),
userDAO: dao.NewUserDAO(),
licenseDAO: dao.NewLicenseDAO(),
timeRecordDAO: dao.NewTimeRecordDAO(),
systemSettingsDAO: dao.NewSystemSettingsDAO(),
tenantDAO: dao.NewTenantDAO(),
userTenantDAO: dao.NewUserTenantDAO(),
tenantLLMDAO: dao.NewTenantLLMDAO(),
fileDAO: dao.NewFileDAO(),
documentDAO: dao.NewDocumentDAO(),
taskDAO: dao.NewTaskDAO(),
kbDAO: dao.NewKnowledgebaseDAO(),
canvasDAO: dao.NewUserCanvasDAO(),
chatDAO: dao.NewChatDAO(),
chatSessionDAO: dao.NewChatSessionDAO(),
apiTokenDAO: dao.NewAPITokenDAO(),
api4ConvDAO: dao.NewAPI4ConversationDAO(),
llmDAO: dao.NewLLMDAO(),
ingestionTaskDAO: dao.NewIngestionTaskDAO(),
ingestionTaskLogDao: dao.NewIngestionTaskLogDAO(),
}
}
@@ -96,51 +101,99 @@ func (s *Service) Logout(user interface{}) error {
}
// ListTasks
func (s *Service) ListTasks() ([]map[string]interface{}, error) {
func (s *Service) ListIngestionTasks() ([]map[string]interface{}, error) {
//tasks, err := s.taskDAO.GetAllTasks()
//if err != nil {
// return nil, err
//}
//
//var result []map[string]interface{}
//for _, task := range tasks {
// // task.ChunkIDs is a string, delimiter is space, count the word count
// ChunkCount := strings.Count(*task.ChunkIDs, " ")
// result = append(result, map[string]interface{}{
// "id": task.ID,
// "task_type": task.TaskType,
// "document_id": task.DocID,
// "chunk_count": ChunkCount,
// "from_page": task.FromPage,
// "to_page": task.ToPage,
// "priority": task.Priority,
// "duration": task.ProcessDuration,
// "progress": task.Progress,
// //"message": *task.ProgressMsg,
// "retry_count": task.RetryCount,
// "digest": task.Digest,
// })
//}
ingestionMgr := GetIngestionManager()
ingestionTasks, err := ingestionMgr.ListIngestionTasks()
ingestionTasks, err := s.ingestionTaskDAO.GetAllTasks(0, 0)
if err != nil {
return nil, fmt.Errorf("fail to list ingestion tasks")
return nil, err
}
return ingestionTasks, nil
showTasks := []map[string]interface{}{}
for _, task := range ingestionTasks {
var user *entity.User
user, err = s.userDAO.GetByTenantID(task.UserID)
if err != nil {
return nil, err
}
//var document *entity.Document
//document, err = s.documentDAO.GetByID(task.DocumentID)
//if err != nil {
// return nil, err
//}
var showTask map[string]interface{}
var latestLog *entity.IngestionTaskLog
latestLog, err = s.ingestionTaskLogDao.LatestLogByTaskID(task.ID)
showTask = map[string]interface{}{
"id": task.ID,
"user_id": task.UserID,
"user": user.Email,
"document_id": task.DocumentID,
"status": task.Status,
}
if err == nil {
showTask = map[string]interface{}{
"id": task.ID,
"user_id": task.UserID,
"user": user.Email,
"document_id": task.DocumentID,
"status": task.Status,
"step": int(latestLog.Checkpoint["current_step"].(float64)),
}
}
showTasks = append(showTasks, showTask)
}
return showTasks, nil
}
func (s *Service) RemoveIngestionTasks(tasks []string) ([]map[string]string, error) {
var deletedTasks []map[string]string
for _, taskID := range tasks {
taskRecord := map[string]string{
"task_id": taskID,
}
_, err := s.ingestionTaskDAO.RemoveByAPIServerOrAdminServer(taskID, nil)
if err != nil {
taskRecord["remove"] = fmt.Sprintf("fail: %s", err.Error())
} else {
taskRecord["remove"] = "success"
}
deletedTasks = append(deletedTasks, taskRecord)
}
return deletedTasks, nil
}
func (s *Service) StopIngestionTasks(tasks []string) ([]*entity.IngestionTask, error) {
var taskResponses []*entity.IngestionTask
for _, taskID := range tasks {
task, err := s.ingestionTaskDAO.SetStoppingByAPIServer(taskID)
if err != nil {
return nil, err
}
if task.Status == common.STOPPING {
msgQueueEngine := engine.GetMessageQueueEngine()
err = msgQueueEngine.PublishTask("tasks.RAGFLOW", []byte(task.ID))
if err != nil {
return nil, err
}
}
taskResponses = append(taskResponses, task)
}
return taskResponses, nil
}
// GetUserByToken get user by access token
func (s *Service) GetUserByToken(token string) (*entity.User, error) {
user, err := s.userDAO.GetByAccessToken(token)
if err != nil {
return nil, ErrInvalidToken
return nil, common.ErrInvalidToken
}
if user.IsSuperuser == nil || !*user.IsSuperuser {
return nil, ErrNotAdmin
return nil, common.ErrNotAdmin
}
if user.IsActive != "1" {
@@ -477,7 +530,7 @@ func (s *Service) GetUserDetails(username string) (map[string]interface{}, error
var user entity.User
err := dao.DB.Where("email = ?", username).First(&user).Error
if err != nil {
return nil, ErrUserNotFound
return nil, common.ErrUserNotFound
}
return map[string]interface{}{
@@ -1050,11 +1103,11 @@ func (s *Service) ListServices() ([]map[string]interface{}, error) {
}
result = append(result, configDict)
}
}
id := len(result)
serverList := GlobalServerStore.ListInfos()
now := time.Now()
for _, serverStatus := range serverList {
serverItem := make(map[string]interface{})
serverItem["name"] = serverStatus.ServerName
@@ -1063,7 +1116,12 @@ func (s *Service) ListServices() ([]map[string]interface{}, error) {
id++
serverItem["host"] = serverStatus.Host
serverItem["port"] = serverStatus.Port
serverItem["status"] = "alive"
// the difference between now and serverStatus.Timestamp is less than 5 seconds, then the server is alive
if now.Sub(serverStatus.Timestamp) < 30*time.Second {
serverItem["status"] = "alive"
} else {
serverItem["status"] = "timeout"
}
result = append(result, serverItem)
}
return result, nil
@@ -1701,11 +1759,6 @@ func (s *Service) HandleHeartbeat(message *common.BaseMessage) (common.ErrorCode
return common.CodeLicenseValid, ""
}
func (s *Service) ListIngestionTasks() ([]map[string]interface{}, error) {
// TODO: Implement with sandbox manager
return []map[string]interface{}{}, nil
}
// InitDefaultAdmin initialize default admin user
// This matches Python's init_default_admin behavior
func (s *Service) InitDefaultAdmin() error {

12
internal/admin/state.go
View File

@@ -17,20 +17,11 @@
package admin
import (
"errors"
"ragflow/internal/common"
"sync"
"time"
)
// Service errors
var (
ErrInvalidToken = errors.New("invalid token")
ErrNotAdmin = errors.New("user is not admin")
ErrUserInactive = errors.New("user is inactive")
ErrUserNotFound = errors.New("user not found")
)
// API server state
// ServerStore is a thread-safe global server status storage
@@ -58,6 +49,9 @@ func (s *ServerStore) UpdateServerInfo(serverName string, status *common.BaseMes
s.servers[serverName] = status
return
case common.ServerTypeIngestion:
s.mu.Lock()
defer s.mu.Unlock()
s.servers[serverName] = status
return
}
}