Configurable best-of

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.")

plane/src/controller/command.rs

@@ -34,6 +34,9 @@ pub struct ControllerOpts {
     /// (after stripping `/ctrl` and everything before it).
     #[clap(long)]
     forward_auth: Option<Url>,
+
+    #[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,
         })
     }
 }

plane/src/controller/core.rs

@@ -15,6 +15,7 @@ pub struct Controller {
     pub id: ControllerName,
     pub client: PlaneClient,
     pub default_cluster: Option<ClusterName>,
+    pub best_of: i64,
 }
 
 pub struct NodeHandle {
@@ -77,6 +78,7 @@ impl Controller {
         id: ControllerName,
         controller_url: Url,
         default_cluster: Option<ClusterName>,
+        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<ConnectResponse, ConnectError> {
         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)

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<i32>,
         cleanup_batch_size: Option<i32>,
         forward_auth: Option<Url>,
+        best_of: i64,
     ) -> Result<Self> {
         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<i32>,
     pub cleanup_batch_size: Option<i32>,
     pub forward_auth: Option<Url>,
+    pub best_of: i64,
 }
 
 pub async fn run_controller(config: ControllerConfig) -> Result<()> {

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<ConnectResponse> {
     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<ConnectResponse> {
     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 {

plane/src/database/drone.rs

@@ -104,7 +104,7 @@
     ) -> sqlx::Result<Vec<DroneWithMetadata>> {
         // Converting to a PgInterval fails if the duration has nanosecond granularity.
         // So let's round up to the nearest microsecond.
         let seen_in_last = if seen_in_last.subsec_nanos() % 1000 == 0 {
             seen_in_last
         } else {
             let nanos = (seen_in_last.subsec_micros() + 1) * 1000;
@@ -175,6 +175,7 @@
         &self,
         cluster: &ClusterName,
         pool: &DronePoolName,
+        best_of: i64,
     ) -> sqlx::Result<Option<DroneForSpawn>> {
         let result = query!(
             r#"
@@ -205,13 +206,14 @@
                 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?;

plane/src/database/mod.rs

@@ -56,23 +56,23 @@
         }
     }
 
     pub fn acme(&self) -> acme::AcmeDatabase {
         AcmeDatabase::new(&self.pool)
     }
 
     pub fn drone(&self) -> DroneDatabase {
         DroneDatabase::new(&self.pool)
     }
 
     pub fn cluster(&self) -> ClusterDatabase {
         ClusterDatabase::new(&self.pool)
     }
 
     pub fn node(&self) -> NodeDatabase {
         NodeDatabase::new(&self.pool)
     }
 
     pub fn backend(&self) -> BackendDatabase {
         BackendDatabase::new(self)
     }
 
@@ -80,11 +80,11 @@
         BackendActionDatabase::new(&self.pool)
     }
 
     pub fn keys(&self) -> backend_key::KeysDatabase {
         KeysDatabase::new(&self.pool)
     }
 
     pub fn controller(&self) -> controller::ControllerDatabase {
         ControllerDatabase::new(&self.pool)
     }
 
@@ -100,8 +100,9 @@
         default_cluster: Option<&ClusterName>,
         request: &ConnectRequest,
         client: &PlaneClient,
+        best_of: i64,
     ) -> Result<ConnectResponse, ConnectError> {
-        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