RUST-830 Consolidate error label storage to Error::labels (#350)

Author: patrickfreed

src/cmap/conn/command.rs

@@ -1,11 +1,11 @@
-use serde::de::DeserializeOwned;
+use serde::{de::DeserializeOwned, Deserialize};
 
 use super::wire::Message;
 use crate::{
     bson::{Bson, Document},
     bson_util,
     client::{options::ServerApi, ClusterTime},
-    error::{CommandError, ErrorKind, Result},
+    error::{CommandError, Error, ErrorKind, Result},
     options::ServerAddress,
     selection_criteria::ReadPreference,
     ClientSession,
@@ -136,16 +136,19 @@ impl CommandResponse {
         }
     }
 
-    /// Retunrs a result indicating whether this response corresponds to a command failure.
+    /// Returns a result indicating whether this response corresponds to a command failure.
     pub(crate) fn validate(&self) -> Result<()> {
         if !self.is_success() {
-            let command_error: CommandError =
+            let error_response: CommandErrorResponse =
                 bson::from_bson(Bson::Document(self.raw_response.clone())).map_err(|_| {
                     ErrorKind::InvalidResponse {
                         message: "invalid server response".to_string(),
                     }
                 })?;
-            Err(ErrorKind::Command(command_error).into())
+            Err(Error::new(
+                ErrorKind::Command(error_response.command_error),
+                error_response.error_labels,
+            ))
         } else {
             Ok(())
         }
@@ -172,3 +175,12 @@ impl CommandResponse {
         &self.source
     }
 }
+
+#[derive(Deserialize, Debug)]
+struct CommandErrorResponse {
+    #[serde(rename = "errorLabels")]
+    error_labels: Option<Vec<String>>,
+
+    #[serde(flatten)]
+    command_error: CommandError,
+}

src/coll/mod.rs

@@ -1,11 +1,11 @@
 mod batch;
 pub mod options;
 
-use std::{borrow::Borrow, fmt, fmt::Debug, sync::Arc};
+use std::{borrow::Borrow, collections::HashSet, fmt, fmt::Debug, sync::Arc};
 
 use futures_util::stream::StreamExt;
 use serde::{
-    de::{DeserializeOwned, Error},
+    de::{DeserializeOwned, Error as DeError},
     Deserialize,
     Deserializer,
     Serialize,
@@ -17,7 +17,7 @@ use crate::{
     bson_util,
     client::session::TransactionState,
     concern::{ReadConcern, WriteConcern},
-    error::{convert_bulk_errors, BulkWriteError, BulkWriteFailure, ErrorKind, Result},
+    error::{convert_bulk_errors, BulkWriteError, BulkWriteFailure, Error, ErrorKind, Result},
     operation::{
         Aggregate,
         Count,
@@ -689,6 +689,7 @@ where
         let ordered = options.as_ref().and_then(|o| o.ordered).unwrap_or(true);
 
         let mut cumulative_failure: Option<BulkWriteFailure> = None;
+        let mut error_labels: HashSet<String> = Default::default();
         let mut cumulative_result: Option<InsertManyResult> = None;
 
         let mut n_attempted = 0;
@@ -736,11 +737,15 @@ where
                             failure_ref.write_concern_error = Some(write_concern_error.clone());
                         }
 
+                        error_labels.extend(e.labels);
+
                         if ordered {
-                            return Err(ErrorKind::BulkWrite(
-                                cumulative_failure.unwrap_or_else(BulkWriteFailure::new),
-                            )
-                            .into());
+                            return Err(Error::new(
+                                ErrorKind::BulkWrite(
+                                    cumulative_failure.unwrap_or_else(BulkWriteFailure::new),
+                                ),
+                                Some(error_labels),
+                            ));
                         }
                     }
                     _ => return Err(e),
@@ -749,7 +754,10 @@ where
         }
 
         match cumulative_failure {
-            Some(failure) => Err(ErrorKind::BulkWrite(failure).into()),
+            Some(failure) => Err(Error::new(
+                ErrorKind::BulkWrite(failure),
+                Some(error_labels),
+            )),
             None => Ok(cumulative_result.unwrap_or_else(InsertManyResult::new)),
         }
     }

src/error.rs

@@ -1,9 +1,6 @@
 //! Contains the `Error` and `Result` types that `mongodb` uses.
 
-use std::{
-    fmt::{self, Debug},
-    sync::Arc,
-};
+use std::{collections::HashSet, fmt::{self, Debug}, sync::Arc};
 
 use serde::Deserialize;
 use thiserror::Error;
@@ -43,10 +40,17 @@ pub type Result<T> = std::result::Result<T, Error>;
 pub struct Error {
     /// The type of error that occurred.
     pub kind: Box<ErrorKind>,
-    labels: Vec<String>,
+    pub(crate) labels: HashSet<String>,
 }
 
 impl Error {
+    pub(crate) fn new(kind: ErrorKind, labels: Option<impl IntoIterator<Item=String>>) -> Self {
+        Self {
+            kind: Box::new(kind),
+            labels: labels.map(|labels| labels.into_iter().collect()).unwrap_or_default(),
+        }
+    }
+
     pub(crate) fn pool_cleared_error(address: &ServerAddress) -> Self {
         ErrorKind::ConnectionPoolCleared {
             message: format!(
@@ -159,19 +163,8 @@ impl Error {
     }
 
     /// Returns the labels for this error.
-    pub fn labels(&self) -> &[String] {
-        match self.kind.as_ref() {
-            ErrorKind::Command(ref err) => &err.labels,
-            ErrorKind::Write(ref err) => match err {
-                WriteFailure::WriteError(_) => &self.labels,
-                WriteFailure::WriteConcernError(ref err) => &err.labels,
-            },
-            ErrorKind::BulkWrite(ref err) => match err.write_concern_error {
-                Some(ref err) => &err.labels,
-                None => &self.labels,
-            },
-            _ => &self.labels,
-        }
+    pub fn labels(&self) -> &HashSet<String> {
+        &self.labels
     }
 
     /// Whether this error contains the specified label.
@@ -184,30 +177,7 @@ impl Error {
     /// Adds the given label to this error.
     pub(crate) fn add_label<T: AsRef<str>>(&mut self, label: T) {
         let label = label.as_ref().to_string();
-        match self.kind.as_mut() {
-            ErrorKind::Command(err) => {
-                err.labels.push(label);
-            }
-            ErrorKind::Write(err) => match err {
-                WriteFailure::WriteError(_) => {
-                    self.labels.push(label);
-                }
-                WriteFailure::WriteConcernError(err) => {
-                    err.labels.push(label);
-                }
-            },
-            ErrorKind::BulkWrite(err) => match err.write_concern_error.as_mut() {
-                Some(write_concern_error) => {
-                    write_concern_error.labels.push(label);
-                }
-                None => {
-                    self.labels.push(label);
-                }
-            },
-            _ => {
-                self.labels.push(label);
-            }
-        }
+        self.labels.insert(label);
     }
 
     pub(crate) fn from_resolve_error(error: trust_dns_resolver::error::ResolveError) -> Self {
@@ -325,7 +295,7 @@ where
     fn from(err: E) -> Self {
         Self {
             kind: Box::new(err.into()),
-            labels: Vec::new(),
+            labels: Default::default(),
         }
     }
 }
@@ -447,10 +417,6 @@ pub struct CommandError {
     /// A description of the error that occurred.
     #[serde(rename = "errmsg")]
     pub message: String,
-
-    /// The error labels that the server returned.
-    #[serde(rename = "errorLabels", default)]
-    pub labels: Vec<String>,
 }
 
 impl fmt::Display for CommandError {
@@ -477,10 +443,6 @@ pub struct WriteConcernError {
     /// A document identifying the write concern setting related to the error.
     #[serde(rename = "errInfo")]
     pub details: Option<Document>,
-
-    /// The error labels that the server returned.
-    #[serde(rename = "errorLabels", default)]
-    pub labels: Vec<String>,
 }
 
 /// An error that occurred during a write operation that wasn't due to being unable to satisfy a
@@ -584,9 +546,9 @@ impl WriteFailure {
 /// untouched.
 pub(crate) fn convert_bulk_errors(error: Error) -> Error {
     match *error.kind {
-        ErrorKind::BulkWrite(ref bulk_failure) => {
-            match WriteFailure::from_bulk_failure(bulk_failure.clone()) {
-                Ok(failure) => ErrorKind::Write(failure).into(),
+        ErrorKind::BulkWrite(bulk_failure) => {
+            match WriteFailure::from_bulk_failure(bulk_failure) {
+                Ok(failure) => Error::new(ErrorKind::Write(failure), Some(error.labels)),
                 Err(e) => e,
             }
         }

src/operation/delete/test.rs

@@ -192,7 +192,6 @@ async fn handle_write_concern_failure() {
                     "wtimeout": 0,
                     "provenance": "clientSupplied"
                 } }),
-                labels: Vec::new(),
             };
             assert_eq!(wc_error, &expected_wc_err);
         }

src/operation/insert/test.rs

@@ -219,7 +219,6 @@ async fn handle_write_failure() {
                     "wtimeout": 0,
                     "provenance": "clientSupplied"
                 } }),
-                labels: Vec::new(),
             };
             assert_eq!(write_concern_error, expected_wc_err);
         }

src/operation/mod.rs

@@ -148,14 +148,18 @@ struct EmptyBody {}
 struct WriteConcernOnlyBody {
     #[serde(rename = "writeConcernError")]
     write_concern_error: Option<WriteConcernError>,
+
+    #[serde(rename = "errorLabels")]
+    labels: Option<Vec<String>>,
 }
 
 impl WriteConcernOnlyBody {
     fn validate(&self) -> Result<()> {
         match self.write_concern_error {
-            Some(ref wc_error) => {
-                Err(ErrorKind::Write(WriteFailure::WriteConcernError(wc_error.clone())).into())
-            }
+            Some(ref wc_error) => Err(Error::new(
+                ErrorKind::Write(WriteFailure::WriteConcernError(wc_error.clone())),
+                self.labels.clone(),
+            )),
             None => Ok(()),
         }
     }
@@ -184,24 +188,15 @@ impl<T> WriteResponseBody<T> {
             return Ok(());
         };
 
-        // Error labels for WriteConcernErrors are sent from the server in a separate field.
-        let write_concern_error = match self.write_concern_error {
-            Some(ref write_concern_error) => {
-                let mut write_concern_error = write_concern_error.clone();
-                if let Some(ref labels) = self.labels {
-                    write_concern_error.labels.append(&mut labels.clone());
-                }
-                Some(write_concern_error)
-            }
-            None => None,
-        };
-
         let failure = BulkWriteFailure {
             write_errors: self.write_errors.clone(),
-            write_concern_error,
+            write_concern_error: self.write_concern_error.clone(),
         };
 
-        Err(ErrorKind::BulkWrite(failure).into())
+        Err(Error::new(
+            ErrorKind::BulkWrite(failure),
+            self.labels.clone(),
+        ))
     }
 }
 

src/operation/update/test.rs

@@ -276,7 +276,6 @@ async fn handle_write_concern_failure() {
                     "wtimeout": 0,
                     "provenance": "clientSupplied"
                 } }),
-                labels: Vec::new(),
             };
             assert_eq!(wc_error, &expected_wc_err);
         }

src/test/spec/retryable_writes/mod.rs

@@ -2,6 +2,7 @@ mod test_file;
 
 use std::time::Duration;
 
+use bson::Bson;
 use futures::stream::TryStreamExt;
 use semver::VersionReq;
 use tokio::sync::{RwLockReadGuard, RwLockWriteGuard};
@@ -27,7 +28,7 @@ use crate::{
 #[cfg_attr(feature = "async-std-runtime", async_std::test)]
 async fn run_spec_tests() {
     async fn run_test(test_file: TestFile) {
-        for test_case in test_file.tests {
+        for mut test_case in test_file.tests {
             if test_case.operation.name == "bulkWrite" {
                 continue;
             }
@@ -63,7 +64,15 @@ async fn run_spec_tests() {
                     .expect(&test_case.description);
             }
 
-            if let Some(ref fail_point) = test_case.fail_point {
+            if let Some(ref mut fail_point) = test_case.fail_point {
+                // TODO: DRIVERS-1385 remove this logic for moving errorLabels to the top level.
+                if let Some(Bson::Document(data)) = fail_point.get_mut("data") {
+                    if let Some(Bson::Document(wc_error)) = data.get_mut("writeConcernError") {
+                        if let Some(labels) = wc_error.remove("errorLabels") {
+                            data.insert("errorLabels", labels);
+                        }
+                    }
+                }
                 client
                     .database("admin")
                     .run_command(fail_point.clone(), None)
@@ -89,9 +98,14 @@ async fn run_spec_tests() {
                 match expected_result {
                     Result::Value(value) => {
                         let description = &test_case.description;
-                        let result = result.unwrap().unwrap_or_else(|| {
-                            panic!("{:?}: operation should succeed", description)
-                        });
+                        let result = result
+                            .unwrap_or_else(|e| {
+                                panic!(
+                                    "{:?}: operation should succeed, got error: {}",
+                                    description, e
+                                )
+                            })
+                            .unwrap();
                         assert_matches(&result, &value, Some(description));
                     }
                     Result::Labels(expected_labels) => {

src/test/spec/v2_runner/mod.rs

@@ -259,13 +259,13 @@ pub async fn run_v2_test(test_file: TestFile) {
                             assert_eq!(error_code_name, code_name);
                         }
                         if let Some(error_labels_contain) = operation_error.error_labels_contain {
-                            let labels = error.labels().to_vec();
+                            let labels = error.labels();
                             error_labels_contain
                                 .iter()
                                 .for_each(|label| assert!(labels.contains(label)));
                         }
                         if let Some(error_labels_omit) = operation_error.error_labels_omit {
-                            let labels = error.labels().to_vec();
+                            let labels = error.labels();
                             error_labels_omit
                                 .iter()
                                 .for_each(|label| assert!(!labels.contains(label)));