89109d2
use std::sync::Arc;
use std::time::Duration;
use config::import::ImportConfig;
use export::{AppDatabase, AppRedis};
use kenshell::KenShellClient;
use kenshell::proto::common::v1::RepositoryRef;
use kenshell::proto::v1 as pb;
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,
ks: KenShellClient,
config: ImportConfig,
cancel: CancellationToken,
}
impl ImportWorker {
#[must_use]
pub fn new(
db: Arc<AppDatabase>,
redis: AppRedis,
ks: KenShellClient,
config: ImportConfig,
cancel: CancellationToken,
) -> Self {
Self {
db,
redis,
ks,
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, fields(job_id = %row.id))]
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 space = git::share::access::space_for_dir(self.db.read(), repo.dir_id)
.await
.map_err(|error| ImportError::Internal(format!("resolve space: {error}")))?
.ok_or_else(|| ImportError::Internal("owning space not found".to_string()))?;
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 auth = match row.token_enc.as_deref() {
Some(stored) => Some(pb::RemoteAuth {
username: options
.credential_username
.clone()
.unwrap_or_else(|| provider.default_credential_username().to_string()),
token: decrypt_token(&self.config.token_secret, stored)?,
authorization_header: String::new(),
ssh_key: String::new(),
known_hosts: String::new(),
resolved_address: String::new(),
}),
None => None,
};
let mut remote = self.ks.remote();
let job = remote
.start_remote_import(pb::StartRemoteImportRequest {
repository: Some(RepositoryRef {
org_name: space.lower_name.clone(),
repo_name: repo.name.clone(),
}),
url: row.clone_url.clone(),
auth,
actor: None,
idempotency_key: row.id.to_string(),
})
.await
.map_err(|status| ImportError::Internal(status.message().to_string()))?
.into_inner()
.job
.ok_or_else(|| ImportError::Internal("import job missing".to_string()))?;
self.check_canceled(row.id).await?;
self.set_stage(row.id, ImportStage::Sync).await?;
let result = self.wait_import_job(&job.id, row.id).await?;
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 = if result.default_branch.is_empty() {
fresh
.default_branch
.strip_prefix("refs/heads/")
.unwrap_or(&fresh.default_branch)
.to_string()
} else {
result
.default_branch
.strip_prefix("refs/heads/")
.unwrap_or(&result.default_branch)
.to_string()
};
let stats = ImportStatsDto {
branches: i64::from(fresh.nums_branches),
tags: i64::from(fresh.nums_tags),
commits: i64::from(fresh.nums_commits).max(result.commit_count as i64),
size_bytes: fresh.git_size,
default_branch: Some(default_branch),
warnings: Vec::new(),
};
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 wait_import_job(
&self,
job_id: &str,
local_id: Uuid,
) -> Result<pb::RemoteJobResult, ImportError> {
let mut client = self.ks.job();
loop {
self.check_canceled(local_id).await.map_err(|error| {
if matches!(error, ImportError::Canceled) {
error
} else {
ImportError::Internal(error.to_string())
}
})?;
let job = client
.get_job(pb::GetJobRequest {
id: job_id.to_string(),
})
.await
.map_err(|status| ImportError::Internal(status.message().to_string()))?
.into_inner()
.job
.ok_or_else(|| ImportError::Internal("import job missing".to_string()))?;
match pb::JobState::try_from(job.state).unwrap_or(pb::JobState::Unspecified) {
pb::JobState::Succeeded => {
return match job.result {
Some(pb::job::Result::Remote(remote)) => Ok(remote),
_ => Err(ImportError::Internal(
"import job returned no result".to_string(),
)),
};
}
pb::JobState::Failed | pb::JobState::Canceled => {
return Err(ImportError::Internal(if job.error_message.is_empty() {
"import job failed".to_string()
} else {
job.error_message
}));
}
_ => {}
}
tokio::time::sleep(Duration::from_millis(500)).await;
}
}
}