src/main.rs
Rama main · réplica de lectura del repositorio original.
mod batch;
mod config;
mod meta;
mod music;
mod state;
use std::net::SocketAddr;
use std::sync::Arc;
use axum::extract::{Path, Query, State};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::routing::{get, get_service, post};
use axum::{Json, Router};
use melodl::SongInfo;
use serde::Deserialize;
use serde_json::{json, Value};
use tower_http::services::{ServeDir, ServeFile};
use crate::batch::BatchStore;
use crate::config::Settings;
use crate::state::{now_ts, spawn_worker, CreateOutcome, SearchStore, TaskStore, TtlCache};
type ApiResult = Result<Json<Value>, ApiError>;
struct ApiError(StatusCode, String);
impl IntoResponse for ApiError {
fn into_response(self) -> axum::response::Response {
(self.0, Json(json!({"detail": self.1}))).into_response()
}
}
fn err(status: StatusCode, msg: impl Into<String>) -> ApiError {
ApiError(status, msg.into())
}
#[derive(Clone)]
struct AppState {
cfg: Arc<Settings>,
cache: Arc<TtlCache<SongInfo>>,
search_store: Arc<SearchStore>,
tasks: Arc<TaskStore>,
batches: Arc<BatchStore>,
meta: Arc<meta::Meta>,
search_tx: tokio::sync::mpsc::UnboundedSender<(String, String)>,
dl_tx: tokio::sync::mpsc::UnboundedSender<(String, SongInfo)>,
batch_tx: tokio::sync::mpsc::UnboundedSender<BatchItem>,
}
type BatchItem = (String, usize, String, String, String); // batch_id, idx, title, artist, task_id
#[tokio::main]
async fn main() {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| "info".into()),
)
.init();
let cfg = Arc::new(Settings::from_env());
let cache = Arc::new(TtlCache::new(cfg.cache_ttl));
let search_store = Arc::new(SearchStore::new(cfg.cache_ttl));
let tasks = Arc::new(TaskStore::new());
let batches = Arc::new(BatchStore::new());
let meta = Arc::new(meta::Meta::new(cfg.meta_ttl));
// 搜索队列 worker：执行搜索并写入结果
let search_tx = spawn_worker({
let (cfg, cache, store) = (cfg.clone(), cache.clone(), search_store.clone());
move |(job_id, keyword): (String, String)| {
let (cfg, cache, store) = (cfg.clone(), cache.clone(), store.clone());
async move {
store.update(&job_id, |j| {
j.status = "running".into();
j.started_at = Some(now_ts());
});
let outcome = music::search(cfg.clone(), keyword).await;
let mut data = Vec::new();
let mut raws = Vec::new();
for (token, view, raw) in outcome.items {
cache.put(token, raw.clone());
data.push(view);
raws.push(raw);
}
store.update(&job_id, |j| {
j.status = "done".into();
j.data = data;
j.raws = raws;
j.sources_ok = outcome.sources_ok;
j.sources_failed = outcome.sources_failed;
j.elapsed = Some(outcome.elapsed);
j.finished_at = Some(now_ts());
});
}
}
});
// 下载队列 worker
let dl_tx = spawn_worker({
let (cfg, tasks) = (cfg.clone(), tasks.clone());
move |(task_id, raw): (String, SongInfo)| {
let (cfg, tasks) = (cfg.clone(), tasks.clone());
async move {
tasks.update(&task_id, |t| t.status = "downloading".into());
match music::download(cfg.clone(), raw).await {
Ok(path) => tasks.update(&task_id, |t| {
t.status = "done".into();
t.file = Some(path.display().to_string());
}),
Err(e) => tasks.update(&task_id, |t| {
t.status = "failed".into();
t.error = Some(e);
}),
}
}
}
});
// 批量下载队列 worker
let batch_tx = spawn_worker({
let (cfg, batches, tasks) = (cfg.clone(), batches.clone(), tasks.clone());
move |(batch_id, idx, title, artist, task_id): BatchItem| {
let (cfg, batches, tasks) = (cfg.clone(), batches.clone(), tasks.clone());
async move {
batch::run_batch_song(cfg, batches, tasks, batch_id, idx, title, artist, task_id)
.await;
}
}
});
let state = AppState {
cfg: cfg.clone(),
cache,
search_store,
tasks,
batches,
meta,
search_tx,
dl_tx,
batch_tx,
};
let app = Router::new()
.route("/", get_service(ServeFile::new(cfg.static_dir.join("index.html"))))
.nest_service("/static", ServeDir::new(&cfg.static_dir))
.route("/api/health", get(health))
.route("/api/search", post(api_search))
.route("/api/search/{job_id}", get(get_search))
.route("/api/download", post(api_download))
.route("/api/tasks", get(list_tasks))
.route("/api/tasks/{task_id}", get(get_task))
.route("/api/meta/charts", get(meta_charts))
.route("/api/meta/playlist/search", get(meta_playlist_search))
.route("/api/meta/playlist/{pid}", get(meta_playlist))
.route("/api/meta/artist", get(meta_artist))
.route("/api/meta/album/search", get(meta_album_search))
.route("/api/meta/album/{aid}", get(meta_album))
.route("/api/meta/download", post(meta_download))
.route("/api/meta/download/{batch_id}", get(meta_download_status))
.with_state(state);
let addr = SocketAddr::from(([0, 0, 0, 0], cfg.port));
let listener = tokio::net::TcpListener::bind(addr)
.await
.expect("bind failed");
tracing::info!("listening on {addr}, music dir: {}", cfg.music_dir.display());
axum::serve(listener, app).await.expect("serve failed");
}
async fn health(State(s): State<AppState>) -> Json<Value> {
Json(json!({
"status": "ok",
"sources": s.cfg.music_sources,
}))
}
#[derive(Deserialize)]
struct SearchRequest {
keyword: String,
}
async fn api_search(
State(s): State<AppState>,
Json(req): Json<SearchRequest>,
) -> ApiResult {
let keyword = req.keyword.trim().to_string();
if keyword.is_empty() {
return Err(err(StatusCode::BAD_REQUEST, "keyword required"));
}
match s.search_store.create_or_get(&keyword, &s.cache) {
CreateOutcome::Cached { job_id, job } => {
let mut v = serde_json::to_value(&job).map_err(|e| {
err(StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
})?;
let obj = v.as_object_mut().unwrap();
obj.insert("search_id".into(), json!(job_id));
obj.insert("status".into(), json!("done"));
obj.insert("cached".into(), json!(true));
Ok(Json(v))
}
CreateOutcome::Running { job_id } => Ok(Json(json!({
"search_id": job_id,
"status": "running",
"cached": false,
}))),
CreateOutcome::Created { job_id } => {
let _ = s.search_tx.send((job_id.clone(), keyword));
Ok(Json(json!({
"search_id": job_id,
"status": "running",
"cached": false,
})))
}
}
}
async fn get_search(
State(s): State<AppState>,
Path(job_id): Path<String>,
) -> Result<Json<Value>, ApiError> {
let job = s
.search_store
.get(&job_id)
.ok_or_else(|| err(StatusCode::NOT_FOUND, "search not found"))?;
Ok(Json(
serde_json::to_value(&job).map_err(|e| err(StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?,
))
}
#[derive(Deserialize)]
struct DownloadRequest {
token: String,
}
async fn api_download(
State(s): State<AppState>,
Json(req): Json<DownloadRequest>,
) -> ApiResult {
let raw = s
.cache
.get(&req.token)
.ok_or_else(|| err(StatusCode::NOT_FOUND, "下载链接已过期，请重新搜索"))?;
let title = if raw.song_name.is_empty() {
"unknown".to_string()
} else {
raw.song_name.clone()
};
let task_id = s.tasks.create(title, None);
let _ = s.dl_tx.send((task_id.clone(), raw));
Ok(Json(json!({"task_id": task_id})))
}
async fn list_tasks(State(s): State<AppState>) -> Json<Value> {
Json(json!({"data": s.tasks.list(20)}))
}
async fn get_task(
State(s): State<AppState>,
Path(task_id): Path<String>,
) -> Result<Json<Value>, ApiError> {
let task = s
.tasks
.get(&task_id)
.ok_or_else(|| err(StatusCode::NOT_FOUND, "task not found"))?;
Ok(Json(
serde_json::to_value(&task)
.map_err(|e| err(StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?,
))
}
// ---------- 元数据浏览（网易云） ----------
#[derive(Deserialize)]
struct PageQuery {
q: Option<String>,
limit: Option<usize>,
offset: Option<usize>,
}
async fn meta_charts(State(s): State<AppState>) -> Json<Value> {
Json(json!({"data": s.meta.charts()}))
}
async fn meta_playlist_search(
State(s): State<AppState>,
Query(p): Query<PageQuery>,
) -> ApiResult {
let q = p
.q
.filter(|q| !q.trim().is_empty())
.ok_or_else(|| err(StatusCode::BAD_REQUEST, "q required"))?;
let out = s
.meta
.playlist_search(&q, p.limit.unwrap_or(20), p.offset.unwrap_or(0))
.await
.map_err(|e| err(StatusCode::INTERNAL_SERVER_ERROR, e))?;
Ok(Json(out))
}
async fn meta_playlist(
State(s): State<AppState>,
Path(pid): Path<u64>,
Query(p): Query<PageQuery>,
) -> ApiResult {
let out = s
.meta
.playlist_detail(pid, p.limit.unwrap_or(50), p.offset.unwrap_or(0))
.await
.map_err(|e| err(StatusCode::INTERNAL_SERVER_ERROR, e))?;
Ok(Json(out))
}
async fn meta_artist(State(s): State<AppState>, Query(p): Query<PageQuery>) -> ApiResult {
let q = p
.q
.filter(|q| !q.trim().is_empty())
.ok_or_else(|| err(StatusCode::BAD_REQUEST, "q required"))?;
let out = s
.meta
.artist_hot(&q, p.limit.unwrap_or(50))
.await
.map_err(|e| err(StatusCode::INTERNAL_SERVER_ERROR, e))?;
Ok(Json(json!({"data": out})))
}
async fn meta_album_search(
State(s): State<AppState>,
Query(p): Query<PageQuery>,
) -> ApiResult {
let q = p
.q
.filter(|q| !q.trim().is_empty())
.ok_or_else(|| err(StatusCode::BAD_REQUEST, "q required"))?;
let out = s
.meta
.album_search(&q, p.limit.unwrap_or(20))
.await
.map_err(|e| err(StatusCode::INTERNAL_SERVER_ERROR, e))?;
Ok(Json(out))
}
async fn meta_album(State(s): State<AppState>, Path(aid): Path<u64>) -> ApiResult {
let out = s
.meta
.album_detail(aid)
.await
.map_err(|e| err(StatusCode::INTERNAL_SERVER_ERROR, e))?;
Ok(Json(json!({"data": out})))
}
#[derive(Deserialize)]
struct MetaSong {
#[serde(default)]
title: String,
#[serde(default)]
artist: String,
}
#[derive(Deserialize)]
struct MetaDownloadRequest {
songs: Vec<MetaSong>,
}
async fn meta_download(
State(s): State<AppState>,
Json(req): Json<MetaDownloadRequest>,
) -> ApiResult {
if req.songs.is_empty() {
return Err(err(StatusCode::BAD_REQUEST, "songs required"));
}
let songs: Vec<(String, String)> = req
.songs
.iter()
.map(|s| (s.title.clone(), s.artist.clone()))
.collect();
let batch_id = s.batches.create(songs);
let count = req.songs.len();
for (idx, song) in req.songs.iter().enumerate() {
let title = format!("{} - {}", song.title, song.artist);
let title = title.trim_matches(&[' ', '-'][..]).to_string();
let task_id = s.tasks.create(title, Some(batch_id.clone()));
let _ = s.batch_tx.send((
batch_id.clone(),
idx,
song.title.clone(),
song.artist.clone(),
task_id,
));
}
Ok(Json(json!({"batch_id": batch_id, "count": count})))
}
async fn meta_download_status(
State(s): State<AppState>,
Path(batch_id): Path<String>,
) -> ApiResult {
let batch = s
.batches
.get(&batch_id)
.ok_or_else(|| err(StatusCode::NOT_FOUND, "batch not found"))?;
Ok(Json(batch))
}