Add service counter for service level server metrics (#4540)

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here -->

Adds a service counter for guarded decrement of outstanding requests. In
case something happens before the response backswing, we will want to
ensure it decrements.

## Description
<!--- Describe your changes in detail -->

## Testing
<!--- Please describe in detail how you tested your changes -->
<!--- Include details of your testing environment, and the tests you ran
to -->
<!--- see how your change affects other areas of the code, etc. -->

## Checklist
<!--- If a checkbox below is not applicable, then please DELETE it
rather than leaving it unchecked -->
- [ ] For changes to the smithy-rs codegen or runtime crates, I have
created a changelog entry Markdown file in the `.changelog` directory,
specifying "client," "server," or both in the `applies_to` key.
- [ ] For changes to the AWS SDK, generated SDK code, or SDK runtime
crates, I have created a changelog entry Markdown file in the
`.changelog` directory, specifying "aws-sdk-rust" in the `applies_to`
key.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

Author: jasgin

rust-runtime/aws-smithy-http-server-metrics/src/default.rs

@@ -9,7 +9,6 @@
 //! to collect standard metrics automatically.
 
 use std::fmt::Debug;
-use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::time::Duration;
@@ -18,12 +17,21 @@ use metrique::unit_of_work::metrics;
 use metrique::Slot;
 use metrique::SlotGuard;
 
+use crate::default::service_counter::ServiceCounter;
+
+pub(crate) mod service_counter;
+
 /// Container for keeping track of state across all requests for the service (all operations) for default metrics
 ///
 /// This type is not intended for direct use. See [`DefaultMetricsPlugin`](crate::plugin::DefaultMetricsPlugin).
+#[derive(Debug, Default, Clone)]
+pub struct DefaultMetricsServiceCounters {
+    pub(crate) outstanding_requests_counter: ServiceCounter,
+}
+
 #[derive(Debug, Default, Clone)]
 pub struct DefaultMetricsServiceState {
-    pub(crate) outstanding_requests_counter: Arc<AtomicUsize>,
+    pub(crate) outstanding_requests: usize,
 }
 
 /// Container for default request and response metrics.

rust-runtime/aws-smithy-http-server-metrics/src/default/service_counter.rs

@@ -0,0 +1,41 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
+
+/// A service-level counter for metrics that need state beyond a single request
+#[derive(Default, Debug, Clone)]
+pub(crate) struct ServiceCounter {
+    inner: Arc<ServiceCounterInner>,
+}
+impl ServiceCounter {
+    /// Increments the global count by 1, returning a guard that
+    /// decrements the count on drop, and the new value
+    pub(crate) fn increment(&self) -> (ServiceCounterGuard, usize) {
+        let count = self.inner.increment();
+        (ServiceCounterGuard(Arc::clone(&self.inner)), count)
+    }
+}
+
+#[derive(Default, Debug, Clone)]
+pub(crate) struct ServiceCounterInner {
+    count: Arc<AtomicUsize>,
+}
+impl ServiceCounterInner {
+    fn increment(&self) -> usize {
+        self.count.fetch_add(1, Ordering::Relaxed) + 1
+    }
+}
+
+/// Guard for [`ServiceCounter`] that decrements its count on drop
+pub(crate) struct ServiceCounterGuard(Arc<ServiceCounterInner>);
+
+impl Drop for ServiceCounterGuard {
+    fn drop(&mut self) {
+        self.0.count.fetch_sub(1, Ordering::Relaxed);
+    }
+}

rust-runtime/aws-smithy-http-server-metrics/src/layer.rs

@@ -12,6 +12,7 @@ use thiserror::Error;
 use tower::Layer;
 
 use crate::default::DefaultMetrics;
+use crate::default::DefaultMetricsServiceCounters;
 use crate::default::DefaultMetricsServiceState;
 use crate::default::DefaultRequestMetricsConfig;
 use crate::default::DefaultResponseMetricsConfig;
@@ -226,7 +227,7 @@ where
             default_metrics_extension_fn: self.default_metrics_extension_fn,
             default_req_metrics_config: self.default_req_metrics_config.clone(),
             default_res_metrics_config: self.default_res_metrics_config.clone(),
-            default_service_state: DefaultMetricsServiceState::default(),
+            default_service_counters: DefaultMetricsServiceCounters::default(),
 
             _entry_sink: PhantomData,
         }

rust-runtime/aws-smithy-http-server-metrics/src/plugin.rs

@@ -5,7 +5,6 @@
 
 use std::future::Future;
 use std::pin::Pin;
-use std::sync::atomic::Ordering;
 use std::task::Context;
 use std::task::Poll;
 use std::time::Duration;
@@ -336,11 +335,8 @@ fn extend_default_request_metrics(
         return DefaultRequestMetrics::default();
     }
 
-    let outstanding_requests = (!config.disable_outstanding_requests).then_some(
-        ext.service_state
-            .outstanding_requests_counter
-            .load(Ordering::Relaxed),
-    );
+    let outstanding_requests =
+        (!config.disable_outstanding_requests).then_some(ext.service_state.outstanding_requests);
 
     DefaultRequestMetrics {
         service: metrics.service.filter(|_| !config.disable_service),

rust-runtime/aws-smithy-http-server-metrics/src/service.rs

@@ -6,13 +6,14 @@
 use std::future::Future;
 use std::marker::PhantomData;
 use std::pin::Pin;
-use std::sync::atomic::Ordering;
 use std::task::Context;
 use std::task::Poll;
 
 use pin_project_lite::pin_project;
 use tower::Service;
 
+use crate::default::service_counter::ServiceCounterGuard;
+use crate::default::DefaultMetricsServiceCounters;
 use crate::default::DefaultMetricsServiceState;
 use crate::default::DefaultRequestMetricsConfig;
 use crate::default::DefaultResponseMetricsConfig;
@@ -38,7 +39,8 @@ pin_project! {
         inner: F,
         metrics: metrique::AppendAndCloseOnDrop<Entry, Sink>,
         response_metrics: Option<Res>,
-        default_service_state: DefaultMetricsServiceState
+        default_service_state: DefaultMetricsServiceCounters,
+        outstanding_requests_counter_guard: ServiceCounterGuard
     }
 }
 
@@ -57,23 +59,13 @@ where
 
         match this.inner.poll(cx) {
             Poll::Ready(Ok(mut res)) => {
-                this.default_service_state
-                    .outstanding_requests_counter
-                    .fetch_sub(1, Ordering::Relaxed);
-
                 if let Some(response_metrics) = this.response_metrics {
                     (response_metrics)(&mut res, this.metrics);
                 }
 
                 Poll::Ready(Ok(res))
             }
-            Poll::Ready(Err(e)) => {
-                this.default_service_state
-                    .outstanding_requests_counter
-                    .fetch_sub(1, Ordering::Relaxed);
-
-                Poll::Ready(Err(e))
-            }
+            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
             Poll::Pending => Poll::Pending,
         }
     }
@@ -101,7 +93,7 @@ where
     ),
     pub(crate) default_req_metrics_config: DefaultRequestMetricsConfig,
     pub(crate) default_res_metrics_config: DefaultResponseMetricsConfig,
-    pub(crate) default_service_state: DefaultMetricsServiceState,
+    pub(crate) default_service_counters: DefaultMetricsServiceCounters,
 
     pub(crate) _entry_sink: PhantomData<Sink>,
 }
@@ -121,7 +113,7 @@ where
             default_metrics_extension_fn: self.default_metrics_extension_fn,
             default_req_metrics_config: self.default_req_metrics_config.clone(),
             default_res_metrics_config: self.default_res_metrics_config.clone(),
-            default_service_state: self.default_service_state.clone(),
+            default_service_counters: self.default_service_counters.clone(),
 
             _entry_sink: PhantomData,
         }
@@ -148,23 +140,30 @@ where
     fn call(&mut self, mut req: HttpRequest) -> Self::Future {
         let mut metrics = (self.init_metrics)(&mut req);
 
-        self.default_service_state
+        // We increment the outstanding requests and get the count at this layer
+        // (typically outer middleware) so we can get this as close to the time
+        // this request entered the system as possible
+        let (outstanding_requests_counter_guard, outstanding_requests) = self
+            .default_service_counters
             .outstanding_requests_counter
-            .fetch_add(1, Ordering::Relaxed);
+            .increment();
 
         (self.default_metrics_extension_fn)(
             &mut req,
             &mut metrics,
             self.default_req_metrics_config.clone(),
             self.default_res_metrics_config.clone(),
-            self.default_service_state.clone(),
+            DefaultMetricsServiceState {
+                outstanding_requests,
+            },
         );
 
         MetricsLayerServiceFuture {
             inner: self.inner.call(req),
             metrics,
             response_metrics: self.response_metrics.clone(),
-            default_service_state: self.default_service_state.clone(),
+            default_service_state: self.default_service_counters.clone(),
+            outstanding_requests_counter_guard,
         }
     }
 }