feat: Add transcription worker and state management
Introduce an `AppState` struct to hold shared application state, including configuration and the transcription sender. Modify `create_router` and `create_router_with_state` to use `AppState`. Update `main.rs` to initialize the transcription worker and pass the `AppState` to the router. Implement the transcription worker logic in `server/src/transcribe/worker.rs`, handling job queues, FFmpeg remuxing, and Whisper transcription. Update `server/src/auth.rs` to use `Arc<Config>` from the shared state. Modify `server/src/routes/upload.rs` to enqueue transcription jobs and handle potential errors. Add `TranscribeJob` and channel types in `server/src/transcribe/mod.rs`.
This commit is contained in:
+12
-9
@@ -1,7 +1,7 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::FromRequestParts;
|
||||
use axum::extract::{FromRef, FromRequestParts};
|
||||
use axum::http::request::Parts;
|
||||
|
||||
use crate::config::Config;
|
||||
@@ -14,31 +14,34 @@ pub struct AuthenticatedUser {
|
||||
pub data_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl FromRequestParts<Arc<Config>> for AuthenticatedUser {
|
||||
impl<S> FromRequestParts<S> for AuthenticatedUser
|
||||
where
|
||||
S: Send + Sync,
|
||||
Arc<Config>: FromRef<S>,
|
||||
{
|
||||
type Rejection = AppError;
|
||||
|
||||
async fn from_request_parts(
|
||||
parts: &mut Parts,
|
||||
state: &Arc<Config>,
|
||||
) -> Result<Self, Self::Rejection> {
|
||||
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
|
||||
let config = Arc::<Config>::from_ref(state);
|
||||
|
||||
let api_key = parts
|
||||
.headers
|
||||
.get("X-API-Key")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.ok_or(AppError::Unauthorized)?;
|
||||
|
||||
let slug = state
|
||||
let slug = config
|
||||
.api_keys
|
||||
.get(api_key)
|
||||
.ok_or(AppError::Unauthorized)?;
|
||||
|
||||
let user = state
|
||||
let user = config
|
||||
.users
|
||||
.iter()
|
||||
.find(|u| u.slug == *slug)
|
||||
.ok_or(AppError::Unauthorized)?;
|
||||
|
||||
let data_dir = state.data_path.join(slug);
|
||||
let data_dir = config.data_path.join(slug);
|
||||
tokio::fs::create_dir_all(data_dir.join("open")).await?;
|
||||
tokio::fs::create_dir_all(data_dir.join("done")).await?;
|
||||
|
||||
|
||||
+35
-3
@@ -7,11 +7,43 @@ pub mod transcribe;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::FromRef;
|
||||
use axum::Router;
|
||||
|
||||
use config::Config;
|
||||
use transcribe::TranscribeSender;
|
||||
|
||||
/// Build the application router. Used by main.rs and integration tests.
|
||||
pub fn create_router(config: Arc<Config>) -> Router {
|
||||
routes::api_router().with_state(config)
|
||||
/// Shared application state. Clonable — all fields are cheap to clone
|
||||
/// (`Arc` and `mpsc::Sender`).
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub config: Arc<Config>,
|
||||
pub transcribe_tx: TranscribeSender,
|
||||
}
|
||||
|
||||
impl FromRef<AppState> for Arc<Config> {
|
||||
fn from_ref(state: &AppState) -> Self {
|
||||
state.config.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl FromRef<AppState> for TranscribeSender {
|
||||
fn from_ref(state: &AppState) -> Self {
|
||||
state.transcribe_tx.clone()
|
||||
}
|
||||
}
|
||||
|
||||
/// Test/simple entrypoint: jobs pushed into the transcribe channel are dropped
|
||||
/// because the receiver is not retained. Use [`create_router_with_state`] from
|
||||
/// `main.rs` where a real worker owns the receiver.
|
||||
pub fn create_router(config: Arc<Config>) -> Router {
|
||||
let (tx, _rx) = transcribe::channel();
|
||||
create_router_with_state(AppState {
|
||||
config,
|
||||
transcribe_tx: tx,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn create_router_with_state(state: AppState) -> Router {
|
||||
routes::api_router().with_state(state)
|
||||
}
|
||||
|
||||
+19
-1
@@ -9,6 +9,8 @@ use tracing_subscriber::util::SubscriberInitExt;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
use doctate_server::config::Config;
|
||||
use doctate_server::transcribe;
|
||||
use doctate_server::AppState;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
@@ -37,10 +39,26 @@ async fn main() {
|
||||
)
|
||||
.init();
|
||||
|
||||
// Transcription pipeline: shared HTTP client + worker task.
|
||||
let (transcribe_tx, transcribe_rx) = transcribe::channel();
|
||||
let http_client = reqwest::Client::builder()
|
||||
.build()
|
||||
.expect("reqwest client build failed");
|
||||
tokio::spawn(transcribe::worker::run(
|
||||
transcribe_rx,
|
||||
config.clone(),
|
||||
http_client,
|
||||
));
|
||||
|
||||
let state = AppState {
|
||||
config: config.clone(),
|
||||
transcribe_tx,
|
||||
};
|
||||
|
||||
let addr = format!("0.0.0.0:{}", config.server_port);
|
||||
info!(port = config.server_port, "Server starting");
|
||||
|
||||
let app = doctate_server::create_router(config)
|
||||
let app = doctate_server::create_router_with_state(state)
|
||||
.layer(TraceLayer::new_for_http())
|
||||
.layer(RequestBodyLimitLayer::new(50 * 1024 * 1024)); // 50 MB
|
||||
|
||||
|
||||
@@ -3,14 +3,12 @@ mod health;
|
||||
mod upload;
|
||||
mod web;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::routing::{get, post};
|
||||
use axum::Router;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::AppState;
|
||||
|
||||
pub fn api_router() -> Router<Arc<Config>> {
|
||||
pub fn api_router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/health", get(health::handle_health))
|
||||
.route("/api/debug/whoami", get(debug::handle_whoami))
|
||||
|
||||
@@ -1,15 +1,18 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use axum::extract::Multipart;
|
||||
use axum::extract::{Multipart, State};
|
||||
use axum::Json;
|
||||
use tracing::info;
|
||||
use tokio::sync::mpsc::error::TrySendError;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::auth::AuthenticatedUser;
|
||||
use crate::error::AppError;
|
||||
use crate::models::{AckResponse, AckStatus};
|
||||
use crate::transcribe::{TranscribeJob, TranscribeSender};
|
||||
|
||||
pub async fn handle_upload(
|
||||
user: AuthenticatedUser,
|
||||
State(transcribe_tx): State<TranscribeSender>,
|
||||
mut multipart: Multipart,
|
||||
) -> Result<Json<AckResponse>, AppError> {
|
||||
let mut case_id: Option<String> = None;
|
||||
@@ -81,6 +84,18 @@ pub async fn handle_upload(
|
||||
"Recording received"
|
||||
);
|
||||
|
||||
// Enqueue transcription. Failure here is non-fatal: on restart the recovery
|
||||
// scan will re-enqueue any .m4a without a .transcript.txt sibling.
|
||||
match transcribe_tx.try_send(TranscribeJob { audio_path: file_path.clone() }) {
|
||||
Ok(()) => {}
|
||||
Err(TrySendError::Full(_)) => {
|
||||
warn!(file = %file_path.display(), "Transcription queue full; will be picked up on restart");
|
||||
}
|
||||
Err(TrySendError::Closed(_)) => {
|
||||
warn!(file = %file_path.display(), "Transcription queue closed (no worker)");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Json(AckResponse {
|
||||
case_id,
|
||||
recorded_at,
|
||||
|
||||
@@ -1,2 +1,21 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
pub mod ffmpeg;
|
||||
pub mod whisper;
|
||||
pub mod worker;
|
||||
|
||||
pub const QUEUE_DEPTH: usize = 64;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TranscribeJob {
|
||||
pub audio_path: PathBuf,
|
||||
}
|
||||
|
||||
pub type TranscribeSender = mpsc::Sender<TranscribeJob>;
|
||||
pub type TranscribeReceiver = mpsc::Receiver<TranscribeJob>;
|
||||
|
||||
pub fn channel() -> (TranscribeSender, TranscribeReceiver) {
|
||||
mpsc::channel(QUEUE_DEPTH)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use super::{ffmpeg, whisper, TranscribeReceiver};
|
||||
use crate::config::Config;
|
||||
|
||||
/// Consume transcription jobs sequentially. Whisper holds the GPU, so parallelism
|
||||
/// here would only cause contention. One job in flight at a time.
|
||||
pub async fn run(mut rx: TranscribeReceiver, config: Arc<Config>, client: reqwest::Client) {
|
||||
info!("Transcription worker started");
|
||||
let timeout = Duration::from_secs(config.whisper_timeout_seconds);
|
||||
|
||||
while let Some(job) = rx.recv().await {
|
||||
let audio_path = job.audio_path;
|
||||
let transcript_path = audio_path.with_extension("transcript.txt");
|
||||
|
||||
if transcript_path.exists() {
|
||||
continue;
|
||||
}
|
||||
|
||||
info!(audio = %audio_path.display(), "Transcribing");
|
||||
|
||||
let remuxed = match ffmpeg::remux_faststart(&audio_path).await {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
error!(audio = %audio_path.display(), error = %e, "ffmpeg remux failed");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let text = match whisper::transcribe(&client, &config.whisper_url, remuxed.path(), timeout).await {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
error!(audio = %audio_path.display(), error = %e, "whisper call failed");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = tokio::fs::write(&transcript_path, &text).await {
|
||||
error!(transcript = %transcript_path.display(), error = %e, "writing transcript failed");
|
||||
continue;
|
||||
}
|
||||
|
||||
info!(
|
||||
audio = %audio_path.display(),
|
||||
bytes = text.len(),
|
||||
"Transcript written"
|
||||
);
|
||||
}
|
||||
|
||||
warn!("Transcription worker stopped (channel closed)");
|
||||
}
|
||||
Reference in New Issue
Block a user