enhancement(sources): add config option to limit source tcp connectio…

…ns (#10491)

benches/topology.rs

@@ -466,6 +466,7 @@ fn benchmark_real_world_1(c: &mut Criterion) {
                         keepalive: None,
                         tls: None,
                         receive_buffer_bytes: None,
+                        connection_limit: None
                     }),
                 );
 

src/sources/fluent/mod.rs

@@ -34,6 +34,7 @@ pub struct FluentConfig {
     receive_buffer_bytes: Option<usize>,
     #[serde(default, deserialize_with = "bool_or_struct")]
     acknowledgements: AcknowledgementsConfig,
+    connection_limit: Option<u32>,
 }
 
 inventory::submit! {
@@ -48,6 +49,7 @@ impl GenerateConfig for FluentConfig {
             tls: None,
             receive_buffer_bytes: None,
             acknowledgements: Default::default(),
+            connection_limit: Some(2),
         })
         .unwrap()
     }
@@ -68,6 +70,7 @@ impl SourceConfig for FluentConfig {
             self.receive_buffer_bytes,
             cx,
             self.acknowledgements,
+            self.connection_limit,
         )
     }
 
@@ -306,20 +309,14 @@ impl Decoder for FluentDecoder {
             let res = Deserialize::deserialize(&mut des).map_err(DecodeError::Decode);
 
             // check for unexpected EOF to indicate that we need more data
-            match res {
-                // can use or-patterns in 1.53
-                // https://github.com/rust-lang/rust/pull/79278
-                Err(DecodeError::Decode(decode::Error::InvalidDataRead(ref custom))) => {
-                    if custom.kind() == io::ErrorKind::UnexpectedEof {
-                        return Ok(None);
-                    }
-                }
-                Err(DecodeError::Decode(decode::Error::InvalidMarkerRead(ref custom))) => {
-                    if custom.kind() == io::ErrorKind::UnexpectedEof {
-                        return Ok(None);
-                    }
+            if let Err(DecodeError::Decode(
+                decode::Error::InvalidDataRead(ref custom)
+                | decode::Error::InvalidMarkerRead(ref custom),
+            )) = res
+            {
+                if custom.kind() == io::ErrorKind::UnexpectedEof {
+                    return Ok(None);
                 }
-                _ => {}
             }
 
             (des.position() as usize, res)
@@ -809,6 +806,7 @@ mod tests {
             keepalive: None,
             receive_buffer_bytes: None,
             acknowledgements: true.into(),
+            connection_limit: None,
         }
         .build(SourceContext::new_test(sender))
         .await
@@ -1053,6 +1051,7 @@ mod integration_tests {
                 keepalive: None,
                 receive_buffer_bytes: None,
                 acknowledgements: false.into(),
+                connection_limit: None,
             }
             .build(SourceContext::new_test(sender))
             .await

src/sources/logstash.rs

@@ -32,6 +32,7 @@ pub struct LogstashConfig {
     receive_buffer_bytes: Option<usize>,
     #[serde(default, deserialize_with = "bool_or_struct")]
     acknowledgements: AcknowledgementsConfig,
+    connection_limit: Option<u32>,
 }
 
 inventory::submit! {
@@ -46,6 +47,7 @@ impl GenerateConfig for LogstashConfig {
             tls: None,
             receive_buffer_bytes: None,
             acknowledgements: Default::default(),
+            connection_limit: None,
         })
         .unwrap()
     }
@@ -68,6 +70,7 @@ impl SourceConfig for LogstashConfig {
             self.receive_buffer_bytes,
             cx,
             self.acknowledgements,
+            self.connection_limit,
         )
     }
 
@@ -591,6 +594,7 @@ mod test {
             keepalive: None,
             receive_buffer_bytes: None,
             acknowledgements: true.into(),
+            connection_limit: None,
         }
         .build(SourceContext::new_test(sender))
         .await
@@ -800,6 +804,7 @@ output {
                 keepalive: None,
                 receive_buffer_bytes: None,
                 acknowledgements: false.into(),
+                connection_limit: None,
             }
             .build(SourceContext::new_test(sender))
             .await

src/sources/socket/mod.rs

@@ -112,6 +112,7 @@ impl SourceConfig for SocketConfig {
                     config.receive_buffer_bytes(),
                     cx,
                     false.into(),
+                    config.connection_limit,
                 )
             }
             Mode::Udp(config) => {

src/sources/socket/tcp.rs

@@ -40,6 +40,7 @@ pub struct TcpConfig {
     #[serde(default = "default_decoding")]
     #[getset(get = "pub", set = "pub")]
     decoding: Box<dyn DeserializerConfig>,
+    pub connection_limit: Option<u32>,
 }
 
 const fn default_shutdown_timeout_secs() -> u64 {
@@ -57,6 +58,7 @@ impl TcpConfig {
         receive_buffer_bytes: Option<usize>,
         framing: Option<Box<dyn FramingConfig>>,
         decoding: Box<dyn DeserializerConfig>,
+        connection_limit: Option<u32>,
     ) -> Self {
         Self {
             address,
@@ -68,6 +70,7 @@ impl TcpConfig {
             receive_buffer_bytes,
             framing,
             decoding,
+            connection_limit,
         }
     }
 
@@ -82,6 +85,7 @@ impl TcpConfig {
             receive_buffer_bytes: None,
             framing: None,
             decoding: default_decoding(),
+            connection_limit: None,
         }
     }
 }

src/sources/statsd/mod.rs

@@ -61,6 +61,7 @@ struct TcpConfig {
     #[serde(default = "default_shutdown_timeout_secs")]
     shutdown_timeout_secs: u64,
     receive_buffer_bytes: Option<usize>,
+    connection_limit: Option<u32>,
 }
 
 impl TcpConfig {
@@ -73,6 +74,7 @@ impl TcpConfig {
             tls: None,
             shutdown_timeout_secs: default_shutdown_timeout_secs(),
             receive_buffer_bytes: None,
+            connection_limit: None,
         }
     }
 }
@@ -112,6 +114,7 @@ impl SourceConfig for StatsdConfig {
                     config.receive_buffer_bytes,
                     cx,
                     false.into(),
+                    config.connection_limit,
                 )
             }
             #[cfg(unix)]

src/sources/syslog.rs

@@ -49,6 +49,7 @@ pub enum Mode {
         keepalive: Option<TcpKeepaliveConfig>,
         tls: Option<TlsConfig>,
         receive_buffer_bytes: Option<usize>,
+        connection_limit: Option<u32>,
     },
     Udp {
         address: SocketAddr,
@@ -80,6 +81,7 @@ impl GenerateConfig for SyslogConfig {
                 keepalive: None,
                 tls: None,
                 receive_buffer_bytes: None,
+                connection_limit: None,
             },
             host_key: None,
             max_length: crate::serde::default_max_length(),
@@ -103,6 +105,7 @@ impl SourceConfig for SyslogConfig {
                 keepalive,
                 tls,
                 receive_buffer_bytes,
+                connection_limit,
             } => {
                 let source = SyslogTcpSource {
                     max_length: self.max_length,
@@ -118,6 +121,7 @@ impl SourceConfig for SyslogConfig {
                     receive_buffer_bytes,
                     cx,
                     false.into(),
+                    connection_limit,
                 )
             }
             Mode::Udp {

src/sources/util/tcp.rs

@@ -1,17 +1,11 @@
-use std::{
-    fmt, io,
-    mem::drop,
-    net::{IpAddr, SocketAddr},
-    sync::Arc,
-    time::Duration,
-};
-
 use bytes::Bytes;
 use futures::{future::BoxFuture, stream, FutureExt, Sink, SinkExt, StreamExt};
 use listenfd::ListenFd;
 use serde::{de, Deserialize, Deserializer, Serialize};
 use smallvec::SmallVec;
 use socket2::SockRef;
+use std::net::{IpAddr, SocketAddr};
+use std::{fmt, io, mem::drop, sync::Arc, time::Duration};
 use tokio::{
     io::AsyncWriteExt,
     net::{TcpListener, TcpStream},
@@ -118,6 +112,7 @@ where
         receive_buffer_bytes: Option<usize>,
         cx: SourceContext,
         acknowledgements: AcknowledgementsConfig,
+        max_connections: Option<u32>,
     ) -> crate::Result<crate::sources::Source> {
         let out = cx
             .out
@@ -150,9 +145,9 @@ where
             let shutdown_clone = cx.shutdown.clone();
 
             listener
-                .accept_stream()
+                .accept_stream_limited(max_connections)
                 .take_until(shutdown_clone)
-                .for_each(move |connection| {
+                .for_each(move |(connection, permit)| {
                     let shutdown_signal = cx.shutdown.clone();
                     let tripwire = tripwire.clone();
                     let source = self.clone();
@@ -202,7 +197,11 @@ where
                             );
 
                             tokio::spawn(
-                                fut.map(move |()| drop(open_token)).instrument(span.clone()),
+                                fut.map(move |()| {
+                                    drop(open_token);
+                                    drop(permit);
+                                })
+                                .instrument(span.clone()),
                             );
                         });
                     }

src/sources/vector/v1.rs

@@ -66,6 +66,7 @@ impl VectorConfig {
             self.receive_buffer_bytes,
             cx,
             false.into(),
+            None,
         )
     }
 

src/tls/incoming.rs

@@ -1,3 +1,4 @@
+use std::sync::Arc;
 use std::{
     future::Future,
     net::SocketAddr,
@@ -8,6 +9,7 @@ use std::{
 use futures::{future::BoxFuture, stream, FutureExt, Stream};
 use openssl::ssl::{Ssl, SslAcceptor, SslMethod};
 use snafu::ResultExt;
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use tokio::{
     io::{self, AsyncRead, AsyncWrite, ReadBuf},
     net::{TcpListener, TcpStream},
@@ -72,6 +74,7 @@ impl MaybeTlsListener {
         (self.accept().await, self)
     }
 
+    #[allow(unused)]
     pub(crate) fn accept_stream(
         self,
     ) -> impl Stream<Item = crate::tls::Result<MaybeTlsIncomingStream<TcpStream>>> {
@@ -85,6 +88,44 @@ impl MaybeTlsListener {
         })
     }
 
+    #[allow(unused)]
+    pub(crate) fn accept_stream_limited(
+        self,
+        max_connections: Option<u32>,
+    ) -> impl Stream<
+        Item = (
+            crate::tls::Result<MaybeTlsIncomingStream<TcpStream>>,
+            Option<OwnedSemaphorePermit>,
+        ),
+    > {
+        let connection_semaphore =
+            max_connections.map(|max| Arc::new(Semaphore::new(max as usize)));
+
+        let mut semaphore_future = connection_semaphore
+            .clone()
+            .map(|x| Box::pin(x.acquire_owned()));
+        let mut accept = Box::pin(self.into_accept());
+        stream::poll_fn(move |context| {
+            let permit = match semaphore_future.as_mut() {
+                Some(semaphore) => match semaphore.as_mut().poll(context) {
+                    Poll::Ready(permit) => {
+                        semaphore.set(connection_semaphore.clone().unwrap().acquire_owned());
+                        permit.ok()
+                    }
+                    Poll::Pending => return Poll::Pending,
+                },
+                None => None,
+            };
+            match accept.as_mut().poll(context) {
+                Poll::Ready((item, this)) => {
+                    accept.set(this.into_accept());
+                    Poll::Ready(Some((item, permit)))
+                }
+                Poll::Pending => Poll::Pending,
+            }
+        })
+    }
+
     #[cfg(feature = "listenfd")]
     pub(crate) fn local_addr(&self) -> Result<SocketAddr, std::io::Error> {
         self.listener.local_addr()

tests/syslog.rs

@@ -37,6 +37,7 @@ async fn test_tcp_syslog() {
             keepalive: None,
             tls: None,
             receive_buffer_bytes: None,
+            connection_limit: None,
         }),
     );
     config.add_sink("out", &["in"], tcp_json_sink(out_addr.to_string()));
@@ -155,6 +156,7 @@ async fn test_octet_counting_syslog() {
             keepalive: None,
             tls: None,
             receive_buffer_bytes: None,
+            connection_limit: None,
         }),
     );
     config.add_sink("out", &["in"], tcp_json_sink(out_addr.to_string()));

website/cue/reference/components/sources/fluent.cue

@@ -58,6 +58,16 @@ components: sources: fluent: {
 				examples: ["0.0.0.0:\(_port)"]
 			}
 		}
+		connection_limit: {
+			common:        false
+			description:   "The max number of TCP connections that will be processed."
+			relevant_when: "mode = `tcp`"
+			required:      false
+			type: uint: {
+				default: null
+				unit:    "concurrency"
+			}
+		}
 	}
 
 	output: logs: line: {

website/cue/reference/components/sources/logstash.cue

@@ -58,6 +58,16 @@ components: sources: logstash: {
 				examples: ["0.0.0.0:\(_port)"]
 			}
 		}
+		connection_limit: {
+			common:        false
+			description:   "The max number of TCP connections that will be processed."
+			relevant_when: "mode = `tcp`"
+			required:      false
+			type: uint: {
+				default: null
+				unit:    "concurrency"
+			}
+		}
 	}
 
 	output: logs: line: {