f0a1552
ZhenYi
feat(chat): add chat module with room management, message handling, and session revocation
f0a1552
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::search::index_job;
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QueryOrder};
use time::OffsetDateTime;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};
use crate::outbox;
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.service.queue_initial().await {
warn!(error = %error, "search initial queue failed");
}
let delay = Duration::from_secs(self.service.config().poll_interval_secs);
loop {
tokio::select! {
_ = self.cancel.cancelled() => break,
_ = tokio::time::sleep(delay) => match self.run_once().await {
Ok(true) => {}
Ok(false) => {}
Err(error) => warn!(error = %error, "search worker iteration failed"),
}
}
}
info!("search worker stopped");
}
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 result = self.process(&job).await;
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, "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);
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
}
"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?;
}
Ok(())
}
}