Make GrpcErrorsAsFailures more flexible

examples/tonic-key-value-store/src/main.rs

@@ -29,6 +29,7 @@ use tokio_stream::{
 use tonic::{async_trait, body::BoxBody, transport::Channel, Code, Request, Response, Status};
 use tower::{BoxError, Service, ServiceBuilder};
 use tower_http::{
+    classify::{GrpcCode, GrpcErrorsAsFailures, SharedClassifier},
     compression::CompressionLayer,
     decompression::DecompressionLayer,
     sensitive_headers::SetSensitiveHeadersLayer,
@@ -165,6 +166,12 @@ async fn serve_forever(listener: TcpListener) -> Result<(), Box<dyn std::error::
     // Build our tonic `Service`
     let service = key_value_store_server::KeyValueStoreServer::new(ServerImpl { db, tx });
 
+    // Response classifier that doesn't consider `Ok`, `Invalid Argument`, or `Not Found` as
+    // failures
+    let classifier = GrpcErrorsAsFailures::new()
+        .with_success(GrpcCode::InvalidArgument)
+        .with_success(GrpcCode::NotFound);
+
     // Build our middleware stack
     let layer = ServiceBuilder::new()
         // Set a timeout
@@ -175,7 +182,8 @@ async fn serve_forever(listener: TcpListener) -> Result<(), Box<dyn std::error::
         .layer(SetSensitiveHeadersLayer::new(once(header::AUTHORIZATION)))
         // Log all requests and responses
         .layer(
-            TraceLayer::new_for_grpc().make_span_with(DefaultMakeSpan::new().include_headers(true)),
+            TraceLayer::new(SharedClassifier::new(classifier))
+                .make_span_with(DefaultMakeSpan::new().include_headers(true)),
         )
         .into_inner();
 

tower-http/CHANGELOG.md

@@ -33,6 +33,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   adding middleware from tower-http.
 - Add `SetRequestId` and `PropagateRequestId` middleware ([#150])
 - Add `LatencyUnit::Seconds` for formatting latencies as seconds.
+- Support customizing which status codes are considered failures by `GrpcErrorsAsFailures` ([#189])
 
 [#124]: https://github.com/tower-rs/tower-http/pull/124
 [#148]: https://github.com/tower-rs/tower-http/pull/148
@@ -44,6 +45,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#172]: https://github.com/tower-rs/tower-http/pull/172
 [#182]: https://github.com/tower-rs/tower-http/pull/182
 [#187]: https://github.com/tower-rs/tower-http/pull/187
+[#189]: https://github.com/tower-rs/tower-http/pull/189
 
 # 0.1.2 (November 13, 2021)
 

tower-http/Cargo.toml

@@ -13,6 +13,7 @@ keywords = ["io", "async", "futures", "service", "http"]
 rust-version = "1.51"
 
 [dependencies]
+bitflags = "1.0"
 bytes = "1"
 futures-core = "0.3"
 futures-util = { version = "0.3", default_features = false, features = [] }