gitks/actor/handler.rs

use std::collections::HashSet;
use async_trait::async_trait;
use ractor::pg;
use ractor::{Actor, ActorProcessingErr, ActorRef, SupervisionEvent};
use crate::actor::message::{GitNodeMessage, NodeHealth, RouteDecision};
use crate::pb::RepositoryHeader;
use crate::server::GitksService;

#[derive(Clone)]
pub struct GitNodeActor {
    pub version: String,
    pub service: GitksService,
}

impl GitNodeActor {
    pub fn init(service: GitksService) -> Self {
        GitNodeActor {
            version: env!("CARGO_PKG_VERSION").to_string(),
            service,
        }
    }
}

pub struct GitNodeArgs {
    pub storage_name: String,
    pub grpc_addr: String,
}

pub struct GitNodeState {
    storage_name: String,
    actor_name: String,
    grpc_addr: String,
    registered_repos: HashSet<String>,
}

#[async_trait]
impl Actor for GitNodeActor {
    type Msg = GitNodeMessage;
    type State = GitNodeState;
    type Arguments = GitNodeArgs;

    async fn pre_start(
        &self,
        myself: ActorRef<Self::Msg>,
        args: Self::Arguments,
    ) -> Result<Self::State, ActorProcessingErr> {
        let actor_name = format!("git_node_{}", args.storage_name);
        pg::join("gitks_nodes".to_string(), vec![myself.get_cell()]);
        pg::join_scoped(args.storage_name.clone(), "node".to_string(), vec![myself.get_cell()]);
        tracing::info!(storage_name = %args.storage_name, actor_name = %actor_name, grpc_addr = %args.grpc_addr, "GitNodeActor started");
        Ok(GitNodeState {
            storage_name: args.storage_name,
            actor_name,
            grpc_addr: args.grpc_addr,
            registered_repos: HashSet::new(),
        })
    }

    async fn handle(
        &self,
        myself: ActorRef<Self::Msg>,
        message: Self::Msg,
        state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        match message {
            GitNodeMessage::ScanAndRegister => {
                let repos = self.service.scan_all_repo()?;
                tracing::info!(storage_name = %state.storage_name, found = repos.len(), "scanning local repositories");
                for repo_path in repos {
                    let relative_path = repo_path
                        .strip_prefix(self.service.repo_prefix.to_string_lossy().as_ref())
                        .unwrap_or(&repo_path)
                        .trim_start_matches('/')
                        .to_string();
                    register_repo(&myself, state, relative_path);
                }
            }

            GitNodeMessage::RegisterRepository(header) => {
                register_repo(&myself, state, header.relative_path);
            }

            GitNodeMessage::RemoveRepository(header) => {
                state.registered_repos.remove(&header.relative_path);
                tracing::info!(
                    storage_name = %state.storage_name,
                    relative_path = %header.relative_path,
                    "repository route removed"
                );
            }

            GitNodeMessage::RouteRepository(header, reply) => {
                let found = state.registered_repos.contains(&header.relative_path);
                reply.send(RouteDecision {
                    found,
                    storage_name: state.storage_name.clone(),
                    relative_path: header.relative_path,
                    actor_name: if found { state.actor_name.clone() } else { String::new() },
                    grpc_addr: if found { state.grpc_addr.clone() } else { String::new() },
                }).ok();
            }

            GitNodeMessage::ListRepositoryPaths(reply) => {
                let paths: Vec<String> = state.registered_repos.iter().cloned().collect();
                reply.send(paths.join("\n")).ok();
            }

            GitNodeMessage::RepositoryExists(header, reply) => {
                reply.send(state.registered_repos.contains(&header.relative_path)).ok();
            }

            GitNodeMessage::GetNodeHealth(reply) => {
                reply.send(NodeHealth {
                    storage_name: state.storage_name.clone(),
                    repo_count: state.registered_repos.len() as u64,
                    healthy: true,
                    version: self.version.clone(),
                }).ok();
            }
        }
        Ok(())
    }

    async fn handle_supervisor_evt(
        &self,
        _myself: ActorRef<Self::Msg>,
        evt: SupervisionEvent,
        _state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        match evt {
            SupervisionEvent::ActorStarted(who) => tracing::debug!(actor = ?who.get_id(), "child started"),
            SupervisionEvent::ActorTerminated(who, _, reason) => {
                tracing::warn!(actor = ?who.get_id(), reason = ?reason, "child terminated")
            }
            SupervisionEvent::ActorFailed(who, panic_msg) => {
                tracing::error!(actor = ?who.get_id(), msg = %panic_msg, "child panicked")
            }
            SupervisionEvent::ProcessGroupChanged(group) => tracing::info!(group = ?group, "PG membership changed"),
            _ => {}
        }
        Ok(())
    }

    async fn post_stop(
        &self,
        _myself: ActorRef<Self::Msg>,
        state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        tracing::info!(storage_name = %state.storage_name, "GitNodeActor stopped");
        Ok(())
    }
}

fn register_repo(myself: &ActorRef<GitNodeMessage>, state: &mut GitNodeState, relative_path: String) {
    let category = extract_category(&relative_path);
    pg::join_scoped(state.storage_name.clone(), category.to_string(), vec![myself.get_cell()]);
    state.registered_repos.insert(relative_path.clone());
    tracing::info!(
        storage_name = %state.storage_name,
        category = %category,
        relative_path = %relative_path,
        actor_name = %state.actor_name,
        "repository route registered"
    );
}

fn extract_category(relative_path: &str) -> &str {
    relative_path.split('/').next().unwrap_or("root")
}

pub async fn start_node_actor(
    service: GitksService,
    storage_name: String,
    grpc_addr: String,
) -> Result<(ActorRef<GitNodeMessage>, tokio::task::JoinHandle<()>), ractor::SpawnErr> {
    let actor = GitNodeActor::init(service);
    let (actor_ref, handle) = Actor::spawn(
        Some(format!("git_node_{storage_name}")),
        actor,
        GitNodeArgs { storage_name, grpc_addr },
    ).await?;
    actor_ref.cast(GitNodeMessage::ScanAndRegister).ok();
    Ok((actor_ref, handle))
}

pub fn get_cluster_nodes(storage_name: &str) -> Vec<ractor::ActorCell> {
    pg::get_scoped_members(&storage_name.to_string(), &"node".to_string())
}

pub fn get_category_members(storage_name: &str, category: &str) -> Vec<ractor::ActorCell> {
    pg::get_scoped_members(&storage_name.to_string(), &category.to_string())
}

pub fn route_group_for(header: &RepositoryHeader) -> String {
    extract_category(&header.relative_path).to_string()
}

pub fn list_all_groups() -> Vec<String> {
    pg::which_groups()
}