Merge pull request #743 from wprzytula/fix-routing-info-consistency

Fix RoutingInfo consistency

Author: piodul

scylla-cql/benches/benchmark.rs

@@ -1,17 +1,19 @@
+use std::borrow::Cow;
+
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 
-use scylla_cql::frame::request::Request;
+use scylla_cql::frame::request::SerializableRequest;
 use scylla_cql::frame::value::SerializedValues;
 use scylla_cql::frame::value::ValueList;
 use scylla_cql::frame::{request::query, Compression, SerializedRequest};
 
 fn make_query<'a>(contents: &'a str, values: &'a SerializedValues) -> query::Query<'a> {
     query::Query {
-        contents,
+        contents: Cow::Borrowed(contents),
         parameters: query::QueryParameters {
             consistency: scylla_cql::Consistency::LocalQuorum,
             serial_consistency: None,
-            values,
+            values: Cow::Borrowed(values),
             page_size: None,
             paging_state: None,
             timestamp: None,
@@ -25,7 +27,7 @@ fn serialized_request_make_bench(c: &mut Criterion) {
         ("INSERT foo INTO ks.table_name (?)", &(1234,).serialized().unwrap()),
         ("INSERT foo, bar, baz INTO ks.table_name (?, ?, ?)", &(1234, "a value", "i am storing a string").serialized().unwrap()),
         (
-            "INSERT foo, bar, baz, boop, blah INTO longer_keyspace.a_big_table_name (?, ?, ?, ?, 1000)", 
+            "INSERT foo, bar, baz, boop, blah INTO longer_keyspace.a_big_table_name (?, ?, ?, ?, 1000)",
             &(1234, "a value", "i am storing a string", "dc0c8cd7-d954-47c1-8722-a857941c43fb").serialized().unwrap()
         ),
     ];

scylla-cql/src/frame/mod.rs

@@ -16,7 +16,7 @@ use uuid::Uuid;
 
 use std::convert::TryFrom;
 
-use request::Request;
+use request::SerializableRequest;
 use response::ResponseOpcode;
 
 const HEADER_SIZE: usize = 9;
@@ -60,7 +60,7 @@ pub struct SerializedRequest {
 }
 
 impl SerializedRequest {
-    pub fn make<R: Request>(
+    pub fn make<R: SerializableRequest>(
         req: &R,
         compression: Option<Compression>,
         tracing: bool,

scylla-cql/src/frame/request/auth_response.rs

@@ -1,15 +1,15 @@
 use crate::frame::frame_errors::ParseError;
 use bytes::BufMut;
 
-use crate::frame::request::{Request, RequestOpcode};
+use crate::frame::request::{RequestOpcode, SerializableRequest};
 use crate::frame::types::write_bytes_opt;
 
 // Implements Authenticate Response
 pub struct AuthResponse {
     pub response: Option<Vec<u8>>,
 }
 
-impl Request for AuthResponse {
+impl SerializableRequest for AuthResponse {
     const OPCODE: RequestOpcode = RequestOpcode::AuthResponse;
 
     fn serialize(&self, buf: &mut impl BufMut) -> Result<(), ParseError> {

scylla-cql/src/frame/request/batch.rs

@@ -1,24 +1,28 @@
-use crate::frame::{frame_errors::ParseError, value::BatchValuesIterator};
-use bytes::{BufMut, Bytes};
-use std::convert::TryInto;
+use bytes::{Buf, BufMut};
+use std::{borrow::Cow, convert::TryInto};
 
 use crate::frame::{
-    request::{Request, RequestOpcode},
+    frame_errors::ParseError,
+    request::{RequestOpcode, SerializableRequest},
     types,
-    value::BatchValues,
+    value::{BatchValues, BatchValuesIterator, SerializedValues},
 };
 
+use super::DeserializableRequest;
+
 // Batch flags
 const FLAG_WITH_SERIAL_CONSISTENCY: u8 = 0x10;
 const FLAG_WITH_DEFAULT_TIMESTAMP: u8 = 0x20;
+const ALL_FLAGS: u8 = FLAG_WITH_SERIAL_CONSISTENCY | FLAG_WITH_DEFAULT_TIMESTAMP;
 
-pub struct Batch<'a, StatementsIter, Values>
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
+pub struct Batch<'b, Statement, Values>
 where
-    StatementsIter: Iterator<Item = BatchStatement<'a>> + Clone,
+    BatchStatement<'b>: From<&'b Statement>,
+    Statement: Clone,
     Values: BatchValues,
 {
-    pub statements: StatementsIter,
-    pub statements_count: usize,
+    pub statements: Cow<'b, [Statement]>,
     pub batch_type: BatchType,
     pub consistency: types::Consistency,
     pub serial_consistency: Option<types::SerialConsistency>,
@@ -28,21 +32,46 @@ where
 
 /// The type of a batch.
 #[derive(Clone, Copy)]
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
 pub enum BatchType {
     Logged = 0,
     Unlogged = 1,
     Counter = 2,
 }
 
-#[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Ord)]
+pub struct BatchTypeParseError {
+    value: u8,
+}
+
+impl From<BatchTypeParseError> for ParseError {
+    fn from(err: BatchTypeParseError) -> Self {
+        Self::BadIncomingData(format!("Bad BatchType value: {}", err.value))
+    }
+}
+
+impl TryFrom<u8> for BatchType {
+    type Error = BatchTypeParseError;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(Self::Logged),
+            1 => Ok(Self::Unlogged),
+            2 => Ok(Self::Counter),
+            _ => Err(BatchTypeParseError { value }),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord)]
 pub enum BatchStatement<'a> {
-    Query { text: &'a str },
-    Prepared { id: &'a Bytes },
+    Query { text: Cow<'a, str> },
+    Prepared { id: Cow<'a, [u8]> },
 }
 
-impl<'a, StatementsIter, Values> Request for Batch<'a, StatementsIter, Values>
+impl<Statement, Values> SerializableRequest for Batch<'_, Statement, Values>
 where
-    StatementsIter: Iterator<Item = BatchStatement<'a>> + Clone,
+    for<'s> BatchStatement<'s>: From<&'s Statement>,
+    Statement: Clone,
     Values: BatchValues,
 {
     const OPCODE: RequestOpcode = RequestOpcode::Batch;
@@ -52,7 +81,7 @@ where
         buf.put_u8(self.batch_type as u8);
 
         // Serializing queries
-        types::write_short(self.statements_count.try_into()?, buf);
+        types::write_short(self.statements.len().try_into()?, buf);
 
         let counts_mismatch_err = |n_values: usize, n_statements: usize| {
             ParseError::BadDataToSerialize(format!(
@@ -62,26 +91,27 @@ where
         };
         let mut n_serialized_statements = 0usize;
         let mut value_lists = self.values.batch_values_iter();
-        for (idx, statement) in self.statements.clone().enumerate() {
-            statement.serialize(buf)?;
+        for (idx, statement) in self.statements.iter().enumerate() {
+            BatchStatement::from(statement).serialize(buf)?;
             value_lists
                 .write_next_to_request(buf)
-                .ok_or_else(|| counts_mismatch_err(idx, self.statements.clone().count()))??;
+                .ok_or_else(|| counts_mismatch_err(idx, self.statements.len()))??;
             n_serialized_statements += 1;
         }
+        // At this point, we have all statements serialized. If any values are still left, we have a mismatch.
         if value_lists.skip_next().is_some() {
             return Err(counts_mismatch_err(
-                std::iter::from_fn(|| value_lists.skip_next()).count() + 1,
+                n_serialized_statements + 1 /*skipped above*/ + value_lists.count(),
                 n_serialized_statements,
             ));
         }
-        if n_serialized_statements != self.statements_count {
+        if n_serialized_statements != self.statements.len() {
             // We want to check this to avoid propagating an invalid construction of self.statements_count as a
             // hard-to-debug silent fail
             return Err(ParseError::BadDataToSerialize(format!(
                 "Invalid Batch constructed: not as many statements serialized as announced \
                     (batch.statement_count: {announced_statement_count}, {n_serialized_statements}",
-                announced_statement_count = self.statements_count
+                announced_statement_count = self.statements.len()
             )));
         }
 
@@ -110,19 +140,115 @@ where
     }
 }
 
+impl BatchStatement<'_> {
+    fn deserialize(buf: &mut &[u8]) -> Result<Self, ParseError> {
+        let kind = buf.get_u8();
+        match kind {
+            0 => {
+                let text = Cow::Owned(types::read_long_string(buf)?.to_owned());
+                Ok(BatchStatement::Query { text })
+            }
+            1 => {
+                let id = types::read_short_bytes(buf)?.to_vec().into();
+                Ok(BatchStatement::Prepared { id })
+            }
+            _ => Err(ParseError::BadIncomingData(format!(
+                "Unexpected batch statement kind: {}",
+                kind
+            ))),
+        }
+    }
+}
+
 impl BatchStatement<'_> {
     fn serialize(&self, buf: &mut impl BufMut) -> Result<(), ParseError> {
         match self {
-            BatchStatement::Query { text } => {
+            Self::Query { text } => {
                 buf.put_u8(0);
                 types::write_long_string(text, buf)?;
             }
-            BatchStatement::Prepared { id } => {
+            Self::Prepared { id } => {
                 buf.put_u8(1);
-                types::write_short_bytes(&id[..], buf)?;
+                types::write_short_bytes(id, buf)?;
             }
         }
 
         Ok(())
     }
 }
+
+impl<'s, 'b> From<&'s BatchStatement<'b>> for BatchStatement<'s> {
+    fn from(value: &'s BatchStatement) -> Self {
+        match value {
+            BatchStatement::Query { text } => BatchStatement::Query { text: text.clone() },
+            BatchStatement::Prepared { id } => BatchStatement::Prepared { id: id.clone() },
+        }
+    }
+}
+
+impl<'b> DeserializableRequest for Batch<'b, BatchStatement<'b>, Vec<SerializedValues>> {
+    fn deserialize(buf: &mut &[u8]) -> Result<Self, ParseError> {
+        let batch_type = buf.get_u8().try_into()?;
+
+        let statements_count: usize = types::read_short(buf)?.try_into()?;
+        let statements_with_values = (0..statements_count)
+            .map(|_| {
+                let batch_statement = BatchStatement::deserialize(buf)?;
+
+                // As stated in CQL protocol v4 specification, values names in Batch are broken and should be never used.
+                let values = SerializedValues::new_from_frame(buf, false)?;
+
+                Ok((batch_statement, values))
+            })
+            .collect::<Result<Vec<_>, ParseError>>()?;
+
+        let consistency = match types::read_consistency(buf)? {
+            types::LegacyConsistency::Regular(reg) => Ok(reg),
+            types::LegacyConsistency::Serial(ser) => Err(ParseError::BadIncomingData(format!(
+                "Expected regular Consistency, got SerialConsistency {}",
+                ser
+            ))),
+        }?;
+
+        let flags = buf.get_u8();
+        let unknown_flags = flags & (!ALL_FLAGS);
+        if unknown_flags != 0 {
+            return Err(ParseError::BadIncomingData(format!(
+                "Specified flags are not recognised: {:02x}",
+                unknown_flags
+            )));
+        }
+        let serial_consistency_flag = (flags & FLAG_WITH_SERIAL_CONSISTENCY) != 0;
+        let default_timestamp_flag = (flags & FLAG_WITH_DEFAULT_TIMESTAMP) != 0;
+
+        let serial_consistency = serial_consistency_flag
+            .then(|| types::read_consistency(buf))
+            .transpose()?
+            .map(|legacy_consistency| match legacy_consistency {
+                types::LegacyConsistency::Regular(reg) => {
+                    Err(ParseError::BadIncomingData(format!(
+                        "Expected SerialConsistency, got regular Consistency {}",
+                        reg
+                    )))
+                }
+                types::LegacyConsistency::Serial(ser) => Ok(ser),
+            })
+            .transpose()?;
+
+        let timestamp = default_timestamp_flag
+            .then(|| types::read_long(buf))
+            .transpose()?;
+
+        let (statements, values): (Vec<BatchStatement>, Vec<SerializedValues>) =
+            statements_with_values.into_iter().unzip();
+
+        Ok(Self {
+            batch_type,
+            consistency,
+            serial_consistency,
+            timestamp,
+            statements: Cow::Owned(statements),
+            values,
+        })
+    }
+}

scylla-cql/src/frame/request/execute.rs

@@ -2,16 +2,19 @@ use crate::frame::frame_errors::ParseError;
 use bytes::{BufMut, Bytes};
 
 use crate::{
-    frame::request::{query, Request, RequestOpcode},
+    frame::request::{query, RequestOpcode, SerializableRequest},
     frame::types,
 };
 
+use super::{query::QueryParameters, DeserializableRequest};
+
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
 pub struct Execute<'a> {
     pub id: Bytes,
     pub parameters: query::QueryParameters<'a>,
 }
 
-impl Request for Execute<'_> {
+impl SerializableRequest for Execute<'_> {
     const OPCODE: RequestOpcode = RequestOpcode::Execute;
 
     fn serialize(&self, buf: &mut impl BufMut) -> Result<(), ParseError> {
@@ -23,3 +26,12 @@ impl Request for Execute<'_> {
         Ok(())
     }
 }
+
+impl<'e> DeserializableRequest for Execute<'e> {
+    fn deserialize(buf: &mut &[u8]) -> Result<Self, ParseError> {
+        let id = types::read_short_bytes(buf)?.to_vec().into();
+        let parameters = QueryParameters::deserialize(buf)?;
+
+        Ok(Self { id, parameters })
+    }
+}