Line data    Source code

       1              : use std::io::ErrorKind;
       2              : 
       3              : use anyhow::Ok;
       4              : use pq_proto::{id_to_cancel_key, CancelKeyData};
       5              : use serde::{Deserialize, Serialize};
       6              : 
       7              : pub mod keyspace {
       8              :     pub const CANCEL_PREFIX: &str = "cancel";
       9              : }
      10              : 
      11            0 : #[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
      12              : pub(crate) enum KeyPrefix {
      13              :     #[serde(untagged)]
      14              :     Cancel(CancelKeyData),
      15              : }
      16              : 
      17              : impl KeyPrefix {
      18            1 :     pub(crate) fn build_redis_key(&self) -> String {
      19            1 :         match self {
      20            1 :             KeyPrefix::Cancel(key) => {
      21            1 :                 let hi = (key.backend_pid as u64) << 32;
      22            1 :                 let lo = (key.cancel_key as u64) & 0xffff_ffff;
      23            1 :                 let id = hi | lo;
      24            1 :                 let keyspace = keyspace::CANCEL_PREFIX;
      25            1 :                 format!("{keyspace}:{id:x}")
      26            1 :             }
      27            1 :         }
      28            1 :     }
      29              : 
      30              :     #[allow(dead_code)]
      31            2 :     pub(crate) fn as_str(&self) -> &'static str {
      32            2 :         match self {
      33            2 :             KeyPrefix::Cancel(_) => keyspace::CANCEL_PREFIX,
      34            2 :         }
      35            2 :     }
      36              : }
      37              : 
      38              : #[allow(dead_code)]
      39            1 : pub(crate) fn parse_redis_key(key: &str) -> anyhow::Result<KeyPrefix> {
      40            1 :     let (prefix, key_str) = key.split_once(':').ok_or_else(|| {
      41            0 :         anyhow::anyhow!(std::io::Error::new(
      42            0 :             ErrorKind::InvalidData,
      43            0 :             "missing prefix"
      44            0 :         ))
      45            1 :     })?;
      46              : 
      47            1 :     match prefix {
      48            1 :         keyspace::CANCEL_PREFIX => {
      49            1 :             let id = u64::from_str_radix(key_str, 16)?;
      50              : 
      51            1 :             Ok(KeyPrefix::Cancel(id_to_cancel_key(id)))
      52              :         }
      53            0 :         _ => Err(anyhow::anyhow!(std::io::Error::new(
      54            0 :             ErrorKind::InvalidData,
      55            0 :             "unknown prefix"
      56            0 :         ))),
      57              :     }
      58            1 : }
      59              : 
      60              : #[cfg(test)]
      61              : mod tests {
      62              :     use super::*;
      63              : 
      64              :     #[test]
      65            1 :     fn test_build_redis_key() {
      66            1 :         let cancel_key: KeyPrefix = KeyPrefix::Cancel(CancelKeyData {
      67            1 :             backend_pid: 12345,
      68            1 :             cancel_key: 54321,
      69            1 :         });
      70            1 : 
      71            1 :         let redis_key = cancel_key.build_redis_key();
      72            1 :         assert_eq!(redis_key, "cancel:30390000d431");
      73            1 :     }
      74              : 
      75              :     #[test]
      76            1 :     fn test_parse_redis_key() {
      77            1 :         let redis_key = "cancel:30390000d431";
      78            1 :         let key: KeyPrefix = parse_redis_key(redis_key).expect("Failed to parse key");
      79            1 : 
      80            1 :         let ref_key = CancelKeyData {
      81            1 :             backend_pid: 12345,
      82            1 :             cancel_key: 54321,
      83            1 :         };
      84            1 : 
      85            1 :         assert_eq!(key.as_str(), KeyPrefix::Cancel(ref_key).as_str());
      86            1 :         let KeyPrefix::Cancel(cancel_key) = key;
      87            1 :         assert_eq!(ref_key, cancel_key);
      88            1 :     }
      89              : }