enhancement(operations): abstracts runtime into runtime.rs (#1098)

* Initial runtime::Runtime wrapper

* Replaces tokio::Runtime with wrapper

#1085

* Add TaskExecutor

Signed-off-by: Lucio Franco <luciofranco14@gmail.com>

Author: Michael-J-Ward

benches/bench.rs

@@ -10,7 +10,7 @@ use vector::test_util::{
 };
 use vector::topology::config::TransformConfig;
 use vector::topology::{self, config};
-use vector::{sinks, sources, transforms};
+use vector::{runtime, sinks, sources, transforms};
 
 mod batch;
 mod buffering;
@@ -61,7 +61,7 @@ fn benchmark_simple_pipe(c: &mut Criterion) {
                         sinks::tcp::TcpSinkConfig::new(out_addr.to_string()),
                     );
 
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines = count_receive(&out_addr);
 
@@ -107,7 +107,7 @@ fn benchmark_simple_pipe_with_tiny_lines(c: &mut Criterion) {
                         sinks::tcp::TcpSinkConfig::new(out_addr.to_string()),
                     );
 
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines = count_receive(&out_addr);
 
@@ -153,7 +153,7 @@ fn benchmark_simple_pipe_with_huge_lines(c: &mut Criterion) {
                         sinks::tcp::TcpSinkConfig::new(out_addr.to_string()),
                     );
 
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines = count_receive(&out_addr);
 
@@ -200,7 +200,7 @@ fn benchmark_simple_pipe_with_many_writers(c: &mut Criterion) {
                         sinks::tcp::TcpSinkConfig::new(out_addr.to_string()),
                     );
 
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines = count_receive(&out_addr);
 
@@ -264,7 +264,7 @@ fn benchmark_interconnected(c: &mut Criterion) {
                         sinks::tcp::TcpSinkConfig::new(out_addr2.to_string()),
                     );
 
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines1 = count_receive(&out_addr1);
                     let output_lines2 = count_receive(&out_addr2);
@@ -331,7 +331,7 @@ fn benchmark_transforms(c: &mut Criterion) {
                         &["filter"],
                         sinks::tcp::TcpSinkConfig::new(out_addr.to_string()),
                     );
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines = count_receive(&out_addr);
 
@@ -485,7 +485,7 @@ fn benchmark_complex(c: &mut Criterion) {
                         &["filter_500"],
                         sinks::tcp::TcpSinkConfig::new(out_addr_500.to_string()),
                     );
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines_all = count_receive(&out_addr_all);
                     let output_lines_sampled = count_receive(&out_addr_sampled);

benches/buffering.rs

@@ -5,7 +5,7 @@ use vector::test_util::{
     block_on, count_receive, next_addr, send_lines, shutdown_on_idle, wait_for_tcp,
 };
 use vector::topology::{self, config};
-use vector::{buffers::BufferConfig, sinks, sources};
+use vector::{buffers::BufferConfig, runtime, sinks, sources};
 
 fn benchmark_buffers(c: &mut Criterion) {
     let num_lines: usize = 100_000;
@@ -35,7 +35,7 @@ fn benchmark_buffers(c: &mut Criterion) {
                         when_full: Default::default(),
                     };
 
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines = count_receive(&out_addr);
 
@@ -72,7 +72,7 @@ fn benchmark_buffers(c: &mut Criterion) {
                     .into();
                     config.global.data_dir = Some(data_dir.clone());
 
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines = count_receive(&out_addr);
 
@@ -108,7 +108,7 @@ fn benchmark_buffers(c: &mut Criterion) {
                     };
                     config.global.data_dir = Some(data_dir2.clone());
 
-                    let mut rt = tokio::runtime::Runtime::new().unwrap();
+                    let mut rt = runtime::Runtime::new().unwrap();
 
                     let output_lines = count_receive(&out_addr);
 

benches/files.rs

@@ -7,7 +7,7 @@ use tokio::codec::{BytesCodec, FramedWrite};
 use tokio::fs::OpenOptions;
 use vector::test_util::random_lines;
 use vector::{
-    sinks, sources,
+    runtime, sinks, sources,
     topology::{self, config},
 };
 
@@ -54,7 +54,7 @@ fn benchmark_files_without_partitions(c: &mut Criterion) {
                     },
                 );
 
-                let mut rt = tokio::runtime::Runtime::new().unwrap();
+                let mut rt = runtime::Runtime::new().unwrap();
                 let (topology, _crash) = topology::start(config, &mut rt, false).unwrap();
 
                 let mut options = OpenOptions::new();

benches/http.rs

@@ -5,7 +5,7 @@ use hyper::{Body, Response, Server};
 use std::net::SocketAddr;
 use vector::test_util::{next_addr, random_lines, send_lines, wait_for_tcp};
 use vector::{
-    sinks, sources,
+    runtime, sinks, sources,
     topology::{self, config},
 };
 
@@ -33,7 +33,7 @@ fn benchmark_http_no_compression(c: &mut Criterion) {
                     },
                 );
 
-                let mut rt = tokio::runtime::Runtime::new().unwrap();
+                let mut rt = runtime::Runtime::new().unwrap();
 
                 let (topology, _crash) = topology::start(config, &mut rt, false).unwrap();
                 wait_for_tcp(in_addr);
@@ -80,7 +80,7 @@ fn benchmark_http_gzip(c: &mut Criterion) {
                     },
                 );
 
-                let mut rt = tokio::runtime::Runtime::new().unwrap();
+                let mut rt = runtime::Runtime::new().unwrap();
 
                 let (topology, _crash) = topology::start(config, &mut rt, false).unwrap();
                 wait_for_tcp(in_addr);

lib/tracing-limit/benches/limit.rs

@@ -41,7 +41,7 @@ fn bench(c: &mut Criterion) {
     c.bench_function_over_inputs(
         "Limit 5 seconds",
         |b, n| {
-            let sub = VisitingSubscriber(Mutex::new(String::from(""))).with(Limit::default());;
+            let sub = VisitingSubscriber(Mutex::new(String::from(""))).with(Limit::default());
             let n = black_box(n);
             tracing::subscriber::with_default(sub, || {
                 b.iter(|| {

src/lib.rs

@@ -15,6 +15,7 @@ pub mod buffers;
 pub mod event;
 pub mod metrics;
 pub mod region;
+pub mod runtime;
 pub mod sinks;
 pub mod sources;
 pub mod template;

src/main.rs

@@ -12,7 +12,7 @@ use structopt::StructOpt;
 use tokio_signal::unix::{Signal, SIGHUP, SIGINT, SIGQUIT, SIGTERM};
 use topology::Config;
 use tracing_futures::Instrument;
-use vector::{metrics, topology, trace};
+use vector::{metrics, runtime, topology, trace};
 
 #[derive(StructOpt, Debug)]
 #[structopt(rename_all = "kebab-case")]
@@ -206,12 +206,9 @@ fn main() {
     });
 
     let mut rt = {
-        let mut builder = tokio::runtime::Builder::new();
-
         let threads = opts.threads.unwrap_or(max(1, num_cpus::get()));
-        builder.core_threads(min(4, threads));
-
-        builder.build().expect("Unable to create async runtime")
+        let num_threads = min(4, threads);
+        runtime::Runtime::with_thread_count(num_threads).expect("Unable to create async runtime")
     };
 
     let (metrics_trigger, metrics_tripwire) = stream_cancel::Tripwire::new();

src/runtime.rs

@@ -0,0 +1,70 @@
+use futures::future::{ExecuteError, Executor, Future};
+use std::io;
+use tokio::runtime::Builder;
+
+pub struct Runtime {
+    rt: tokio::runtime::Runtime,
+}
+
+impl Runtime {
+    pub fn new() -> io::Result<Self> {
+        Ok(Runtime {
+            rt: tokio::runtime::Runtime::new()?,
+        })
+    }
+
+    pub fn single_threaded() -> io::Result<Self> {
+        Self::with_thread_count(1)
+    }
+
+    pub fn with_thread_count(number: usize) -> io::Result<Self> {
+        Ok(Runtime {
+            rt: Builder::new().core_threads(number).build()?,
+        })
+    }
+
+    pub fn spawn<F>(&mut self, future: F) -> &mut Self
+    where
+        F: Future<Item = (), Error = ()> + Send + 'static,
+    {
+        self.rt.spawn(future);
+        self
+    }
+
+    pub fn executor(&self) -> TaskExecutor {
+        TaskExecutor {
+            inner: self.rt.executor(),
+        }
+    }
+
+    pub fn block_on<F, R, E>(&mut self, future: F) -> Result<R, E>
+    where
+        F: Send + 'static + Future<Item = R, Error = E>,
+        R: Send + 'static,
+        E: Send + 'static,
+    {
+        self.rt.block_on(future)
+    }
+
+    pub fn shutdown_on_idle(self) -> impl Future<Item = (), Error = ()> {
+        self.rt.shutdown_on_idle()
+    }
+
+    pub fn shutdown_now(self) -> impl Future<Item = (), Error = ()> {
+        self.rt.shutdown_now()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct TaskExecutor {
+    inner: tokio::runtime::TaskExecutor,
+}
+
+impl<F> Executor<F> for TaskExecutor
+where
+    F: Future<Item = (), Error = ()> + Send + 'static,
+{
+    fn execute(&self, future: F) -> Result<(), ExecuteError<F>> {
+        self.inner.execute(future)
+    }
+}

src/sinks/aws_kinesis_streams.rs

@@ -330,13 +330,13 @@ mod integration_tests {
     use crate::{
         buffers::Acker,
         region::RegionOrEndpoint,
+        runtime,
         test_util::{random_lines_with_stream, random_string},
     };
     use futures::{Future, Sink};
     use rusoto_core::Region;
     use rusoto_kinesis::{Kinesis, KinesisClient};
     use std::sync::Arc;
-    use tokio::runtime::Runtime;
 
     #[test]
     fn kinesis_put_records() {
@@ -356,7 +356,7 @@ mod integration_tests {
             ..Default::default()
         };
 
-        let mut rt = Runtime::new().unwrap();
+        let mut rt = runtime::Runtime::new().unwrap();
 
         let sink = KinesisService::new(config, Acker::Null).unwrap();
 

src/sinks/aws_s3.rs

@@ -400,6 +400,7 @@ mod integration_tests {
         assert_downcast_matches,
         event::Event,
         region::RegionOrEndpoint,
+        runtime::Runtime,
         sinks::aws_s3::{S3Sink, S3SinkConfig},
         test_util::{block_on, random_lines_with_stream, random_string},
     };
@@ -502,7 +503,7 @@ mod integration_tests {
         let (tx, rx) = futures::sync::mpsc::channel(1);
         let pump = sink.send_all(rx).map(|_| ()).map_err(|_| ());
 
-        let mut rt = tokio::runtime::Runtime::new().unwrap();
+        let mut rt = Runtime::new().unwrap();
         rt.spawn(pump);
 
         let mut tx = tx.wait();
@@ -588,15 +589,15 @@ mod integration_tests {
 
     #[test]
     fn s3_healthchecks() {
-        let mut rt = tokio::runtime::Runtime::new().unwrap();
+        let mut rt = Runtime::new().unwrap();
 
         let healthcheck = S3Sink::healthcheck(&config()).unwrap();
         rt.block_on(healthcheck).unwrap();
     }
 
     #[test]
     fn s3_healthchecks_invalid_bucket() {
-        let mut rt = tokio::runtime::Runtime::new().unwrap();
+        let mut rt = Runtime::new().unwrap();
 
         let config = S3SinkConfig {
             bucket: "asdflkjadskdaadsfadf".to_string(),

src/sinks/http.rs

@@ -293,6 +293,7 @@ mod tests {
     use crate::buffers::Acker;
     use crate::{
         assert_downcast_matches,
+        runtime::Runtime,
         sinks::http::HttpSinkConfig,
         test_util::{next_addr, random_lines_with_stream, shutdown_on_idle},
         topology::config::SinkConfig,
@@ -387,7 +388,7 @@ mod tests {
         let (input_lines, events) = random_lines_with_stream(100, num_lines);
         let pump = sink.send_all(events);
 
-        let mut rt = tokio::runtime::Runtime::new().unwrap();
+        let mut rt = Runtime::new().unwrap();
         rt.spawn(server);
 
         rt.block_on(pump).unwrap();
@@ -445,7 +446,7 @@ mod tests {
         let (input_lines, events) = random_lines_with_stream(100, num_lines);
         let pump = sink.send_all(events);
 
-        let mut rt = tokio::runtime::Runtime::new().unwrap();
+        let mut rt = Runtime::new().unwrap();
         rt.spawn(server);
 
         rt.block_on(pump).unwrap();
@@ -503,7 +504,7 @@ mod tests {
         let (input_lines, events) = random_lines_with_stream(100, num_lines);
         let pump = sink.send_all(events);
 
-        let mut rt = tokio::runtime::Runtime::new().unwrap();
+        let mut rt = Runtime::new().unwrap();
         rt.spawn(server);
 
         rt.block_on(pump).unwrap();

src/sinks/util/http.rs

@@ -224,7 +224,7 @@ mod test {
             .serve(new_service)
             .map_err(|e| eprintln!("server error: {}", e));
 
-        let mut rt = tokio::runtime::Runtime::new().unwrap();
+        let mut rt = crate::runtime::Runtime::new().unwrap();
 
         rt.spawn(server);
 

src/sinks/util/mod.rs

@@ -219,10 +219,10 @@ where
 mod test {
     use super::BatchServiceSink;
     use crate::buffers::Acker;
+    use crate::runtime::Runtime;
     use crate::test_util::wait_for;
     use futures::{stream, sync::oneshot, Future, Poll, Sink};
     use std::sync::{atomic::Ordering, Arc, Mutex};
-    use tokio::runtime::Runtime;
     use tower::Service;
 
     struct FakeService {