Line data    Source code

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