ignite_client/

cache.rs

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use crate::protocol::codec::read_bool;
use crate::protocol::error::{ProtocolError, value_type_name};
use crate::protocol::messages::{
    decode_cache_get_all_response, decode_cache_get_size_response, decode_cache_value_response,
    encode_cache_create_with_name, encode_cache_destroy_req, encode_cache_get_size,
    encode_cache_key_req, encode_cache_kv_req, encode_cache_multi_key_req,
    encode_cache_multi_kv_req,
};
use crate::protocol::op_code;
use crate::protocol::{ExpiryPolicy, IgniteValue};
use crate::transport::{IgniteConnection, next_request_id};
use bytes::Bytes;
use uuid::Uuid;

use crate::affinity::AffinityContext;
use crate::binary::{ReadBinary, WriteBinary};
use crate::channel::ChannelRegistry;
use crate::error::{IgniteError, Result};
use crate::protocol::binary::client_ops::{
    BinaryTypeCache, RegistryResolver, fetch_binary_type, metadata_compatible,
    prefetch_nested_schemas, register_binary_type, schema_field_ids_with_refetch,
};
use crate::protocol::binary::header::{BinaryHeader, flags};
use crate::protocol::binary::reader::BinaryObjectReader;

// ─── CacheSource ─────────────────────────────────────────────────────────────

/// Backing connection source for an [`IgniteCache`] handle.
///
/// - `Routed` — non-transactional: each operation routes through the channel
///   registry, preferring the key's owning node (partition awareness) and
///   falling back to the default channel.
/// - `Tx` — transactional: every operation goes through the transaction's
///   dedicated connection and embeds the `tx_id` in the cache-header flags byte.
#[derive(Clone)]
pub(crate) enum CacheSource {
    Routed {
        registry: Arc<ChannelRegistry>,
        affinity: Arc<AffinityContext>,
    },
    Tx {
        tx_id: i32,
        conn: Arc<IgniteConnection>,
    },
}

// ─── IgniteCache ─────────────────────────────────────────────────────────────

/// A handle to an Ignite cache, obtained via [`crate::IgniteClient::cache`],
/// [`crate::IgniteClient::get_or_create_cache`], or [`crate::transaction::Transaction::cache`].
///
/// Cheap to clone — internally holds an `i32` cache-id and either a pool
/// reference (Arc-backed) or a transaction connection (Arc-backed).
#[derive(Clone)]
pub struct IgniteCache {
    pub(crate) cache_id: i32,
    source: CacheSource,
    /// Optional expiry policy applied to every operation on this handle.
    expiry: Option<ExpiryPolicy>,
    /// Client-side cache of binary-type metadata, shared with the
    /// [`crate::IgniteClient`] this handle was obtained from (all clones and
    /// all cache handles refer to the same logical client). Used by
    /// [`Self::get_binary`]/[`Self::put_binary`] to decode/register
    /// compact-footer binary objects without a metadata round-trip on every
    /// call. Transactional handles ([`CacheSource::Tx`]) get a fresh,
    /// unshared map since they have no registry to fetch/register against.
    binary_types: Arc<BinaryTypeCache>,
}

impl std::fmt::Debug for IgniteCache {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut s = f.debug_struct("IgniteCache");
        s.field("cache_id", &self.cache_id);
        match &self.source {
            CacheSource::Routed { .. } => {
                s.field("source", &"routed");
            }
            CacheSource::Tx { tx_id, .. } => {
                s.field("tx_id", tx_id);
            }
        }
        s.field("expiry", &self.expiry).finish()
    }
}

impl IgniteCache {
    /// Create a non-transactional cache handle backed by the channel registry,
    /// routing keys to their owning node via the shared affinity context.
    pub(crate) fn new(
        cache_id: i32,
        registry: Arc<ChannelRegistry>,
        affinity: Arc<AffinityContext>,
        binary_types: Arc<BinaryTypeCache>,
    ) -> Self {
        Self {
            cache_id,
            source: CacheSource::Routed { registry, affinity },
            expiry: None,
            binary_types,
        }
    }

    /// Create a transactional cache handle that routes all ops through a
    /// transaction's dedicated connection and embeds `tx_id` in every request.
    ///
    /// `binary_types` is a fresh, unshared cache: `Transaction::cache` has no
    /// access to the owning client's registry, so `get_binary`/`put_binary`
    /// error out on this handle anyway (see [`Self::registry`]) and never
    /// populate it.
    pub(crate) fn new_tx(cache_id: i32, tx_id: i32, conn: Arc<IgniteConnection>) -> Self {
        Self {
            cache_id,
            source: CacheSource::Tx { tx_id, conn },
            expiry: None,
            binary_types: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Return a new handle to the same cache whose operations apply `policy`
    /// (per-entry time-to-live on create / update / access).  Cheap to clone.
    ///
    /// ```no_run
    /// # use ignite_client::{IgniteCache, ExpiryPolicy, ExpiryDuration, IgniteValue};
    /// # async fn f(cache: IgniteCache) -> ignite_client::Result<()> {
    /// // New entries live 60s; updates and reads leave the TTL unchanged.
    /// let ttl = ExpiryPolicy::new(
    ///     ExpiryDuration::Millis(60_000),
    ///     ExpiryDuration::Unchanged,
    ///     ExpiryDuration::Unchanged,
    /// );
    /// cache.with_expiry_policy(ttl).put(IgniteValue::Int(1), IgniteValue::Int(2)).await?;
    /// # Ok(()) }
    /// ```
    pub fn with_expiry_policy(&self, policy: ExpiryPolicy) -> IgniteCache {
        IgniteCache {
            cache_id: self.cache_id,
            source: self.source.clone(),
            expiry: Some(policy),
            binary_types: self.binary_types.clone(),
        }
    }

    // ── Internal helpers ──────────────────────────────────────────────────────

    /// Returns the active transaction ID, or `None` for non-transactional handles.
    fn tx_id(&self) -> Option<i32> {
        match &self.source {
            CacheSource::Routed { .. } => None,
            CacheSource::Tx { tx_id, .. } => Some(*tx_id),
        }
    }

    /// The expiry policy applied to this handle's operations, if any.
    fn expiry(&self) -> Option<&ExpiryPolicy> {
        self.expiry.as_ref()
    }

    /// The channel registry backing this handle, or an error for
    /// transactional handles (which route through a dedicated connection and
    /// carry no registry). Binary-object metadata lookups (`get_binary`/
    /// `put_binary`) need a registry to reach `OP_BINARY_TYPE_GET`/`_PUT`.
    fn registry(&self) -> Result<&ChannelRegistry> {
        match &self.source {
            CacheSource::Routed { registry, .. } => Ok(registry),
            CacheSource::Tx { .. } => Err(IgniteError::UnsupportedOnTransaction(
                "binary-object metadata operations (get_binary/put_binary)",
            )),
        }
    }

    /// Resolve the node owning `key`, lazily refreshing the affinity mapping
    /// first.  `primary = false` (read-only ops) may resolve to a same-data
    /// centre backup owner.  Returns `None` (route to the default channel) for
    /// transactional handles, when PA is disabled, or for unsupported keys.
    async fn route(&self, key: &IgniteValue, primary: bool) -> Option<Uuid> {
        match &self.source {
            CacheSource::Tx { .. } => None,
            CacheSource::Routed { registry, affinity } => {
                registry.ensure_affinity(affinity, self.cache_id).await;
                affinity.affinity_node(self.cache_id, key, primary)
            }
        }
    }

    /// Send `payload` to `target` (its owning node when known, else the default
    /// channel), returning the post-header response bytes.
    /// `IgniteConnection::request` already strips and validates the response
    /// header, so the returned `Bytes` contains only the operation payload.
    async fn send(&self, req_id: i64, payload: Bytes, target: Option<Uuid>) -> Result<Bytes> {
        match &self.source {
            CacheSource::Routed { registry, .. } => {
                let conn = registry.get(target).await?;
                let out = conn
                    .request(req_id, payload)
                    .await
                    .map_err(IgniteError::Transport)?;
                registry.observe_topology(&conn);
                Ok(out)
            }
            CacheSource::Tx { conn, .. } => conn
                .request(req_id, payload)
                .await
                .map_err(IgniteError::Transport),
        }
    }

    // ── Public API ────────────────────────────────────────────────────────────

    /// Retrieve a value by key.  Returns `IgniteValue::Null` if the key is not present.
    pub async fn get(&self, key: IgniteValue) -> Result<IgniteValue> {
        let req_id = next_request_id();
        // Read-only op: a same-DC backup owner is acceptable.
        let target = self.route(&key, false).await;
        let payload = encode_cache_key_req(
            op_code::CACHE_GET,
            req_id,
            self.cache_id,
            &key,
            self.tx_id(),
            self.expiry(),
        );
        let mut resp = self.send(req_id, payload, target).await?;
        decode_cache_value_response(&mut resp).map_err(IgniteError::Protocol)
    }

    /// Store a key-value pair.  Overwrites any existing value.
    #[must_use = "futures do nothing unless you `.await` them"]
    pub async fn put(&self, key: IgniteValue, value: IgniteValue) -> Result<()> {
        let req_id = next_request_id();
        let target = self.route(&key, true).await;
        let payload = encode_cache_kv_req(
            op_code::CACHE_PUT,
            req_id,
            self.cache_id,
            &key,
            &value,
            self.tx_id(),
            self.expiry(),
        );
        self.send(req_id, payload, target).await?;
        Ok(())
    }

    /// Store a key-value pair **only if the key is not already present**.
    /// Returns `true` if the value was stored, `false` if the key already existed.
    pub async fn put_if_absent(&self, key: IgniteValue, value: IgniteValue) -> Result<bool> {
        let req_id = next_request_id();
        let target = self.route(&key, true).await;
        let payload = encode_cache_kv_req(
            op_code::CACHE_PUT_IF_ABSENT,
            req_id,
            self.cache_id,
            &key,
            &value,
            self.tx_id(),
            self.expiry(),
        );
        let mut resp = self.send(req_id, payload, target).await?;
        read_bool(&mut resp).map_err(IgniteError::Protocol)
    }

    /// Retrieve values for multiple keys.
    /// Returns only the pairs for keys that exist in the cache; absent keys are omitted.
    pub async fn get_all(&self, keys: Vec<IgniteValue>) -> Result<Vec<(IgniteValue, IgniteValue)>> {
        let req_id = next_request_id();
        let payload = encode_cache_multi_key_req(
            op_code::CACHE_GET_ALL,
            req_id,
            self.cache_id,
            &keys,
            self.tx_id(),
            self.expiry(),
        );
        // Multi-key requests span partitions; route via the default channel.
        let mut resp = self.send(req_id, payload, None).await?;
        decode_cache_get_all_response(&mut resp).map_err(IgniteError::Protocol)
    }

    /// Store multiple key-value pairs.
    #[must_use = "futures do nothing unless you `.await` them"]
    pub async fn put_all(&self, entries: Vec<(IgniteValue, IgniteValue)>) -> Result<()> {
        let req_id = next_request_id();
        let payload = encode_cache_multi_kv_req(
            op_code::CACHE_PUT_ALL,
            req_id,
            self.cache_id,
            &entries,
            self.tx_id(),
            self.expiry(),
        );
        self.send(req_id, payload, None).await?;
        Ok(())
    }

    /// Returns `true` if the cache contains the given key.
    pub async fn contains_key(&self, key: IgniteValue) -> Result<bool> {
        let req_id = next_request_id();
        // Read-only op: a same-DC backup owner is acceptable.
        let target = self.route(&key, false).await;
        let payload = encode_cache_key_req(
            op_code::CACHE_CONTAINS_KEY,
            req_id,
            self.cache_id,
            &key,
            self.tx_id(),
            self.expiry(),
        );
        let mut resp = self.send(req_id, payload, target).await?;
        read_bool(&mut resp).map_err(IgniteError::Protocol)
    }

    /// Remove a key.
    ///
    /// Implemented via `CACHE_GET_AND_REMOVE` (op 1007) because the standalone
    /// `CACHE_REMOVE_KEY` (op 1019) does not remove only the specified key across
    /// all Ignite 2.x server configurations; the returned previous value is
    /// discarded.
    #[must_use = "futures do nothing unless you `.await` them"]
    pub async fn remove(&self, key: IgniteValue) -> Result<()> {
        self.get_and_remove(key).await?;
        Ok(())
    }

    /// Replace the value only if the key is already present.
    /// Returns `true` if the value was replaced, `false` if the key was absent.
    pub async fn replace(&self, key: IgniteValue, value: IgniteValue) -> Result<bool> {
        let req_id = next_request_id();
        let target = self.route(&key, true).await;
        let payload = encode_cache_kv_req(
            op_code::CACHE_REPLACE,
            req_id,
            self.cache_id,
            &key,
            &value,
            self.tx_id(),
            self.expiry(),
        );
        let mut resp = self.send(req_id, payload, target).await?;
        read_bool(&mut resp).map_err(IgniteError::Protocol)
    }

    /// Atomically store a new value and return the previous value.
    /// Returns `IgniteValue::Null` if the key was not previously present.
    pub async fn get_and_put(&self, key: IgniteValue, value: IgniteValue) -> Result<IgniteValue> {
        let req_id = next_request_id();
        let target = self.route(&key, true).await;
        let payload = encode_cache_kv_req(
            op_code::CACHE_GET_AND_PUT,
            req_id,
            self.cache_id,
            &key,
            &value,
            self.tx_id(),
            self.expiry(),
        );
        let mut resp = self.send(req_id, payload, target).await?;
        decode_cache_value_response(&mut resp).map_err(IgniteError::Protocol)
    }

    /// Atomically remove a key and return its previous value.
    /// Returns `IgniteValue::Null` if the key was not present.
    pub async fn get_and_remove(&self, key: IgniteValue) -> Result<IgniteValue> {
        let req_id = next_request_id();
        let target = self.route(&key, true).await;
        let payload = encode_cache_key_req(
            op_code::CACHE_GET_AND_REMOVE,
            req_id,
            self.cache_id,
            &key,
            self.tx_id(),
            self.expiry(),
        );
        let mut resp = self.send(req_id, payload, target).await?;
        decode_cache_value_response(&mut resp).map_err(IgniteError::Protocol)
    }

    /// Replace the value only if the key is already present, and return the old value.
    /// Returns `IgniteValue::Null` if the key was absent (no change is made).
    pub async fn get_and_replace(
        &self,
        key: IgniteValue,
        value: IgniteValue,
    ) -> Result<IgniteValue> {
        let req_id = next_request_id();
        let target = self.route(&key, true).await;
        let payload = encode_cache_kv_req(
            op_code::CACHE_GET_AND_REPLACE,
            req_id,
            self.cache_id,
            &key,
            &value,
            self.tx_id(),
            self.expiry(),
        );
        let mut resp = self.send(req_id, payload, target).await?;
        decode_cache_value_response(&mut resp).map_err(IgniteError::Protocol)
    }

    /// Remove all specified keys from the cache.
    ///
    /// Implemented as individual [`Self::remove`] calls because the server-side
    /// bulk `CACHE_REMOVE_KEYS` operation does not behave as expected across all
    /// Ignite 2.x versions with the thin-client wire format.
    #[must_use = "futures do nothing unless you `.await` them"]
    pub async fn remove_all(&self, keys: Vec<IgniteValue>) -> Result<()> {
        for key in keys {
            self.remove(key).await?;
        }
        Ok(())
    }

    /// Return the number of entries in the cache.
    pub async fn get_size(&self) -> Result<i64> {
        let req_id = next_request_id();
        let payload = encode_cache_get_size(
            op_code::CACHE_GET_SIZE,
            req_id,
            self.cache_id,
            self.tx_id(),
            self.expiry(),
        );
        let mut resp = self.send(req_id, payload, None).await?;
        decode_cache_get_size_response(&mut resp).map_err(IgniteError::Protocol)
    }

    // ── Binary-object KV API ─────────────────────────────────────────────────

    /// Retrieve a value stored as an Ignite binary (complex) object, decoded
    /// via `V`'s [`ReadBinary`] impl. `key` is looked up as an
    /// `IgniteValue::String`. Returns `Ok(None)` if the key is absent.
    ///
    /// If the stored object uses a compact footer (the Java thin client's
    /// default), this fetches the object's registered
    /// [`BinaryType`](crate::binary::BinaryType) metadata
    /// (cached after the first lookup) to recover the schema's field ids —
    /// full (non-compact) footers carry field ids inline and need no lookup.
    ///
    /// A binary-object cache **value** (as opposed to a nested field of type
    /// Object) comes back from `CACHE_GET` as `IgniteValue::RawObject`, not
    /// `IgniteValue::Object`: real Ignite wraps top-level complex-object
    /// values in the "wrapped binary object" envelope (type code 27 —
    /// `[i32 len][payload][i32 offset]`), whereas `Object` is decoded
    /// straight off a raw `COMPLEX_OBJECT` (103) frame as seen when a complex
    /// object is *nested* inside another one's field data. Both cases carry
    /// the identical self-describing frame once unwrapped, so both are
    /// accepted here.
    ///
    /// Only supported on non-transactional (`Routed`) handles: transactional
    /// handles carry no channel registry to fetch metadata with, so this
    /// returns [`IgniteError::UnsupportedOnTransaction`] on a `Transaction::cache`
    /// handle.
    pub async fn get_binary<V: ReadBinary>(&self, key: &str) -> Result<Option<V>> {
        let registry = self.registry()?;
        let value = self.get(IgniteValue::String(key.to_string())).await?;
        let frame = match value {
            IgniteValue::Null => return Ok(None),
            IgniteValue::Object(o) => o.bytes,
            IgniteValue::RawObject(data) => Bytes::from(data),
            other => {
                return Err(IgniteError::Protocol(ProtocolError::TypeMismatch {
                    expected: "Object",
                    got: value_type_name(&other),
                }));
            }
        };
        let header = BinaryHeader::read(&mut frame.clone()).map_err(IgniteError::Protocol)?;
        let reader = if header.flags & flags::COMPACT_FOOTER != 0 {
            // Resolve the schema's field ids, refetching metadata once if our
            // cached view doesn't know this schema (a peer may have evolved the
            // type since we cached it).
            let resolver = RegistryResolver { registry, cache: &self.binary_types };
            let field_ids =
                schema_field_ids_with_refetch(&resolver, header.type_id, header.schema_id).await?;
            BinaryObjectReader::with_schema(frame, &field_ids)
        } else {
            BinaryObjectReader::new(frame)
        }
        .map_err(IgniteError::Protocol)?;

        // A nested `Object` field may itself be compact-footer-encoded (real
        // Ignite peers use compact footers throughout, not just at the top
        // level), which needs that nested type's own schema to decode —
        // discoverable only from its own bytes, not ahead of time like the
        // top-level type above. Walk the object graph now (while `.await` is
        // still available) and install what's found for the synchronous
        // `V::read` below. See `crate::binary::with_nested_schemas` for the
        // full rationale.
        let mut nested_schemas = HashMap::new();
        prefetch_nested_schemas(registry, &self.binary_types, &reader, &mut nested_schemas)
            .await?;

        crate::binary::with_nested_schemas(nested_schemas, || V::read(&reader)).map(Some)
    }

    /// Store `v` as an Ignite binary (complex) object under `key` (looked up
    /// as an `IgniteValue::String`), encoded via `V`'s [`WriteBinary`] impl.
    ///
    /// Registers `V::binary_type()` with the cluster first so other clients
    /// (e.g. a Java peer) can decode it. Ignite merges compatible metadata
    /// idempotently, so re-registering an existing type with a matching
    /// subset of fields is safe; if the server rejects the PUT (e.g. an
    /// incompatible field type for an already-registered field), this falls
    /// back to treating it as non-fatal when metadata already exists for the
    /// type id (some other writer already registered it) — logging a warning
    /// when that existing metadata is actually incompatible — and only surfaces
    /// the error when nothing is registered at all.
    ///
    /// Only supported on non-transactional (`Routed`) handles — see
    /// [`Self::get_binary`].
    #[must_use = "futures do nothing unless you `.await` them"]
    pub async fn put_binary<V: WriteBinary>(&self, key: &str, v: &V) -> Result<()> {
        let registry = self.registry()?;
        let type_def = V::binary_type();
        if let Err(e) = register_binary_type(registry, &self.binary_types, &type_def).await {
            // The server rejected the metadata registration. If the type is
            // entirely unregistered, propagate — nothing could decode this
            // object later. Otherwise proceed: a binary object is self-typed on
            // the wire, so the write itself is valid (verified by the
            // cross-language gate). But warn when the registered metadata is
            // actually incompatible with what we're writing, rather than
            // swallowing it silently, since SQL/queryable typing may disagree.
            match fetch_binary_type(registry, &self.binary_types, type_def.type_id).await? {
                None => return Err(e),
                Some(existing) if !metadata_compatible(&existing, &type_def) => {
                    tracing::warn!(
                        type_name = %type_def.type_name,
                        error = %e,
                        "binary-type metadata differs from the cluster's; writing anyway",
                    );
                }
                Some(_) => {}
            }
        }

        let obj = v.to_binary();
        self.put(
            IgniteValue::String(key.to_string()),
            IgniteValue::Object(obj),
        )
        .await
    }
}

// ─── Cache management helpers (used by IgniteClient) ─────────────────────────

/// Send `CACHE_GET_OR_CREATE_WITH_NAME` and return an `IgniteCache` handle.
/// `IgniteConnection::request` strips the response header automatically.
pub(crate) async fn get_or_create_cache_by_name(
    name: &str,
    registry: &Arc<ChannelRegistry>,
    affinity: &Arc<AffinityContext>,
    binary_types: &Arc<BinaryTypeCache>,
) -> Result<IgniteCache> {
    let req_id = next_request_id();
    let payload =
        encode_cache_create_with_name(op_code::CACHE_GET_OR_CREATE_WITH_NAME, req_id, name);

    let conn = registry.get(None).await?;
    // request() strips the response header; for this op the remaining body is empty.
    conn.request(req_id, payload)
        .await
        .map_err(IgniteError::Transport)?;

    let cid = crate::protocol::cache_id(name);
    Ok(IgniteCache::new(
        cid,
        registry.clone(),
        affinity.clone(),
        binary_types.clone(),
    ))
}

/// Send `CACHE_DESTROY` for the given cache name.
/// `IgniteConnection::request` strips the response header automatically.
pub(crate) async fn destroy_cache_by_name(name: &str, registry: &ChannelRegistry) -> Result<()> {
    let cid = crate::protocol::cache_id(name);
    let req_id = next_request_id();
    let payload = encode_cache_destroy_req(op_code::CACHE_DESTROY, req_id, cid);

    let conn = registry.get(None).await?;
    // request() strips the response header; for this op the remaining body is empty.
    conn.request(req_id, payload)
        .await
        .map_err(IgniteError::Transport)?;
    Ok(())
}