fix: make push to drainer generic and add application metrics for KV (#2563)

This commit is contained in:
Kartikeya Hegde
2023-10-17 11:21:18 +00:00
committed by GitHub
parent cecea8718a
commit 274a78343e
11 changed files with 258 additions and 334 deletions

View File

@ -3,9 +3,10 @@ use std::{fmt::Debug, sync::Arc};
use common_utils::errors::CustomResult;
use redis_interface::errors::RedisError;
use router_derive::TryGetEnumVariant;
use router_env::logger;
use serde::de;
use crate::{consts, KVRouterStore};
use crate::{consts, metrics, store::kv::TypedSql, KVRouterStore};
pub trait KvStorePartition {
fn partition_number(key: PartitionKey<'_>, num_partitions: u8) -> u32 {
@ -48,10 +49,11 @@ pub trait RedisConnInterface {
) -> error_stack::Result<Arc<redis_interface::RedisConnectionPool>, RedisError>;
}
/// An enum to represent what operation to do on
pub enum KvOperation<'a, S: serde::Serialize + Debug> {
Hset((&'a str, String)),
SetNx(S),
HSetNx(&'a str, S),
Hset((&'a str, String), TypedSql),
SetNx(&'a S, TypedSql),
HSetNx(&'a str, &'a S, TypedSql),
HGet(&'a str),
Get,
Scan(&'a str),
@ -68,6 +70,22 @@ pub enum KvResult<T: de::DeserializeOwned> {
Scan(Vec<T>),
}
impl<T> std::fmt::Display for KvOperation<'_, T>
where
T: serde::Serialize + Debug,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
KvOperation::Hset(_, _) => f.write_str("Hset"),
KvOperation::SetNx(_, _) => f.write_str("Setnx"),
KvOperation::HSetNx(_, _, _) => f.write_str("HSetNx"),
KvOperation::HGet(_) => f.write_str("Hget"),
KvOperation::Get => f.write_str("Get"),
KvOperation::Scan(_) => f.write_str("Scan"),
}
}
}
pub async fn kv_wrapper<'a, T, D, S>(
store: &KVRouterStore<D>,
op: KvOperation<'a, S>,
@ -76,45 +94,101 @@ pub async fn kv_wrapper<'a, T, D, S>(
where
T: de::DeserializeOwned,
D: crate::database::store::DatabaseStore,
S: serde::Serialize + Debug,
S: serde::Serialize + Debug + KvStorePartition,
{
let redis_conn = store.get_redis_conn()?;
let key = key.as_ref();
let type_name = std::any::type_name::<T>();
let operation = op.to_string();
match op {
KvOperation::Hset(value) => {
redis_conn
.set_hash_fields(key, value, Some(consts::KV_TTL))
.await?;
Ok(KvResult::Hset(()))
let partition_key = PartitionKey::MerchantIdPaymentIdCombination { combination: key };
let result = async {
match op {
KvOperation::Hset(value, sql) => {
logger::debug!("Operation: {operation} value: {value:?}");
redis_conn
.set_hash_fields(key, value, Some(consts::KV_TTL))
.await?;
store
.push_to_drainer_stream::<S>(sql, partition_key)
.await?;
Ok(KvResult::Hset(()))
}
KvOperation::HGet(field) => {
let result = redis_conn
.get_hash_field_and_deserialize(key, field, type_name)
.await?;
Ok(KvResult::HGet(result))
}
KvOperation::Scan(pattern) => {
let result: Vec<T> = redis_conn.hscan_and_deserialize(key, pattern, None).await?;
Ok(KvResult::Scan(result))
}
KvOperation::HSetNx(field, value, sql) => {
logger::debug!("Operation: {operation} value: {value:?}");
let result = redis_conn
.serialize_and_set_hash_field_if_not_exist(
key,
field,
value,
Some(consts::KV_TTL),
)
.await?;
if matches!(result, redis_interface::HsetnxReply::KeySet) {
store
.push_to_drainer_stream::<S>(sql, partition_key)
.await?;
}
Ok(KvResult::HSetNx(result))
}
KvOperation::SetNx(value, sql) => {
logger::debug!("Operation: {operation} value: {value:?}");
let result = redis_conn
.serialize_and_set_key_if_not_exist(key, value, Some(consts::KV_TTL.into()))
.await?;
if matches!(result, redis_interface::SetnxReply::KeySet) {
store
.push_to_drainer_stream::<S>(sql, partition_key)
.await?;
}
Ok(KvResult::SetNx(result))
}
KvOperation::Get => {
let result = redis_conn.get_and_deserialize_key(key, type_name).await?;
Ok(KvResult::Get(result))
}
}
KvOperation::HGet(field) => {
let result = redis_conn
.get_hash_field_and_deserialize(key, field, type_name)
.await?;
Ok(KvResult::HGet(result))
}
KvOperation::Scan(pattern) => {
let result: Vec<T> = redis_conn.hscan_and_deserialize(key, pattern, None).await?;
Ok(KvResult::Scan(result))
}
KvOperation::HSetNx(field, value) => {
let result = redis_conn
.serialize_and_set_hash_field_if_not_exist(key, field, value, Some(consts::KV_TTL))
.await?;
Ok(KvResult::HSetNx(result))
}
KvOperation::SetNx(value) => {
let result = redis_conn
.serialize_and_set_key_if_not_exist(key, value, Some(consts::KV_TTL.into()))
.await?;
Ok(KvResult::SetNx(result))
}
KvOperation::Get => {
let result = redis_conn.get_and_deserialize_key(key, type_name).await?;
Ok(KvResult::Get(result))
}
}
};
result
.await
.map(|result| {
logger::debug!("KvOperation {operation} succeeded");
let keyvalue = router_env::opentelemetry::KeyValue::new("operation", operation.clone());
metrics::KV_OPERATION_SUCCESSFUL.add(&metrics::CONTEXT, 1, &[keyvalue]);
result
})
.map_err(|err| {
logger::error!("KvOperation for {operation} failed with {err:?}");
let keyvalue = router_env::opentelemetry::KeyValue::new("operation", operation);
metrics::KV_OPERATION_FAILED.add(&metrics::CONTEXT, 1, &[keyvalue]);
err
})
}