80f9e01
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use cache::CacheClient;
use config::git::SyncSettings;
use config::import::ImportConfig;
use export::{AppDatabase, AppRedis};
use git::service::clone::{GitCredential, mirror_fetch};
use git::sync::RepoSync;
use models::import::job;
use models::repo::repos;
use sea_orm::sea_query::Expr;
use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, QueryOrder, Set};
use time::OffsetDateTime;
use tokio::sync::Semaphore;
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use uuid::Uuid;
use crate::errors::ImportError;
use crate::identity;
use crate::providers::ImportProvider;
use crate::service::{ImportJobState, ImportOptions, ImportStage, ImportStatsDto, decrypt_token};
#[derive(Clone)]
pub struct ImportWorker {
db: Arc<AppDatabase>,
redis: AppRedis,
cache: Arc<CacheClient>,
repo_root: PathBuf,
settings: SyncSettings,
config: ImportConfig,
cancel: CancellationToken,
}
impl ImportWorker {
#[must_use]
pub fn new(
db: Arc<AppDatabase>,
redis: AppRedis,
cache: Arc<CacheClient>,
repo_root: PathBuf,
settings: SyncSettings,
config: ImportConfig,
cancel: CancellationToken,
) -> Self {
Self {
db,
redis,
cache,
repo_root,
settings,
config,
cancel,
}
}
pub async fn run(self) {
let permits = Arc::new(Semaphore::new(self.config.max_concurrent_jobs.max(1)));
let delay = Duration::from_secs(self.config.poll_interval_secs.max(1));
info!(
concurrency = self.config.max_concurrent_jobs.max(1),
"import worker started"
);
loop {
tokio::select! {
_ = self.cancel.cancelled() => break,
_ = tokio::time::sleep(delay) => {
if let Err(error) = self.reap_stale().await {
warn!(%error, "import stale job reap failed");
}
loop {
let Ok(permit) = permits.clone().try_acquire_owned() else {
break;
};
match self.claim().await {
Ok(Some(row)) => {
let worker = self.clone();
tokio::spawn(async move {
let _permit = permit;
worker.process(row).await;
});
}
Ok(None) => {
drop(permit);
break;
}
Err(error) => {
drop(permit);
warn!(%error, "import job claim failed");
break;
}
}
}
}
}
}
info!("import worker stopped");
}
async fn claim(&self) -> Result<Option<job::Model>, ImportError> {
let now = OffsetDateTime::now_utc();
let Some(row) = job::Entity::find()
.filter(job::Column::State.eq(ImportJobState::Queued.as_str()))
.filter(job::Column::RunAfter.lte(now))
.order_by_asc(job::Column::CreatedAt)
.one(self.db.read())
.await?
else {
return Ok(None);
};
let claimed = job::Entity::update_many()
.col_expr(
job::Column::State,
Expr::value(ImportJobState::Processing.as_str()),
)
.col_expr(job::Column::LockedAt, Expr::value(now))
.col_expr(job::Column::StartedAt, Expr::value(now))
.col_expr(job::Column::UpdatedAt, Expr::value(now))
.filter(job::Column::Id.eq(row.id))
.filter(job::Column::State.eq(ImportJobState::Queued.as_str()))
.exec(self.db.write())
.await?;
if claimed.rows_affected == 0 {
return Ok(None);
}
Ok(job::Entity::find_by_id(row.id).one(self.db.read()).await?)
}
async fn reap_stale(&self) -> Result<(), ImportError> {
let cutoff = OffsetDateTime::now_utc()
- time::Duration::seconds(
i64::try_from(self.config.stale_lock_secs.max(1)).unwrap_or(1800),
);
let now = OffsetDateTime::now_utc();
let result = job::Entity::update_many()
.col_expr(
job::Column::State,
Expr::value(ImportJobState::Queued.as_str()),
)
.col_expr(
job::Column::LockedAt,
Expr::value(Option::<OffsetDateTime>::None),
)
.col_expr(job::Column::RunAfter, Expr::value(now))
.col_expr(job::Column::UpdatedAt, Expr::value(now))
.filter(job::Column::State.eq(ImportJobState::Processing.as_str()))
.filter(job::Column::LockedAt.lt(cutoff))
.exec(self.db.write())
.await?;
if result.rows_affected > 0 {
warn!(count = result.rows_affected, "requeued stale import jobs");
}
Ok(())
}
fn spawn_heartbeat(&self, job_id: Uuid) -> tokio::task::JoinHandle<()> {
let db = self.db.clone();
let cancel = self.cancel.clone();
let interval = Duration::from_secs((self.config.stale_lock_secs.max(3) / 3).max(1));
tokio::spawn(async move {
loop {
tokio::select! {
_ = cancel.cancelled() => break,
_ = tokio::time::sleep(interval) => {
let now = OffsetDateTime::now_utc();
let result = job::Entity::update_many()
.col_expr(job::Column::LockedAt, Expr::value(now))
.col_expr(job::Column::UpdatedAt, Expr::value(now))
.filter(job::Column::Id.eq(job_id))
.filter(job::Column::State.eq(ImportJobState::Processing.as_str()))
.exec(db.write())
.await;
if let Err(error) = result {
warn!(%error, "import heartbeat failed");
}
}
}
}
})
}
#[tracing::instrument(name = "import.process", skip_all)]
async fn process(&self, row: job::Model) {
let heartbeat = self.spawn_heartbeat(row.id);
let outcome = self.run_job(&row).await;
heartbeat.abort();
export::metrics::import_job(match &outcome {
Ok(_) => "ok",
Err(ImportError::Canceled) => "canceled",
Err(_) => "failed",
});
let stage = ImportStage::parse(&row.stage);
match outcome {
Ok(stats) => {
if let Err(error) = self
.finish(
row.id,
ImportJobState::Succeeded,
ImportStage::Done,
None,
Some(stats),
)
.await
{
warn!(%error, job_id = %row.id, "failed to mark import job succeeded");
return;
}
info!(job_id = %row.id, repo = %row.repo_name, "repository import completed");
}
Err(ImportError::Canceled) => {
let _ = self
.finish(row.id, ImportJobState::Canceled, stage, None, None)
.await;
info!(job_id = %row.id, "repository import canceled");
}
Err(error) => {
let attempts = row.attempts + 1;
let exhausted = attempts >= self.config.max_attempts.max(1);
if error.is_permanent() || exhausted {
let _ = self
.finish(
row.id,
ImportJobState::Failed,
stage,
Some(error.to_string()),
None,
)
.await;
warn!(job_id = %row.id, attempts, %error, "repository import failed");
} else {
let _ = self.requeue(row.id, attempts, error.to_string()).await;
warn!(job_id = %row.id, attempts, %error, "repository import retried later");
}
}
}
}
async fn set_stage(&self, id: Uuid, stage: ImportStage) -> Result<(), ImportError> {
let now = OffsetDateTime::now_utc();
job::ActiveModel {
id: Set(id),
stage: Set(stage.as_str().to_string()),
updated_at: Set(now),
..Default::default()
}
.update(self.db.write())
.await?;
Ok(())
}
async fn finish(
&self,
id: Uuid,
state: ImportJobState,
stage: ImportStage,
last_error: Option<String>,
stats: Option<ImportStatsDto>,
) -> Result<(), ImportError> {
let now = OffsetDateTime::now_utc();
let mut model = job::ActiveModel {
id: Set(id),
state: Set(state.as_str().to_string()),
stage: Set(stage.as_str().to_string()),
last_error: Set(last_error),
token_enc: Set(None),
locked_at: Set(None),
finished_at: Set(Some(now)),
updated_at: Set(now),
..Default::default()
};
if let Some(stats) = stats {
model.stats = Set(serde_json::to_value(&stats).unwrap_or_default());
}
model.update(self.db.write()).await?;
Ok(())
}
async fn requeue(&self, id: Uuid, attempts: i32, error: String) -> Result<(), ImportError> {
let now = OffsetDateTime::now_utc();
let backoff = time::Duration::seconds(i64::from(attempts.clamp(1, 8)) * 15);
job::ActiveModel {
id: Set(id),
state: Set(ImportJobState::Queued.as_str().to_string()),
attempts: Set(attempts),
last_error: Set(Some(error)),
locked_at: Set(None),
run_after: Set(now + backoff),
updated_at: Set(now),
..Default::default()
}
.update(self.db.write())
.await?;
Ok(())
}
async fn check_canceled(&self, id: Uuid) -> Result<(), ImportError> {
let row = job::Entity::find_by_id(id).one(self.db.read()).await?;
match row {
None => Err(ImportError::Internal(
"import job disappeared while processing".to_string(),
)),
Some(row) if ImportJobState::parse(&row.state) == ImportJobState::Canceled => {
Err(ImportError::Canceled)
}
Some(_) => Ok(()),
}
}
async fn run_job(&self, row: &job::Model) -> Result<ImportStatsDto, ImportError> {
let repo_id = row.repo_id.ok_or_else(|| {
ImportError::Internal("import job is not attached to a repository".to_string())
})?;
let repo = repos::Entity::find_by_id(repo_id)
.one(self.db.read())
.await?
.ok_or_else(|| {
ImportError::Internal(format!("repository {repo_id} no longer exists"))
})?;
let provider = ImportProvider::parse(&row.provider)?;
let options: ImportOptions =
serde_json::from_value(row.options.clone()).unwrap_or_default();
self.check_canceled(row.id).await?;
self.set_stage(row.id, ImportStage::Clone).await?;
let credential = match row.token_enc.as_deref() {
Some(stored) => {
let token = decrypt_token(&self.config.token_secret, stored)?;
let username = options
.credential_username
.clone()
.unwrap_or_else(|| provider.default_credential_username().to_string());
Some(GitCredential::new(username, token))
}
None => None,
};
let dir = git::context::facade::repo_dir_for(&self.repo_root, repo_id);
let result = mirror_fetch(
"import::mirror",
&dir,
&row.clone_url,
credential.as_ref(),
self.config.clone_timeout_secs.max(1),
)
.await;
drop(credential);
result?;
self.check_canceled(row.id).await?;
self.set_stage(row.id, ImportStage::Sync).await?;
let sync = RepoSync::new(
self.db.clone(),
self.redis.clone(),
self.cache.clone(),
repo,
self.repo_root.clone(),
self.settings.clone(),
);
sync.sync()
.await
.map_err(|e| ImportError::Internal(format!("repository sync failed: {e}")))?;
self.check_canceled(row.id).await?;
self.set_stage(row.id, ImportStage::Finalize).await?;
let fresh = repos::Entity::find_by_id(repo_id)
.one(self.db.read())
.await?
.ok_or_else(|| {
ImportError::Internal(format!("repository {repo_id} disappeared during import"))
})?;
let default_branch = fresh
.default_branch
.strip_prefix("refs/heads/")
.unwrap_or(&fresh.default_branch)
.to_string();
let mut stats = ImportStatsDto {
branches: i64::from(fresh.nums_branches),
tags: i64::from(fresh.nums_tags),
commits: i64::from(fresh.nums_commits),
size_bytes: fresh.git_size,
default_branch: Some(default_branch),
warnings: Vec::new(),
};
if uses_lfs(
&dir,
&fresh.default_branch,
self.settings.git_command_timeout_secs,
)
.await
{
stats.warnings.push(
"源仓库使用 Git LFS，本次导入只搬运了指针文件，未搬迁 LFS 大对象。".to_string(),
);
}
if let Err(error) = identity::harvest(self.db.write(), provider, repo_id, &options).await {
warn!(%error, job_id = %row.id, "import identity harvest failed");
}
Ok(stats)
}
}
async fn uses_lfs(dir: &Path, default_branch: &str, timeout_secs: u64) -> bool {
let revision = default_branch
.strip_prefix("refs/heads/")
.unwrap_or(default_branch);
if revision.is_empty() {
return false;
}
let output = git::service::gitcmd::run_git_soft(
"import::lfs_probe",
dir,
vec![
"grep".to_string(),
"-l".to_string(),
"--fixed-strings".to_string(),
"filter=lfs".to_string(),
revision.to_string(),
"--".to_string(),
".gitattributes".to_string(),
],
&[],
timeout_secs.clamp(5, 60),
)
.await;
match output {
Ok(output) => output.status.success() && !output.stdout.is_empty(),
Err(_) => false,
}
}