chore: add rate limiting capability to the loggin.

Author: darklight3it

lambda-runtime/src/layers/api_response.rs

@@ -1,18 +1,17 @@
 use crate::{
-    constants::LAMBDA_RUNTIME_INVOCATION_ID,
-    deserializer,
-    requests::{EventCompletionRequest, IntoRequest},
-    runtime::LambdaInvocation,
-    Diagnostic, EventErrorRequest, IntoFunctionResponse, LambdaEvent,
+    Diagnostic, EventErrorRequest, IntoFunctionResponse, LambdaEvent, constants::LAMBDA_RUNTIME_INVOCATION_ID, deserializer, rate_limiter::RateLimiter, requests::{EventCompletionRequest, IntoRequest}, runtime::LambdaInvocation,
 };
 use futures::{ready, Stream};
 use lambda_runtime_api_client::{body::Body, BoxError};
 use pin_project::pin_project;
 use serde::{Deserialize, Serialize};
-use std::{fmt::Debug, future::Future, marker::PhantomData, pin::Pin, task};
+use std::{fmt::Debug, future::Future, marker::PhantomData, pin::Pin, task, time::Duration};
 use tower::Service;
 use tracing::{error, trace, warn};
 
+
+static MALFORMED_INVOCATION_ID_LIMITER: RateLimiter = RateLimiter::new(Duration::from_secs(60));
+
 /// Tower service that turns the result or an error of a handler function into a Lambda Runtime API
 /// response.
 ///
@@ -133,7 +132,14 @@ where
             Some(value) => match value.to_str() {
                 Ok(value) => Some(value.to_owned()),
                 Err(error) => {
-                    warn!(error = ?error, "Ignoring malformed Lambda runtime invocation ID header");
+
+                    if MALFORMED_INVOCATION_ID_LIMITER.allow() {
+                        warn!(
+                            error = ?error,
+                            rate_limit_interval_ms = MALFORMED_INVOCATION_ID_LIMITER.interval().as_millis(),
+                            "Ignoring malformed Lambda runtime invocation ID header; this warning is rate limited"
+                        );
+                    }
                     None
                 }
             },

lambda-runtime/src/lib.rs

@@ -29,6 +29,8 @@ mod constants;
 pub mod diagnostic;
 pub use diagnostic::Diagnostic;
 
+mod rate_limiter;
+
 mod deserializer;
 /// Tower middleware to be applied to runtime invocations.
 pub mod layers;

lambda-runtime/src/rate_limiter.rs

@@ -0,0 +1,109 @@
+//! Thread-safe, process-local rate limiting for infrequent runtime events.
+
+use std::{
+    sync::Mutex,
+    time::{Duration, Instant},
+};
+
+/// Allows an operation at most once during each configured interval.
+///
+/// A `RateLimiter` is intended to be shared by concurrent runtime tasks. When
+/// stored in a `static`, it is initialized once per Lambda execution environment
+/// and retains its state across warm invocations. A new cold-started environment
+/// receives a new limiter.
+pub(crate) struct RateLimiter {
+    /// Minimum duration between allowed operations.
+    interval: Duration,
+    /// Timestamp of the most recent allowed operation.
+    last_allowed: Mutex<Option<Instant>>,
+}
+
+impl RateLimiter {
+    /// Creates a rate limiter with the specified minimum interval.
+    pub(crate) const fn new(interval: Duration) -> RateLimiter {
+        RateLimiter {
+            interval,
+            last_allowed: Mutex::new(None),
+        }
+    }
+
+    /// Returns the minimum duration between allowed operations.
+    pub(crate) const fn interval(&self) -> Duration {
+        self.interval
+    }
+
+    ///
+    /// The first call is allowed. Subsequent calls are rejected until the
+    /// configured interval has elapsed since the previous allowed call.
+    /// Concurrent callers are serialized while checking and updating the last
+    /// allowed timestamp so only one caller crosses the interval boundary.
+    pub(crate) fn allow(&self) -> bool {
+        let mut last_allowed = match self.last_allowed.lock() {
+            Ok(guard) => guard,
+            Err(poisoned) => {
+                // The limiter state is disposable, so reset it and recover instead of
+                // allowing a poisoned mutex to crash the runtime or suppress future logs.
+                let mut guard = poisoned.into_inner();
+                *guard = None;
+                self.last_allowed.clear_poison();
+                guard
+            }
+        };
+
+        if last_allowed
+            .as_ref()
+            .is_some_and(|value| value.elapsed() < self.interval)
+        {
+            return false;
+        }
+
+        *last_allowed = Some(Instant::now());
+
+        true
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::panic::{catch_unwind, AssertUnwindSafe};
+    use std::thread;
+
+    #[test]
+    fn allows_first_call() {
+        let limiter = RateLimiter::new(Duration::from_secs(60));
+
+        assert!(limiter.allow());
+    }
+
+    #[test]
+    fn rejects_calls_inside_interval() {
+        let limiter = RateLimiter::new(Duration::from_secs(60));
+
+        assert!(limiter.allow());
+        assert!(!limiter.allow());
+    }
+
+    #[test]
+    fn allows_call_after_interval() {
+        let limiter = RateLimiter::new(Duration::from_millis(10));
+
+        assert!(limiter.allow());
+        thread::sleep(Duration::from_millis(15));
+
+        assert!(limiter.allow());
+    }
+
+    #[test]
+    fn recovers_from_poisoned_mutex() {
+        let limiter = RateLimiter::new(Duration::from_secs(60));
+
+        let _ = catch_unwind(AssertUnwindSafe(|| {
+            let _guard = limiter.last_allowed.lock().unwrap();
+            panic!("poison the limiter mutex");
+        }));
+
+        assert!(limiter.allow());
+        assert!(!limiter.allow());
+    }
+}