Files
gitks/main.rs
T
zhenyi a40da90ef9 refactor(build): reformat code and add tonic health dependency
- Reformatted build script with proper indentation and line breaks
- Added tonic-health dependency to Cargo.toml and updated lock file
- Improved error handling in disk cache with concurrent deletion checks
- Refactored conditional chains using && and let expressions
- Reformatted struct initialization and function parameter lists
- Added proper spacing and alignment in language stats processing
- Improved assertion formatting in test cases
- Reorganized import statements and code layout in multiple files
- Updated metrics functions with better parameter handling and formatting
2026-06-11 13:56:15 +08:00

324 lines
10 KiB
Rust

use std::path::PathBuf;
use std::time::Duration;
use gitks::actor::init_actor_cluster;
use gitks::cluster::{ClusterConfig, ClusterManager};
use gitks::disk_cache::DiskCache;
use gitks::hooks::HookManager;
use gitks::metrics;
use gitks::server::{GitksService, serve};
use tracing_subscriber::EnvFilter;
use tracing_subscriber::fmt;
use tracing_subscriber::prelude::*;
const DEFAULT_HOST: &str = "0.0.0.0";
const DEFAULT_PORT: &str = "50051";
const DEFAULT_STORAGE_NAME: &str = "default";
fn env_or(key: &str, default: &str) -> String {
std::env::var(key).unwrap_or_else(|_| default.into())
}
fn env_bool(key: &str, default: bool) -> bool {
match std::env::var(key).as_deref() {
Ok("true" | "1" | "yes") => true,
Ok("false" | "0" | "no") => false,
_ => default,
}
}
fn env_u64(key: &str, default: u64) -> u64 {
std::env::var(key)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
fn init_tracing() -> Option<tracing_appender::non_blocking::WorkerGuard> {
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
let log_format = env_or("GITKS_LOG_FORMAT", "pretty");
let fmt_layer = match log_format.as_str() {
"json" => fmt::layer()
.json()
.with_target(true)
.with_file(true)
.with_line_number(true)
.with_thread_ids(true)
.with_span_events(fmt::format::FmtSpan::NEW | fmt::format::FmtSpan::CLOSE)
.boxed(),
_ => fmt::layer()
.pretty()
.with_target(true)
.with_file(true)
.with_line_number(true)
.boxed(),
};
// Optional file output with rotation
if let Ok(log_dir) = std::env::var("GITKS_LOG_DIR") {
let rotation = match env_or("GITKS_LOG_ROTATION", "daily").as_str() {
"hourly" => tracing_appender::rolling::Rotation::HOURLY,
"never" => tracing_appender::rolling::Rotation::NEVER,
_ => tracing_appender::rolling::Rotation::DAILY,
};
let retention = env_u64("GITKS_LOG_RETENTION", 7) as usize;
let mut builder = tracing_appender::rolling::Builder::new()
.rotation(rotation)
.filename_prefix("gitks")
.filename_suffix("log");
if retention > 0 {
builder = builder.max_log_files(retention);
}
let file_appender = builder
.build(&log_dir)
.expect("failed to create log directory");
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
let file_layer = fmt::layer()
.json()
.with_target(true)
.with_file(true)
.with_line_number(true)
.with_writer(non_blocking)
.with_filter(EnvFilter::new("info"))
.boxed();
tracing_subscriber::registry()
.with(env_filter)
.with(fmt_layer)
.with(file_layer)
.init();
Some(guard)
} else {
tracing_subscriber::registry()
.with(env_filter)
.with(fmt_layer)
.init();
None
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
dotenvy::dotenv().ok();
let _log_guard = init_tracing();
tracing::info!(
version = env!("CARGO_PKG_VERSION"),
log_format = %env_or("GITKS_LOG_FORMAT", "pretty"),
"gitks starting up"
);
let host = env_or("GITKS_HOST", DEFAULT_HOST);
let port = env_or("GITKS_PORT", DEFAULT_PORT);
let storage_name = env_or("STORAGE_NAME", DEFAULT_STORAGE_NAME);
let grpc_addr =
std::env::var("GITKS_ADVERTISE_ADDR").unwrap_or_else(|_| format!("http://{host}:{port}"));
let repo_prefix = std::env::var("REPO_PREFIX_PATH")
.map_err(|_| "REPO_PREFIX_PATH environment variable is required (e.g. /data/repos)")?;
let repo_prefix = PathBuf::from(&repo_prefix);
if !repo_prefix.is_absolute() {
return Err("REPO_PREFIX_PATH must be an absolute path".into());
}
if !repo_prefix.exists() {
tracing::info!(path = %repo_prefix.display(), "creating repo prefix directory");
std::fs::create_dir_all(&repo_prefix)?;
}
// Disk cache configuration
let disk_cache_enabled = env_bool("GITKS_DISK_CACHE_ENABLED", false);
let disk_cache_max_age = env_u64("GITKS_DISK_CACHE_MAX_AGE", 300);
let disk_cache = DiskCache::new(
repo_prefix.clone(),
env!("CARGO_PKG_VERSION").to_string(),
disk_cache_max_age,
disk_cache_enabled,
);
if disk_cache_enabled {
tracing::info!(max_age_secs = disk_cache_max_age, "disk cache enabled");
disk_cache.cleanup_on_startup()?;
gitks::disk_cache::start_cache_cleanup_task(disk_cache.clone(), Duration::from_secs(300));
} else {
tracing::info!("disk cache disabled");
}
// Pack cache configuration
let pack_cache_enabled = env_bool("GITKS_PACK_CACHE_ENABLED", false);
let pack_backpressure = env_bool("GITKS_PACK_CACHE_BACKPRESSURE", true);
let pack_cache = if disk_cache_enabled {
tracing::info!(
pack_objects_cache = pack_cache_enabled,
backpressure = pack_backpressure,
"pack cache wrapper enabled"
);
Some(gitks::pack_cache::PackCache::new(
disk_cache.clone(),
pack_backpressure,
))
} else {
None
};
// Hook manager configuration
let hooks_enabled = env_bool("GITKS_HOOKS_ENABLED", true);
let server_hooks_dir = std::env::var("GITKS_SERVER_HOOKS_DIR")
.ok()
.map(PathBuf::from);
let hook_callback_addr = std::env::var("GITKS_HOOK_CALLBACK_ADDR").ok();
let hook_timeout = env_u64("GITKS_HOOK_TIMEOUT", 30);
let allow_custom_hooks = env_bool("GITKS_ALLOW_CUSTOM_HOOKS", true);
let hook_manager = if hooks_enabled {
tracing::info!(
timeout_secs = hook_timeout,
custom_hooks = allow_custom_hooks,
"hooks enabled"
);
Some(HookManager::new(
repo_prefix.clone(),
server_hooks_dir,
hook_callback_addr,
Duration::from_secs(hook_timeout),
allow_custom_hooks,
))
} else {
tracing::info!("hooks disabled");
None
};
// Health check / election configuration
let health_check_interval = env_u64("GITKS_HEALTH_CHECK_INTERVAL", 1);
let max_health_failures = env_u64("GITKS_MAX_HEALTH_FAILURES", 10);
tracing::info!(
interval_secs = health_check_interval,
max_failures = max_health_failures,
"health check configured"
);
let metrics_port = env_u64("GITKS_METRICS_PORT", 9100) as u16;
let http_cancel = tokio_util::sync::CancellationToken::new();
metrics::set_http_cancel_token(http_cancel.clone());
let _metrics_handle = metrics::start_metrics_server(metrics_port);
tracing::info!(port = metrics_port, "metrics server started");
// Slow request threshold
let slow_request_threshold = env_u64("GITKS_SLOW_REQUEST_THRESHOLD_MS", 5000);
metrics::set_slow_request_threshold(slow_request_threshold);
tracing::info!(
threshold_ms = slow_request_threshold,
"slow request detection configured"
);
let etcd_endpoints = std::env::var("GITKS_ETCD_ENDPOINTS")
.ok()
.filter(|s| !s.is_empty())
.map(|s| {
s.split(',')
.map(str::trim)
.map(String::from)
.collect::<Vec<_>>()
});
let cluster_port = env_or("GITKS_CLUSTER_PORT", "4697")
.parse::<u16>()
.unwrap_or(4697);
let cluster_cookie = env_or("GITKS_CLUSTER_COOKIE", "gitks-default-cookie");
let lease_ttl = env_u64("GITKS_LEASE_TTL", 15) as i64;
let connect_timeout_ms = env_u64("GITKS_ETCD_CONNECT_TIMEOUT", 5000);
let cluster_hostname = std::env::var("GITKS_CLUSTER_HOSTNAME")
.or_else(|_| std::env::var("POD_IP"))
.or_else(|_| std::env::var("HOSTNAME"))
.unwrap_or_else(|_| "localhost".to_string());
let _cluster: Option<ClusterManager> = if let Some(endpoints) = etcd_endpoints {
tracing::info!(
?endpoints,
cluster_port,
cluster_hostname = %cluster_hostname,
"starting cluster discovery via etcd"
);
let config = ClusterConfig {
etcd_endpoints: endpoints,
storage_name: storage_name.clone(),
grpc_addr: grpc_addr.clone(),
cluster_port,
cookie: cluster_cookie,
lease_ttl_secs: lease_ttl,
connect_timeout_ms,
cluster_hostname,
};
match ClusterManager::start(config).await {
Ok(cm) => {
tracing::info!("cluster discovery active");
Some(cm)
}
Err(e) => {
tracing::warn!(error = %e, "etcd unavailable, running in standalone mode");
None
}
}
} else {
tracing::info!("GITKS_ETCD_ENDPOINTS not set, running in standalone mode");
None
};
let addr: std::net::SocketAddr = format!("{host}:{port}").parse()?;
let mut svc = GitksService::new(repo_prefix.clone());
if disk_cache_enabled {
svc = svc.with_disk_cache(disk_cache);
}
if let Some(pc) = pack_cache {
svc = svc.with_pack_cache(pc);
}
if let Some(hm) = hook_manager {
svc = svc.with_hook_manager(hm);
}
let raft_data_dir = repo_prefix.join(".gitks_raft");
let (node_actor, node_handle) = init_actor_cluster(
svc.clone(),
storage_name.clone(),
grpc_addr.clone(),
raft_data_dir,
)
.await?;
let svc = svc
.with_actor(node_actor.clone())
.with_grpc_addr(grpc_addr.clone());
tracing::info!(
addr = %addr,
repo_prefix = %repo_prefix.display(),
storage = %storage_name,
advertise = %grpc_addr,
"starting gitks gRPC server"
);
let _route_cache_cleanup = gitks::server::GitksService::start_route_cache_cleanup(svc.clone());
serve(addr, svc).await?;
// Gracefully shut down the HTTP metrics server
http_cancel.cancel();
node_actor.stop(None);
node_handle.await?;
tracing::info!("gitks shut down");
Ok(())
}