scylla_cql/frame/
protocol_features.rs

1//! Implementation and negotiation of extensions to the CQL protocol.
2
3use std::borrow::Cow;
4use std::collections::HashMap;
5
6const RATE_LIMIT_ERROR_EXTENSION: &str = "SCYLLA_RATE_LIMIT_ERROR";
7/// The extension used to add metadata for LWT optimization.
8/// See [ProtocolFeatures::lwt_optimization_meta_bit_mask] and
9/// [related issue](https://github.com/scylladb/scylla-rust-driver/issues/100)
10/// for more details.
11pub const SCYLLA_LWT_ADD_METADATA_MARK_EXTENSION: &str = "SCYLLA_LWT_ADD_METADATA_MARK";
12/// The key of the single entry of the LWT optimization extension,
13/// which entry is a bit mask for the frame flags used to mark LWT frames.
14pub const LWT_OPTIMIZATION_META_BIT_MASK_KEY: &str = "LWT_OPTIMIZATION_META_BIT_MASK";
15const TABLETS_ROUTING_V1_KEY: &str = "TABLETS_ROUTING_V1";
16
17/// Which protocol extensions are supported by the server.
18///
19/// This is used to inform the server about the features that the client supports,
20/// so that the server can adjust its behavior accordingly.
21///
22/// So to draw the picture:
23/// - server responds to an `OPTIONS` frame with a `SUPPORTED` frame with the list of
24///   protocol features it supports;
25/// - client parses the `SUPPORTED` frame by extracting extensions it recognizes (supports)
26///   and creates a `ProtocolFeatures` instance;
27/// - from now on, client uses this instance to determine how to handle certain frames,
28///   e.g. whether to expect a rate limit error or how to handle LWT operations;
29/// - client also adds the features it supports to the `STARTUP` frame and sends it to
30///   the server, which finishes the extensions negotiation process.
31#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
32#[non_exhaustive]
33pub struct ProtocolFeatures {
34    /// The error code to use for rate limit errors, if negotiated.
35    pub rate_limit_error: Option<i32>,
36
37    /// The bit mask used for the LWT optimization, if negotiated.
38    /// This is used to mark PREPARED response frames as related to LWT operations
39    /// in order to to optimize the handling of LWT requests.
40    ///
41    /// The mask is ANDed with the flags of the PREPARED response frame,
42    /// and if the result is equal to the mask, it means that the frame is related
43    /// to an LWT operation.
44    pub lwt_optimization_meta_bit_mask: Option<u32>,
45
46    /// Whether the server supports tablets routing v1.
47    pub tablets_v1_supported: bool,
48}
49
50// TODO: Log information about options which failed to parse
51
52impl ProtocolFeatures {
53    /// Parses the supported protocol features from the `supported` map.
54    pub fn parse_from_supported(supported: &HashMap<String, Vec<String>>) -> Self {
55        Self {
56            rate_limit_error: Self::maybe_parse_rate_limit_error(supported),
57            lwt_optimization_meta_bit_mask: Self::maybe_parse_lwt_optimization_meta_bit_mask(
58                supported,
59            ),
60            tablets_v1_supported: Self::check_tablets_routing_v1_support(supported),
61        }
62    }
63
64    fn maybe_parse_rate_limit_error(supported: &HashMap<String, Vec<String>>) -> Option<i32> {
65        let vals = supported.get(RATE_LIMIT_ERROR_EXTENSION)?;
66        let code_str = Self::get_cql_extension_field(vals.as_slice(), "ERROR_CODE")?;
67        code_str.parse::<i32>().ok()
68    }
69
70    fn maybe_parse_lwt_optimization_meta_bit_mask(
71        supported: &HashMap<String, Vec<String>>,
72    ) -> Option<u32> {
73        let vals = supported.get(SCYLLA_LWT_ADD_METADATA_MARK_EXTENSION)?;
74        let mask_str =
75            Self::get_cql_extension_field(vals.as_slice(), LWT_OPTIMIZATION_META_BIT_MASK_KEY)?;
76        mask_str.parse::<u32>().ok()
77    }
78
79    fn check_tablets_routing_v1_support(supported: &HashMap<String, Vec<String>>) -> bool {
80        supported.contains_key(TABLETS_ROUTING_V1_KEY)
81    }
82
83    // Looks up a field which starts with `key=` and returns the rest
84    fn get_cql_extension_field<'a>(vals: &'a [String], key: &str) -> Option<&'a str> {
85        vals.iter()
86            .find_map(|v| v.as_str().strip_prefix(key)?.strip_prefix('='))
87    }
88
89    /// Adds the protocol features as STARTUP options.
90    pub fn add_startup_options(&self, options: &mut HashMap<Cow<'_, str>, Cow<'_, str>>) {
91        if self.rate_limit_error.is_some() {
92            options.insert(Cow::Borrowed(RATE_LIMIT_ERROR_EXTENSION), Cow::Borrowed(""));
93        }
94        if let Some(mask) = self.lwt_optimization_meta_bit_mask {
95            options.insert(
96                Cow::Borrowed(SCYLLA_LWT_ADD_METADATA_MARK_EXTENSION),
97                Cow::Owned(format!("{LWT_OPTIMIZATION_META_BIT_MASK_KEY}={mask}")),
98            );
99        }
100
101        if self.tablets_v1_supported {
102            options.insert(Cow::Borrowed(TABLETS_ROUTING_V1_KEY), Cow::Borrowed(""));
103        }
104    }
105
106    /// Checks if the given flags of a PREPARED response contain the LWT optimization mark.
107    ///
108    /// If the extension was not negotiated, it conservatively returns `false`.
109    pub fn prepared_flags_contain_lwt_mark(&self, flags: u32) -> bool {
110        self.lwt_optimization_meta_bit_mask
111            .map(|mask| (flags & mask) == mask)
112            .unwrap_or(false)
113    }
114}
scylla_cql/frame/protocol_features.rs

scylla_cql/frame/
protocol_features.rs
