wizardchen
035093dc35
- Increased the default worker pool size for the asynq server from 16 to 32 to improve performance during concurrent uploads. - Updated the .env.example and docker-compose.yml files to reflect the new default concurrency setting. - Enhanced the handling of FAQ entries in the tool results, ensuring distinct identification through standard questions, improving user experience and clarity in the output. - Refactored related components to utilize the new FAQ attributes, ensuring consistent display and interaction across the application.
323 lines
13 KiB
Go
323 lines
13 KiB
Go
package router
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"log"
|
||
"os"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Tencent/WeKnora/internal/application/service"
|
||
"github.com/Tencent/WeKnora/internal/logger"
|
||
"github.com/Tencent/WeKnora/internal/middleware/asynqdl"
|
||
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
|
||
"github.com/Tencent/WeKnora/internal/types"
|
||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
"github.com/hibiken/asynq"
|
||
"go.uber.org/dig"
|
||
)
|
||
|
||
type AsynqTaskParams struct {
|
||
dig.In
|
||
|
||
Server *asynq.Server
|
||
KnowledgeService interfaces.KnowledgeService
|
||
KnowledgeBaseService interfaces.KnowledgeBaseService
|
||
TagService interfaces.KnowledgeTagService
|
||
DataSourceService interfaces.DataSourceService
|
||
ChunkExtractor interfaces.TaskHandler `name:"chunkExtractor"`
|
||
DataTableSummary interfaces.TaskHandler `name:"dataTableSummary"`
|
||
ImageMultimodal interfaces.TaskHandler `name:"imageMultimodal"`
|
||
KnowledgePostProcess interfaces.TaskHandler `name:"knowledgePostProcess"`
|
||
WikiIngest interfaces.TaskHandler `name:"wikiIngest"`
|
||
DeadLetterRepo interfaces.TaskDeadLetterRepository
|
||
SpanTracker service.SpanTracker
|
||
}
|
||
|
||
// defaultRedisOpTimeout is the previous hard-coded read timeout. The 100ms
|
||
// floor was tight enough to cause spurious i/o timeout errors during bursty
|
||
// workloads (large batch uploads, multimodal counter DECRs under load), so we
|
||
// raise the default to 500ms while still allowing operators to tune via env.
|
||
const defaultRedisOpTimeoutMs = 500
|
||
|
||
// readRedisOpTimeoutMs reads WEKNORA_REDIS_OP_TIMEOUT_MS, falling back to
|
||
// defaultRedisOpTimeoutMs on missing/invalid input. Kept as a separate helper
|
||
// so both ReadTimeout and WriteTimeout share the same source of truth.
|
||
func readRedisOpTimeoutMs() int {
|
||
if v := strings.TrimSpace(os.Getenv("WEKNORA_REDIS_OP_TIMEOUT_MS")); v != "" {
|
||
if parsed, err := strconv.Atoi(v); err == nil && parsed > 0 {
|
||
return parsed
|
||
}
|
||
}
|
||
return defaultRedisOpTimeoutMs
|
||
}
|
||
|
||
func getAsynqRedisClientOpt() *asynq.RedisClientOpt {
|
||
db := 0
|
||
if dbStr := os.Getenv("REDIS_DB"); dbStr != "" {
|
||
if parsed, err := strconv.Atoi(dbStr); err == nil {
|
||
db = parsed
|
||
}
|
||
}
|
||
timeoutMs := readRedisOpTimeoutMs()
|
||
opt := &asynq.RedisClientOpt{
|
||
Addr: os.Getenv("REDIS_ADDR"),
|
||
Username: os.Getenv("REDIS_USERNAME"),
|
||
Password: os.Getenv("REDIS_PASSWORD"),
|
||
ReadTimeout: time.Duration(timeoutMs) * time.Millisecond,
|
||
// Writes are typically more sensitive to congestion than reads
|
||
// (RESP pipelining, BRPOPLPUSH on Asynq dequeue), so we keep
|
||
// WriteTimeout slightly larger to absorb head-of-line stalls.
|
||
WriteTimeout: time.Duration(timeoutMs*2) * time.Millisecond,
|
||
DB: db,
|
||
}
|
||
return opt
|
||
}
|
||
|
||
func NewAsyncqClient() (*asynq.Client, error) {
|
||
opt := getAsynqRedisClientOpt()
|
||
client := asynq.NewClient(opt)
|
||
err := client.Ping()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return client, nil
|
||
}
|
||
|
||
// wikiIngestRetryDelay is a fixed, short backoff for wiki ingest lock
|
||
// conflicts. Must be slightly longer than the active-lock TTL's worst-case
|
||
// "just got set" window so the retry is highly likely to succeed without
|
||
// burning through retries; but short enough that users don't feel the stall.
|
||
const wikiIngestRetryDelay = 15 * time.Second
|
||
|
||
// asynqRetryDelayFunc customizes per-task retry backoff.
|
||
//
|
||
// Default asynq backoff is exponential (≈10s, 40s, 90s, 2.5m, ...), which
|
||
// is appropriate for transient errors like remote HTTP failures. But for
|
||
// wiki ingest lock conflicts (ErrWikiIngestConcurrent), exponential
|
||
// backoff is harmful: a freshly orphaned lock expires in ≤60s, so a 15s
|
||
// fixed retry virtually guarantees the next attempt succeeds. Without
|
||
// this override, a crash-restart cycle can leave a KB unable to make
|
||
// progress for 7–10 minutes while the orphan lock expires AND the retry
|
||
// schedule catches up.
|
||
func asynqRetryDelayFunc(n int, e error, t *asynq.Task) time.Duration {
|
||
if errors.Is(e, service.ErrWikiIngestConcurrent) {
|
||
return wikiIngestRetryDelay
|
||
}
|
||
return asynq.DefaultRetryDelayFunc(n, e, t)
|
||
}
|
||
|
||
// defaultAsynqConcurrency is the worker pool size used when
|
||
// WEKNORA_ASYNQ_CONCURRENCY is unset. The asynq library defaults to
|
||
// runtime.NumCPU(), which under-provisions during batch document uploads:
|
||
// a single 4-core container can only process 4 documents in parallel even
|
||
// when 100 are queued, so the queue wait time eats into each task's
|
||
// DocumentProcessTimeout budget. 32 is a safer default for the I/O-bound
|
||
// nature of doc parsing (most time is spent in DocReader / embedding RPCs,
|
||
// not on local CPU).
|
||
const defaultAsynqConcurrency = 32
|
||
|
||
func NewAsynqServer(svc interfaces.SystemSettingService) *asynq.Server {
|
||
opt := getAsynqRedisClientOpt()
|
||
concurrency := defaultAsynqConcurrency
|
||
if svc != nil {
|
||
n := svc.GetInt(context.Background(), "asynq.concurrency", "WEKNORA_ASYNQ_CONCURRENCY", defaultAsynqConcurrency)
|
||
if n > 0 {
|
||
concurrency = int(n)
|
||
}
|
||
}
|
||
log.Printf("asynq server starting with concurrency=%d redis_op_timeout=%dms",
|
||
concurrency, readRedisOpTimeoutMs())
|
||
srv := asynq.NewServer(
|
||
opt,
|
||
asynq.Config{
|
||
Concurrency: concurrency,
|
||
Queues: map[string]int{
|
||
types.QueueCritical: 6, // Highest priority queue
|
||
types.QueueDefault: 3, // Default priority queue
|
||
types.QueueLow: 1, // Lowest priority queue
|
||
types.QueueMultimodal: 1, // Isolated lane for high-volume slow VLM image tasks
|
||
types.QueueGraph: 1, // Isolated lane for high-volume slow graph-extraction tasks
|
||
types.QueueQuestion: 1, // Isolated lane for high-volume slow question-generation tasks
|
||
},
|
||
RetryDelayFunc: asynqRetryDelayFunc,
|
||
},
|
||
)
|
||
return srv
|
||
}
|
||
|
||
func RunAsynqServer(params AsynqTaskParams) *asynq.ServeMux {
|
||
// Create a new mux and register all handlers
|
||
mux := asynq.NewServeMux()
|
||
|
||
// Install the dead-letter middleware FIRST so it sees the raw error
|
||
// returned by the handler, before any other middleware that might
|
||
// transform it. The middleware records one task_dead_letters row per
|
||
// task that exhausts its retry budget — operators can then SQL-query
|
||
// failures by task type, scope, or tenant without scraping logs.
|
||
// Best-effort: a DB failure is logged and swallowed; the original task
|
||
// error always propagates upstream to asynq for retry/archival.
|
||
//
|
||
// The callback flips Knowledge.parse_status to "failed" the moment a
|
||
// document-related task exhausts its retry budget. Without this hook,
|
||
// a permanently-failing task left its parent knowledge stranded in
|
||
// "processing" until housekeeping cron caught it minutes later — the
|
||
// UI signal users actually see.
|
||
knowledgeFailer := newDeadLetterKnowledgeFailer(params.KnowledgeService, params.SpanTracker)
|
||
mux.Use(asynqdl.MiddlewareWithCallback(params.DeadLetterRepo, knowledgeFailer))
|
||
|
||
// Install Langfuse middleware BEFORE handler registration so every task
|
||
// type is automatically wrapped. When Langfuse is disabled the middleware
|
||
// is a pass-through; when enabled it resumes the upstream HTTP trace (if
|
||
// the payload carries one) or opens a standalone trace, then wraps the
|
||
// handler execution in a SPAN so all child generations (embedding / VLM /
|
||
// chat / rerank / ASR) nest correctly in the Langfuse UI.
|
||
mux.Use(langfuse.AsynqMiddleware())
|
||
|
||
// Register extract handlers - router will dispatch to appropriate handler
|
||
mux.HandleFunc(types.TypeChunkExtract, params.ChunkExtractor.Handle)
|
||
mux.HandleFunc(types.TypeDataTableSummary, params.DataTableSummary.Handle)
|
||
|
||
// Register document processing handler
|
||
mux.HandleFunc(types.TypeDocumentProcess, params.KnowledgeService.ProcessDocument)
|
||
|
||
// Register manual knowledge processing handler (cleanup + re-indexing)
|
||
mux.HandleFunc(types.TypeManualProcess, params.KnowledgeService.ProcessManualUpdate)
|
||
|
||
// Register FAQ import handler (includes dry run mode)
|
||
mux.HandleFunc(types.TypeFAQImport, params.KnowledgeService.ProcessFAQImport)
|
||
|
||
// Register question generation handler
|
||
mux.HandleFunc(types.TypeQuestionGeneration, params.KnowledgeService.ProcessQuestionGeneration)
|
||
|
||
// Register summary generation handler
|
||
mux.HandleFunc(types.TypeSummaryGeneration, params.KnowledgeService.ProcessSummaryGeneration)
|
||
|
||
// Register KB clone handler
|
||
mux.HandleFunc(types.TypeKBClone, params.KnowledgeService.ProcessKBClone)
|
||
|
||
// Register knowledge move handler
|
||
mux.HandleFunc(types.TypeKnowledgeMove, params.KnowledgeService.ProcessKnowledgeMove)
|
||
|
||
// Register knowledge list delete handler
|
||
mux.HandleFunc(types.TypeKnowledgeListDelete, params.KnowledgeService.ProcessKnowledgeListDelete)
|
||
|
||
// Register index delete handler
|
||
mux.HandleFunc(types.TypeIndexDelete, params.TagService.ProcessIndexDelete)
|
||
|
||
// Register KB delete handler
|
||
mux.HandleFunc(types.TypeKBDelete, params.KnowledgeBaseService.ProcessKBDelete)
|
||
|
||
// Register image multimodal handler
|
||
mux.HandleFunc(types.TypeImageMultimodal, params.ImageMultimodal.Handle)
|
||
|
||
// Register knowledge post process handler
|
||
mux.HandleFunc(types.TypeKnowledgePostProcess, params.KnowledgePostProcess.Handle)
|
||
|
||
// Register data source sync handler
|
||
mux.HandleFunc(types.TypeDataSourceSync, params.DataSourceService.ProcessSync)
|
||
|
||
// Register wiki ingest handler
|
||
mux.HandleFunc(types.TypeWikiIngest, params.WikiIngest.Handle)
|
||
|
||
go func() {
|
||
// Start the server
|
||
if err := params.Server.Run(mux); err != nil {
|
||
log.Fatalf("could not run server: %v", err)
|
||
}
|
||
}()
|
||
return mux
|
||
}
|
||
|
||
// deadLetterKnowledgePayload extracts only the field we need from any
|
||
// document-related asynq payload. Kept narrow so we don't accidentally
|
||
// depend on the full payload schema and survive future field churn.
|
||
type deadLetterKnowledgePayload struct {
|
||
KnowledgeID string `json:"knowledge_id,omitempty"`
|
||
// Attempt threads through DocumentProcess / ManualProcess /
|
||
// KnowledgePostProcess payloads (added when span tracking shipped)
|
||
// — extracted here so the dead-letter callback can also close the
|
||
// matching root span as failed. Older in-flight payloads without
|
||
// this field decode as 0 and the tracker call no-ops.
|
||
Attempt int `json:"attempt,omitempty"`
|
||
}
|
||
|
||
// taskTypesAffectingKnowledgeStatus enumerates the asynq task types whose
|
||
// dead-letter event should flip the parent Knowledge to "failed". Only
|
||
// terminal task types are listed here:
|
||
//
|
||
// - TypeDocumentProcess: the entry point of the parsing pipeline.
|
||
// - TypeImageMultimodal: a single image hitting dead-letter would have
|
||
// been counted by isFinalAsynqAttempt (see image_multimodal.go), so
|
||
// the parent might still complete via remaining images. We DO NOT mark
|
||
// the parent failed for this case — finalize-on-last-attempt already
|
||
// ensures progress.
|
||
// - TypeKnowledgePostProcess: terminal stage; failure here strands the
|
||
// knowledge in "processing".
|
||
// - TypeManualProcess: same shape as DocumentProcess for re-indexing.
|
||
//
|
||
// Question/Summary generation are NOT included: they run after parse_status
|
||
// has already become "completed" and have their own status fields.
|
||
var taskTypesAffectingKnowledgeStatus = map[string]struct{}{
|
||
types.TypeDocumentProcess: {},
|
||
types.TypeKnowledgePostProcess: {},
|
||
types.TypeManualProcess: {},
|
||
}
|
||
|
||
// newDeadLetterKnowledgeFailer returns the callback wired into the asynq
|
||
// dead-letter middleware. When a document-related task exhausts its retry
|
||
// budget, this callback marks the corresponding Knowledge row as failed so
|
||
// the UI surfaces the error instead of a perpetual spinner.
|
||
//
|
||
// All work is best-effort: missing payload, missing knowledge_id, or DB
|
||
// errors are logged and swallowed. The dead-letter record is the source of
|
||
// truth — this is purely a UX shortcut so users don't wait for the
|
||
// housekeeping cron's next sweep.
|
||
func newDeadLetterKnowledgeFailer(ks interfaces.KnowledgeService, tracker service.SpanTracker) asynqdl.OnDeadLetter {
|
||
if ks == nil {
|
||
return nil
|
||
}
|
||
repo := ks.GetRepository()
|
||
if repo == nil {
|
||
return nil
|
||
}
|
||
return func(ctx context.Context, t *asynq.Task, taskErr error) {
|
||
if t == nil {
|
||
return
|
||
}
|
||
if _, ok := taskTypesAffectingKnowledgeStatus[t.Type()]; !ok {
|
||
return
|
||
}
|
||
var probe deadLetterKnowledgePayload
|
||
if err := json.Unmarshal(t.Payload(), &probe); err != nil || probe.KnowledgeID == "" {
|
||
return
|
||
}
|
||
errMsg := "task " + t.Type() + " exhausted retries: " + taskErr.Error()
|
||
// 8KB is the same cap the dead-letter row uses for last_error.
|
||
if len(errMsg) > 8192 {
|
||
errMsg = errMsg[:8192]
|
||
}
|
||
// Single UPDATE so we never end up with parse_status=failed but
|
||
// stale error_message (or vice versa) when the second write
|
||
// fails.
|
||
if err := repo.UpdateKnowledgeColumns(ctx, probe.KnowledgeID, map[string]interface{}{
|
||
"parse_status": types.ParseStatusFailed,
|
||
"error_message": errMsg,
|
||
}); err != nil {
|
||
logger.Warnf(ctx, "dead-letter callback: failed to mark knowledge %s as failed: %v", probe.KnowledgeID, err)
|
||
return
|
||
}
|
||
// Close the matching root span so the timeline stops showing
|
||
// "进行中" after dead-letter exhaustion. Best-effort: nil
|
||
// tracker / missing attempt / missing root all no-op cleanly.
|
||
if tracker != nil && probe.Attempt > 0 {
|
||
tracker.FinalizeAttempt(ctx, probe.KnowledgeID, probe.Attempt,
|
||
types.SpanStatusFailed, nil, "TASK_TIMEOUT", errMsg)
|
||
}
|
||
logger.Infof(ctx, "dead-letter callback: marked knowledge %s as failed (task=%s)", probe.KnowledgeID, t.Type())
|
||
}
|
||
}
|