feat: Add Push Mode

Cargo.toml

@@ -17,6 +17,7 @@ clap = { version = "4.5.20", features = ["derive"] }
 derive_builder = "0.20.2"
 elegant-departure = { version = "0.3.1", features = ["tokio"] }
 figment = { version = "0.10.19", features = ["env", "yaml", "test"] }
+flume = "0.12.0"
 futures = "0.3.31"
 futures-util = "0.3.31"
 hex = "0.4.3"
@@ -26,8 +27,8 @@ http-body-util = "0.1.2"
 libsqlite3-sys = "0.30.1"
 metrics = "0.24.0"
 metrics-exporter-statsd = "0.9.0"
-prost = "0.13"
-prost-types = "0.13.3"
+prost = "0.14"
+prost-types = "0.14"
 rand = "0.8.5"
 rdkafka = { version = "0.37.0", features = ["cmake-build", "ssl"] }
 sentry = { version = "0.41.0", default-features = false, features = [
@@ -41,16 +42,16 @@ sentry = { version = "0.41.0", default-features = false, features = [
     "tracing",
     "logs"
 ] }
-sentry_protos = "0.4.11"
+sentry_protos = "0.8.5"
 serde = "1.0.214"
 serde_yaml = "0.9.34"
 sha2 = "0.10.8"
 sqlx = { version = "0.8.3", features = ["sqlite", "runtime-tokio", "chrono", "postgres"] }
 tokio = { version = "1.43.1", features = ["full"] }
 tokio-stream = { version = "0.1.16", features = ["full"] }
 tokio-util = "0.7.12"
-tonic = "0.13"
-tonic-health = "0.13"
+tonic = "0.14"
+tonic-health = "0.14"
 tower = "0.5.1"
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = [

README.md

@@ -77,7 +77,7 @@ The test suite is composed of unit and integration tests in Rust, and end-to-end
 
 ```bash
 # Run unit/integration tests
-make test
+make unit-test
 
 # Run end-to-end tests
 make integration-test

src/config.rs

@@ -239,6 +239,27 @@ pub struct Config {
 
     /// Enable additional metrics for the sqlite.
     pub enable_sqlite_status_metrics: bool,
+
+    /// Run the taskbroker in push mode (as opposed to pull mode).
+    pub push_mode: bool,
+
+    /// The number of concurrent dispatchers to run.
+    pub fetch_threads: usize,
+
+    /// The number of concurrent pushers each dispatcher should run.
+    pub push_threads: usize,
+
+    /// The size of the push queue.
+    pub push_queue_size: usize,
+
+    /// The worker service endpoint.
+    pub worker_endpoint: String,
+
+    /// The hostname used to construct `callback_url` for task push requests.
+    pub callback_addr: String,
+
+    /// The port used to construct `callback_url` for task push requests.
+    pub callback_port: u32,
 }
 
 impl Default for Config {
@@ -308,6 +329,13 @@ impl Default for Config {
             full_vacuum_on_upkeep: true,
             vacuum_interval_ms: 30000,
             enable_sqlite_status_metrics: true,
+            push_mode: false,
+            fetch_threads: 1,
+            push_threads: 1,
+            push_queue_size: 1,
+            worker_endpoint: "http://127.0.0.1:50052".into(),
+            callback_addr: "0.0.0.0".into(),
+            callback_port: 50051,
         }
     }
 }
@@ -712,4 +740,48 @@ mod tests {
             Ok(())
         });
     }
+
+    #[test]
+    fn test_default_push_callback_fields() {
+        let config = Config::default();
+        assert_eq!(config.callback_addr, "0.0.0.0");
+        assert_eq!(config.callback_port, 50051);
+    }
+
+    #[test]
+    fn test_from_args_push_callback_fields_from_env() {
+        Jail::expect_with(|jail| {
+            jail.set_env("TASKBROKER_CALLBACK_ADDR", "127.0.0.1");
+            jail.set_env("TASKBROKER_CALLBACK_PORT", "51000");
+
+            let args = Args { config: None };
+            let config = Config::from_args(&args).unwrap();
+            assert_eq!(config.callback_addr, "127.0.0.1");
+            assert_eq!(config.callback_port, 51000);
+
+            Ok(())
+        });
+    }
+
+    #[test]
+    fn test_from_args_push_callback_fields_from_config_file() {
+        Jail::expect_with(|jail| {
+            jail.create_file(
+                "config.yaml",
+                r#"
+                callback_addr: 10.0.0.1
+                callback_port: 52000
+            "#,
+            )?;
+
+            let args = Args {
+                config: Some("config.yaml".to_owned()),
+            };
+            let config = Config::from_args(&args).unwrap();
+            assert_eq!(config.callback_addr, "10.0.0.1");
+            assert_eq!(config.callback_port, 52000);
+
+            Ok(())
+        });
+    }
 }