d3d11f4
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use export::AppDatabase;
use models::issue::issue_item;
use models::pull_request::pull_requests;
use models::repo::repos;
use models::room::rooms;
use models::search::index_job;
use sea_orm::{
ColumnTrait, ConnectionTrait, DatabaseBackend, EntityTrait, QueryFilter, QueryOrder, Statement,
};
use time::OffsetDateTime;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};
use crate::outbox;
/// 回收停滞 processing 任务的等待分钟数。
const STALE_LOCK_MINUTES: i64 = 10;
use crate::{SearchError, SearchService};
#[derive(Clone)]
pub struct SearchWorker {
db: Arc<AppDatabase>,
service: SearchService,
repo_root: PathBuf,
cancel: CancellationToken,
}
impl SearchWorker {
pub fn new(
db: Arc<AppDatabase>,
service: SearchService,
repo_root: PathBuf,
cancel: CancellationToken,
) -> Self {
Self {
db,
service,
repo_root,
cancel,
}
}
pub async fn run(self) {
if let Err(error) = self.requeue_interrupted().await {
warn!(error = %error, "search recover interrupted jobs failed");
}
if let Err(error) = self.service.queue_initial().await {
warn!(error = %error, "search initial queue failed");
}
let delay = Duration::from_secs(self.service.config().poll_interval_secs.max(1));
let mut last_sweep = tokio::time::Instant::now();
loop {
tokio::select! {
_ = self.cancel.cancelled() => break,
_ = tokio::time::sleep(delay) => {
if last_sweep.elapsed() >= Duration::from_secs(60) {
if let Err(error) = self.requeue_interrupted().await {
warn!(error = %error, "search stale job sweep failed");
}
last_sweep = tokio::time::Instant::now();
}
while !self.cancel.is_cancelled() {
match self.run_once().await {
Ok(true) => {}
Ok(false) => break,
Err(error) => {
warn!(error = %error, "search worker iteration failed");
break;
}
}
}
}
}
}
info!("search worker stopped");
}
/// 进程被杀时任务会停在 processing，之后永远不会再被拾取；
/// 只回收锁定超过 STALE_LOCK_MINUTES 的记录，避免多实例互抢在途任务。
async fn requeue_interrupted(&self) -> Result<(), SearchError> {
let now = OffsetDateTime::now_utc();
let stale_before = now - time::Duration::minutes(STALE_LOCK_MINUTES);
// queued 去重索引下，回收时若同目标已有 queued 行，直接丢弃僵尸 processing 行，
// 由 queued 行按最新状态重建，避免 UPDATE 撞唯一索引。
self.db
.write()
.execute_raw(Statement::from_string(
DatabaseBackend::Postgres,
format!(
"DELETE FROM search_index_job p USING search_index_job q \
WHERE p.state = 'processing' AND p.locked_at < to_timestamp({}) \
AND q.state = 'queued' AND q.kind = p.kind \
AND q.target_id IS NOT DISTINCT FROM p.target_id",
stale_before.unix_timestamp()
),
))
.await?;
let recovered = index_job::Entity::update_many()
.col_expr(
index_job::Column::State,
sea_orm::sea_query::Expr::value("queued"),
)
.col_expr(
index_job::Column::LockedAt,
sea_orm::sea_query::Expr::value(Option::<OffsetDateTime>::None),
)
.col_expr(
index_job::Column::RunAfter,
sea_orm::sea_query::Expr::value(now),
)
.col_expr(
index_job::Column::UpdatedAt,
sea_orm::sea_query::Expr::value(now),
)
.filter(index_job::Column::State.eq("processing"))
.filter(index_job::Column::LockedAt.lt(stale_before))
.exec(self.db.write())
.await?;
if recovered.rows_affected > 0 {
info!(
count = recovered.rows_affected,
"search requeued interrupted index jobs"
);
}
Ok(())
}
pub async fn run_once(&self) -> Result<bool, SearchError> {
if !self.service.enabled() {
return Ok(false);
}
let now = OffsetDateTime::now_utc();
let Some(job) = index_job::Entity::find()
.filter(index_job::Column::State.eq("queued"))
.filter(index_job::Column::RunAfter.lte(now))
.order_by_asc(index_job::Column::CreatedAt)
.one(self.db.read())
.await?
else {
return Ok(false);
};
let claimed = index_job::Entity::update_many()
.col_expr(
index_job::Column::State,
sea_orm::sea_query::Expr::value("processing"),
)
.col_expr(
index_job::Column::LockedAt,
sea_orm::sea_query::Expr::value(now),
)
.col_expr(
index_job::Column::UpdatedAt,
sea_orm::sea_query::Expr::value(now),
)
.filter(index_job::Column::Id.eq(job.id))
.filter(index_job::Column::State.eq("queued"))
.exec(self.db.write())
.await?;
if claimed.rows_affected == 0 {
return Ok(false);
}
let started = tokio::time::Instant::now();
let result = self.process(&job).await;
export::metrics::search_index_job(if result.is_ok() { "ok" } else { "error" });
match result {
Ok(()) => {
index_job::Entity::update_many()
.col_expr(
index_job::Column::State,
sea_orm::sea_query::Expr::value("succeeded"),
)
.col_expr(
index_job::Column::UpdatedAt,
sea_orm::sea_query::Expr::value(OffsetDateTime::now_utc()),
)
.filter(index_job::Column::Id.eq(job.id))
.exec(self.db.write())
.await?;
info!(
job_id = %job.id,
kind = %job.kind,
elapsed_ms = started.elapsed().as_millis() as u64,
lag_ms = (OffsetDateTime::now_utc() - job.created_at).whole_milliseconds().max(0) as u64,
"search index job completed"
);
}
Err(error) => {
let attempts = job.attempts + 1;
let retry_after = OffsetDateTime::now_utc()
+ time::Duration::seconds((attempts.min(8) * 5) as i64);
// 处理期间若同目标已有 queued 孪生行，直接丢弃本次失败行，
// 由孪生行按最新状态重试，避免回队时撞 queued 去重索引。
self.db
.write()
.execute_raw(Statement::from_string(
DatabaseBackend::Postgres,
format!(
"DELETE FROM search_index_job f USING search_index_job q \
WHERE f.id = '{}' AND q.state = 'queued' AND q.kind = f.kind \
AND q.target_id IS NOT DISTINCT FROM f.target_id",
job.id
),
))
.await?;
index_job::Entity::update_many()
.col_expr(
index_job::Column::State,
sea_orm::sea_query::Expr::value("queued"),
)
.col_expr(
index_job::Column::Attempts,
sea_orm::sea_query::Expr::value(attempts),
)
.col_expr(
index_job::Column::RunAfter,
sea_orm::sea_query::Expr::value(retry_after),
)
.col_expr(
index_job::Column::LockedAt,
sea_orm::sea_query::Expr::value(Option::<OffsetDateTime>::None),
)
.col_expr(
index_job::Column::LastError,
sea_orm::sea_query::Expr::value(error.to_string()),
)
.col_expr(
index_job::Column::UpdatedAt,
sea_orm::sea_query::Expr::value(OffsetDateTime::now_utc()),
)
.filter(index_job::Column::Id.eq(job.id))
.exec(self.db.write())
.await?;
error!(job_id = %job.id, kind = %job.kind, error = %error, "search index job failed and was requeued");
}
}
Ok(true)
}
async fn process(&self, job: &index_job::Model) -> Result<(), SearchError> {
match job.kind.as_str() {
"repo" => {
let repo_id = job
.target_id
.ok_or_else(|| SearchError::Invalid("repo job has no target".to_string()))?;
self.service.index_repo(repo_id, &self.repo_root).await
}
"issue" => {
let issue_id = job
.target_id
.ok_or_else(|| SearchError::Invalid("issue job has no target".to_string()))?;
self.service.index_issue(issue_id).await
}
"pull" => {
let pull_id = job
.target_id
.ok_or_else(|| SearchError::Invalid("pull job has no target".to_string()))?;
self.service.index_pull(pull_id).await
}
"message" => {
let message_id = job
.target_id
.ok_or_else(|| SearchError::Invalid("message job has no target".to_string()))?;
self.service.index_message(message_id).await
}
"room_messages" => {
let room_id = job
.target_id
.ok_or_else(|| SearchError::Invalid("room job has no target".to_string()))?;
self.service.index_room_messages(room_id).await
}
"all" => self.expand_all().await,
other => Err(SearchError::Invalid(format!(
"unknown index job kind `{other}`"
))),
}
}
async fn expand_all(&self) -> Result<(), SearchError> {
for repo in repos::Entity::find().all(self.db.read()).await? {
outbox::enqueue_repo(self.db.write(), repo.id).await?;
}
for issue in issue_item::Entity::find().all(self.db.read()).await? {
outbox::enqueue_issue(self.db.write(), issue.id).await?;
}
for pull in pull_requests::Entity::find().all(self.db.read()).await? {
outbox::enqueue_pull(self.db.write(), pull.id, pull.repo_id).await?;
}
// 含已停用房间：房间级任务先按 room_id 清空再回填，停用房间因此
// 在每次全量重建时顺带清掉残留文档。
for room in rooms::Entity::find().all(self.db.read()).await? {
outbox::enqueue_room_messages(self.db.write(), room.id).await?;
}
Ok(())
}
}