codetics-ocpp-service 0.1.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Kevin Vavelin, Codetics (https://codetics.fr)

//! Registry of the station connections held by this node.
//!
//! Sharded so that tens of thousands of connections and concurrent API calls
//! do not contend on one lock. Each entry holds the sender side of the
//! connection's command queue; the session task owns everything else.

use std::collections::HashMap;
use std::hash::BuildHasher;
use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use ocpp_core::{DateTime, OcppVersion, RpcError, Violation};
use serde::Serialize;
use tokio::sync::{mpsc, oneshot};

const SHARDS: usize = 64;

#[derive(Debug)]
pub enum Command {
    /// Send a CALL and report the outcome.
    Call {
        action: String,
        payload: serde_json::Value,
        reply: oneshot::Sender<CallOutcome>,
    },
    /// Close the connection (e.g. replaced by a newer one).
    Close { code: u16, reason: &'static str },
}

#[derive(Debug)]
pub enum CallFailure {
    /// The station answered with CALLERROR.
    Error(RpcError),
    /// The CALLRESULT payload does not match the message definition (raw payload attached).
    InvalidResponse {
        violation: Violation,
        raw: serde_json::Value,
    },
    Timeout,
    Disconnected,
}

pub type CallOutcome = Result<serde_json::Value, CallFailure>;

#[derive(Debug, Clone)]
pub struct ConnHandle {
    pub conn_id: u64,
    pub station_id: i64,
    pub version: OcppVersion,
    pub connected_at: DateTime,
    pub remote_addr: Option<String>,
    pub tx: mpsc::Sender<Command>,
}

#[derive(Debug, Serialize)]
pub struct ConnInfo {
    pub ocpp_identity: String,
    pub station_id: i64,
    pub ocpp_version: String,
    pub connected_at: DateTime,
    pub remote_addr: Option<String>,
}

pub enum SendError {
    NotConnected,
    QueueFull,
}

pub struct Hub {
    shards: Vec<RwLock<HashMap<String, ConnHandle>>>,
    hasher: std::collections::hash_map::RandomState,
    next_id: AtomicU64,
}

impl Default for Hub {
    fn default() -> Self {
        Self::new()
    }
}

impl Hub {
    pub fn new() -> Self {
        Hub {
            shards: (0..SHARDS).map(|_| RwLock::new(HashMap::new())).collect(),
            hasher: Default::default(),
            next_id: AtomicU64::new(1),
        }
    }

    fn shard(&self, identity: &str) -> &RwLock<HashMap<String, ConnHandle>> {
        &self.shards[self.hasher.hash_one(identity) as usize % SHARDS]
    }

    pub fn next_conn_id(&self) -> u64 {
        self.next_id.fetch_add(1, Ordering::Relaxed)
    }

    /// Registers a connection, returning the handle it replaced (if any).
    pub fn register(&self, identity: &str, handle: ConnHandle) -> Option<ConnHandle> {
        self.shard(identity)
            .write()
            .unwrap_or_else(|e| e.into_inner())
            .insert(identity.to_owned(), handle)
    }

    /// Removes the entry only if it still belongs to `conn_id`.
    pub fn unregister(&self, identity: &str, conn_id: u64) {
        let mut shard = self
            .shard(identity)
            .write()
            .unwrap_or_else(|e| e.into_inner());
        if shard.get(identity).is_some_and(|h| h.conn_id == conn_id) {
            shard.remove(identity);
        }
    }

    pub fn get(&self, identity: &str) -> Option<ConnHandle> {
        self.shard(identity)
            .read()
            .unwrap_or_else(|e| e.into_inner())
            .get(identity)
            .cloned()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn len(&self) -> usize {
        self.shards
            .iter()
            .map(|s| s.read().unwrap_or_else(|e| e.into_inner()).len())
            .sum()
    }

    pub fn list(&self) -> Vec<ConnInfo> {
        let mut out = Vec::new();
        for s in &self.shards {
            for (id, h) in s.read().unwrap_or_else(|e| e.into_inner()).iter() {
                out.push(ConnInfo {
                    ocpp_identity: id.clone(),
                    station_id: h.station_id,
                    ocpp_version: h.version.to_string(),
                    connected_at: h.connected_at,
                    remote_addr: h.remote_addr.clone(),
                });
            }
        }
        out.sort_by(|a, b| a.ocpp_identity.cmp(&b.ocpp_identity));
        out
    }

    /// Queues a CALL on the station's connection and waits for its outcome.
    /// The session sends CALLs one at a time (OCPP-J §4.1.1), so `wait`
    /// covers queueing plus the call timeout.
    pub async fn call(
        &self,
        identity: &str,
        action: &str,
        payload: serde_json::Value,
        wait: Duration,
    ) -> Result<CallOutcome, SendError> {
        let handle = self.get(identity).ok_or(SendError::NotConnected)?;
        let (reply, rx) = oneshot::channel();
        handle
            .tx
            .try_send(Command::Call {
                action: action.to_owned(),
                payload,
                reply,
            })
            .map_err(|e| match e {
                mpsc::error::TrySendError::Full(_) => SendError::QueueFull,
                mpsc::error::TrySendError::Closed(_) => SendError::NotConnected,
            })?;
        Ok(match tokio::time::timeout(wait, rx).await {
            Ok(Ok(outcome)) => outcome,
            Ok(Err(_)) => Err(CallFailure::Disconnected),
            Err(_) => Err(CallFailure::Timeout),
        })
    }
}