diff --git a/plane/.sqlx/query-cc4bc25932348053c5772632c5be3b9ee1e6a73798340d11950486cc16a6c81c.json b/plane/.sqlx/query-909f3471f8b5f39821d24eeffacbf02ec13a898272581437e00fc6d967def921.json similarity index 92% rename from plane/.sqlx/query-cc4bc25932348053c5772632c5be3b9ee1e6a73798340d11950486cc16a6c81c.json rename to plane/.sqlx/query-909f3471f8b5f39821d24eeffacbf02ec13a898272581437e00fc6d967def921.json index 546ad97e1..c8ed5059c 100644 --- a/plane/.sqlx/query-cc4bc25932348053c5772632c5be3b9ee1e6a73798340d11950486cc16a6c81c.json +++ b/plane/.sqlx/query-909f3471f8b5f39821d24eeffacbf02ec13a898272581437e00fc6d967def921.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n select\n drone.id,\n node.name,\n drone.last_local_time as \"last_local_time!\",\n (\n select\n count(*)\n from backend\n where drone_id = node.id\n and last_status != $4\n ) as \"backend_count\"\n from node\n left join drone\n on node.id = drone.id\n left join controller\n on node.controller = controller.id\n where\n drone.ready = true\n and controller is not null\n and cluster = $1\n and now() - drone.last_heartbeat < $2\n and now() - controller.last_heartbeat < $2\n and controller.is_online = true\n and draining = false\n and last_local_time is not null\n and pool = $3\n order by random()\n limit 2\n ", + "query": "\n select\n drone.id,\n node.name,\n drone.last_local_time as \"last_local_time!\",\n (\n select\n count(*)\n from backend\n where drone_id = node.id\n and last_status != $4\n ) as \"backend_count\"\n from node\n left join drone\n on node.id = drone.id\n left join controller\n on node.controller = controller.id\n where\n drone.ready = true\n and controller is not null\n and cluster = $1\n and now() - drone.last_heartbeat < $2\n and now() - controller.last_heartbeat < $2\n and controller.is_online = true\n and draining = false\n and last_local_time is not null\n and pool = $3\n order by random()\n limit $5\n ", "describe": { "columns": [ { @@ -29,7 +29,8 @@ "Text", "Interval", "Text", - "Text" + "Text", + "Int8" ] }, "nullable": [ @@ -39,5 +40,5 @@ null ] }, - "hash": "cc4bc25932348053c5772632c5be3b9ee1e6a73798340d11950486cc16a6c81c" + "hash": "909f3471f8b5f39821d24eeffacbf02ec13a898272581437e00fc6d967def921" } diff --git a/plane/plane-tests/tests/common/test_env.rs b/plane/plane-tests/tests/common/test_env.rs index ef006ac0e..68586d492 100644 --- a/plane/plane-tests/tests/common/test_env.rs +++ b/plane/plane-tests/tests/common/test_env.rs @@ -121,6 +121,7 @@ impl TestEnvironment { None, None, None, + 2, ) .await .expect("Unable to construct controller.") @@ -179,6 +180,7 @@ impl TestEnvironment { None, None, Some(forward_auth.clone()), + 2, ) .await .expect("Unable to construct controller.") diff --git a/plane/src/controller/command.rs b/plane/src/controller/command.rs index 848a1b46a..613bb0319 100644 --- a/plane/src/controller/command.rs +++ b/plane/src/controller/command.rs @@ -34,6 +34,9 @@ pub struct ControllerOpts { /// (after stripping `/ctrl` and everything before it). #[clap(long)] forward_auth: Option, + + #[clap(long, default_value = "2")] + best_of: i64, } impl ControllerOpts { @@ -56,6 +59,7 @@ impl ControllerOpts { cleanup_min_age_days: self.cleanup_min_age_days, cleanup_batch_size: None, forward_auth: self.forward_auth, + best_of: self.best_of, }) } } diff --git a/plane/src/controller/core.rs b/plane/src/controller/core.rs index 377825a2d..bdabc2831 100644 --- a/plane/src/controller/core.rs +++ b/plane/src/controller/core.rs @@ -15,6 +15,7 @@ pub struct Controller { pub id: ControllerName, pub client: PlaneClient, pub default_cluster: Option, + pub best_of: i64, } pub struct NodeHandle { @@ -77,6 +78,7 @@ impl Controller { id: ControllerName, controller_url: Url, default_cluster: Option, + best_of: i64, ) -> Self { let client = PlaneClient::new(controller_url); @@ -85,6 +87,7 @@ impl Controller { id, client, default_cluster, + best_of, } } @@ -94,7 +97,12 @@ impl Controller { ) -> Result { let response = self .db - .connect(self.default_cluster.as_ref(), connect_request, &self.client) + .connect( + self.default_cluster.as_ref(), + connect_request, + &self.client, + self.best_of, + ) .await?; Ok(response) diff --git a/plane/src/controller/mod.rs b/plane/src/controller/mod.rs index c96adb5ea..e6ce2ff7a 100644 --- a/plane/src/controller/mod.rs +++ b/plane/src/controller/mod.rs @@ -172,6 +172,7 @@ impl ControllerServer { config.cleanup_min_age_days, config.cleanup_batch_size, config.forward_auth, + config.best_of, ) .await } @@ -186,6 +187,7 @@ impl ControllerServer { cleanup_min_age_days: Option, cleanup_batch_size: Option, forward_auth: Option, + best_of: i64, ) -> Result { let bind_addr = listener.local_addr()?; @@ -200,8 +202,14 @@ impl ControllerServer { let (graceful_terminate_sender, graceful_terminate_receiver) = tokio::sync::oneshot::channel::<()>(); - let controller = - Controller::new(db.clone(), id.clone(), controller_url, default_cluster).await; + let controller = Controller::new( + db.clone(), + id.clone(), + controller_url, + default_cluster, + best_of, + ) + .await; let trace_layer = TraceLayer::new_for_http() .make_span_with(DefaultMakeSpan::new().level(Level::INFO)) @@ -353,6 +361,7 @@ pub struct ControllerConfig { pub cleanup_min_age_days: Option, pub cleanup_batch_size: Option, pub forward_auth: Option, + pub best_of: i64, } pub async fn run_controller(config: ControllerConfig) -> Result<()> { diff --git a/plane/src/database/connect.rs b/plane/src/database/connect.rs index 849ffffc6..41aee3939 100644 --- a/plane/src/database/connect.rs +++ b/plane/src/database/connect.rs @@ -243,6 +243,7 @@ async fn attempt_connect( default_cluster: Option<&ClusterName>, request: &ConnectRequest, client: &PlaneClient, + best_of: i64, ) -> Result { let key = if let Some(key) = &request.key { // Request includes a key, so we need to check if it is held. @@ -317,7 +318,7 @@ async fn attempt_connect( .ok_or(ConnectError::NoClusterProvided)?; let drone = DroneDatabase::new(pool) - .pick_drone_for_spawn(cluster, &spawn_config.pool) + .pick_drone_for_spawn(cluster, &spawn_config.pool, best_of) .await? .ok_or(ConnectError::NoDroneAvailable)?; @@ -372,10 +373,11 @@ pub async fn connect( default_cluster: Option<&ClusterName>, request: &ConnectRequest, client: &PlaneClient, + best_of: i64, ) -> Result { let mut attempt = 1; loop { - match attempt_connect(pool, default_cluster, request, client).await { + match attempt_connect(pool, default_cluster, request, client, best_of).await { Ok(response) => return Ok(response), Err(error) => { if !error.retryable() || attempt >= 3 { diff --git a/plane/src/database/drone.rs b/plane/src/database/drone.rs index b2b25d541..dc106c8ab 100644 --- a/plane/src/database/drone.rs +++ b/plane/src/database/drone.rs @@ -175,6 +175,7 @@ impl<'a> DroneDatabase<'a> { &self, cluster: &ClusterName, pool: &DronePoolName, + best_of: i64, ) -> sqlx::Result> { let result = query!( r#" @@ -205,13 +206,14 @@ impl<'a> DroneDatabase<'a> { and last_local_time is not null and pool = $3 order by random() - limit 2 + limit $5 "#, cluster.to_string(), PgInterval::try_from(Duration::from_secs(UNHEALTHY_SECONDS as _)) .expect("valid interval"), pool.to_string(), BackendStatus::Terminated.to_string(), + best_of, ) .fetch_all(self.pool) .await?; diff --git a/plane/src/database/mod.rs b/plane/src/database/mod.rs index 54bfb2bb2..e68026ef6 100644 --- a/plane/src/database/mod.rs +++ b/plane/src/database/mod.rs @@ -100,8 +100,9 @@ impl PlaneDatabase { default_cluster: Option<&ClusterName>, request: &ConnectRequest, client: &PlaneClient, + best_of: i64, ) -> Result { - connect::connect(&self.pool, default_cluster, request, client).await + connect::connect(&self.pool, default_cluster, request, client, best_of).await } pub async fn revoke(&self, request: &RevokeRequest) -> Result<(), ConnectError> { connect::revoke(&self.pool, request).await