WIP: support running different policies

Initial iteration, everything seems to be working, but there's still
some plumbing to do (settings, cli flags,...)

Author: flavio

src/main.rs

@@ -3,15 +3,17 @@ use hyper::service::{make_service_fn, service_fn};
 use hyper::Server;
 use std::{net::SocketAddr, thread};
 use std::process;
-use tokio::{runtime::Runtime, sync::mpsc};
+use tokio::{runtime::Runtime, sync::mpsc::channel};
 
 mod admission_review;
 mod api;
 mod wasm;
 mod wasm_fetcher;
+mod worker;
 
-#[tokio::main]
-async fn main() {
+use crate::wasm::EvalRequest;
+
+fn main() {
     let matches = App::new("policy-server")
         .version("0.0.1")
         .about("Kubernetes admission controller powered by Chimera WASM policies")
@@ -84,6 +86,13 @@ async fn main() {
         }
     };
 
+    let rt = match Runtime::new() {
+        Ok(r) => { r },
+        Err(error) => {
+            return fatal_error(format!("Error initializing tokio runtime: {}", error));
+        }
+    };
+
     let fetcher = match wasm_fetcher::parse_wasm_url(
         matches.value_of("wasm-uri").unwrap(),
         matches.is_present("wasm-remote-insecure"),
@@ -94,42 +103,37 @@ async fn main() {
             return fatal_error(format!("Error parsing arguments: {}", error));
         }
     };
-    let wasm_path = match fetcher.fetch().await {
-        Ok(p) => { p },
-        Err(error) => {
-            return fatal_error(format!("Error fetching WASM module: {}", error));
-        }
+    let wasm_path = match rt.block_on(async { fetcher.fetch().await }) {
+        Ok(p) =>p,
+        Err(error) => { return fatal_error(format!("Error fetching WASM module: {}", error));}
     };
 
-    let (tx, mut rx) = mpsc::channel::<wasm::EvalRequest>(32);
+    let (api_tx, api_rx) = channel::<EvalRequest>(32);
+
+    let mut wasm_modules = Vec::<String>::new();
+    wasm_modules.push(wasm_path);
 
-    let rt = Runtime::new().unwrap();
     let wasm_thread = thread::spawn(move || {
-        let mut policy_evaluator = match wasm::PolicyEvaluator::new(&wasm_path) {
-            Ok(e) => { e },
-            Err(error) => {
-                return fatal_error(format!("Error initializing policy evaluator for {}: {}", wasm_path, error));
-            }
-        };
-        rt.block_on(async move {
-            while let Some(req) = rx.recv().await {
-                let resp = policy_evaluator.validate(req.req);
-                let _ = req.resp_chan.send(resp);
-            }
-        });
-    });
+        let worker_pool = worker::WorkerPool::new(3, wasm_modules.clone(), api_rx).unwrap();
 
-    let make_svc = make_service_fn(|_conn| {
-        let svc_tx = tx.clone();
-        async move { Ok::<_, hyper::Error>(service_fn(move |req| api::route(req, svc_tx.clone()))) }
+        worker_pool.run();
     });
 
-    let server = Server::bind(&addr).serve(make_svc);
-    println!("Started server on {}", addr);
+    rt.block_on( async {
+        let make_svc = make_service_fn(|_conn| {
+            let svc_tx = api_tx.clone();
+            async move { 
+                Ok::<_, hyper::Error>(service_fn(move |req| api::route(req, svc_tx.clone()))) }
+        });
+
+        let server = Server::bind(&addr).serve(make_svc);
+        println!("Started server on {}", addr);
+
+        if let Err(e) = server.await {
+            eprintln!("server error: {}", e);
+        }
+    });
 
-    if let Err(e) = server.await {
-        eprintln!("server error: {}", e);
-    }
     wasm_thread.join().unwrap();
 }
 

src/worker.rs

@@ -0,0 +1,90 @@
+use crate::wasm::{EvalRequest, PolicyEvaluator};
+use anyhow::Result;
+use std::{collections::HashMap, thread, vec::Vec};
+use tokio::sync::mpsc::{channel, Receiver, Sender};
+
+pub(crate) struct Worker {
+  evaluators: HashMap<String, PolicyEvaluator>,
+  channel_rx: Receiver<EvalRequest>,
+}
+
+impl Worker {
+  pub(crate) fn new(rx: Receiver<EvalRequest>, wasm_modules: Vec<String>) -> Result<Worker> {
+    let mut evs: HashMap<String, PolicyEvaluator> = HashMap::new();
+
+    for w in wasm_modules {
+      let policy_evaluator = PolicyEvaluator::new(&w)?;
+      //TODO: no hard coded value
+      evs.insert("validate".into(), policy_evaluator);
+    }
+
+    Ok(Worker {
+      evaluators: evs,
+      channel_rx: rx,
+    })
+  }
+
+  pub(crate) fn run(mut self) {
+    while let Some(req) = self.channel_rx.blocking_recv() {
+      //TODO: handle error
+      let policy_id = "validate".to_string();
+      match self.evaluators.get_mut(&policy_id) {
+        Some(policy_evaluator) => {
+          let resp = policy_evaluator.validate(req.req);
+          let _ = req.resp_chan.send(resp);
+        }
+        None => continue,
+      }
+    }
+  }
+}
+
+pub(crate) struct WorkerPool {
+  pool_size: usize,
+  worker_tx_chans: Vec<Sender<EvalRequest>>,
+  api_rx: Receiver<EvalRequest>,
+}
+
+impl WorkerPool {
+  pub(crate) fn new(
+    size: usize,
+    wasm_modules: Vec<String>,
+    rx: Receiver<EvalRequest>,
+  ) -> Result<WorkerPool> {
+    let mut tx_chans = Vec::<Sender<EvalRequest>>::new();
+
+    for n in 1..=size {
+      let (tx, rx) = channel::<EvalRequest>(32);
+      tx_chans.push(tx);
+      let wasm_modules = wasm_modules.clone();
+
+      thread::spawn(move || {
+        let worker = Worker::new(rx, wasm_modules).unwrap();
+
+        //TODO: better logging
+        println!("worker {} loop start", n);
+        worker.run();
+        println!("worker {} loop exit", n);
+      });
+    }
+
+    Ok(WorkerPool {
+      pool_size: size,
+      worker_tx_chans: tx_chans,
+      api_rx: rx,
+    })
+  }
+
+  pub(crate) fn run(mut self) {
+    let mut next_worker_id = 0;
+    while let Some(req) = self.api_rx.blocking_recv() {
+      let _ = self.worker_tx_chans[next_worker_id].blocking_send(req);
+      next_worker_id += 1;
+      if next_worker_id >= self.pool_size {
+        next_worker_id = 0;
+      }
+    }
+
+    //TODO: should we also `join` the children threads here?
+  }
+}