node, store: Allow passing a list of shards when placing deployments

Author: lutter

node/src/config.rs

@@ -798,7 +798,11 @@ impl Deployment {
 }
 
 impl DeploymentPlacer for Deployment {
-    fn place(&self, name: &str, network: &str) -> Result<Option<(ShardName, Vec<NodeId>)>, String> {
+    fn place(
+        &self,
+        name: &str,
+        network: &str,
+    ) -> Result<Option<(Vec<ShardName>, Vec<NodeId>)>, String> {
         // Errors here are really programming errors. We should have validated
         // everything already so that the various conversions can't fail. We
         // still return errors so that they bubble up to the deployment request
@@ -814,7 +818,7 @@ impl DeploymentPlacer for Deployment {
                             .map_err(|()| format!("{} is not a valid node name", idx))
                     })
                     .collect::<Result<Vec<_>, _>>()?;
-                Some((shard, indexers))
+                Some((vec![shard], indexers))
             }
             None => None,
         };

node/src/manager/commands/config.rs

@@ -15,11 +15,12 @@ pub fn place(placer: &dyn DeploymentPlacer, name: &str, network: &str) -> Result
                 "no matching placement rule; default placement from JSON RPC call would be used"
             );
         }
-        Some((shard, nodes)) => {
+        Some((shards, nodes)) => {
             let nodes: Vec<_> = nodes.into_iter().map(|n| n.to_string()).collect();
+            let shards: Vec<_> = shards.into_iter().map(|s| s.to_string()).collect();
             println!("subgraph: {}", name);
             println!("network:  {}", network);
-            println!("shard:    {}", shard);
+            println!("shard:    {}", shards.join(", "));
             println!("nodes:    {}", nodes.join(", "));
         }
     }

store/postgres/src/subgraph_store.rs

@@ -117,7 +117,8 @@ impl ToSql<Text, Pg> for Shard {
 /// indexers that should index it. The deployment should then be assigned to
 /// one of the returned indexers.
 pub trait DeploymentPlacer {
-    fn place(&self, name: &str, network: &str) -> Result<Option<(Shard, Vec<NodeId>)>, String>;
+    fn place(&self, name: &str, network: &str)
+        -> Result<Option<(Vec<Shard>, Vec<NodeId>)>, String>;
 }
 
 /// Tools for managing unused deployments
@@ -382,6 +383,31 @@ impl SubgraphStoreInner {
         store.find_layout(site)
     }
 
+    fn place_on_node(
+        &self,
+        mut nodes: Vec<NodeId>,
+        default_node: NodeId,
+    ) -> Result<NodeId, StoreError> {
+        match nodes.len() {
+            0 => {
+                // This is really a configuration error
+                Ok(default_node)
+            }
+            1 => Ok(nodes.pop().unwrap()),
+            _ => {
+                let conn = self.primary_conn()?;
+
+                // unwrap is fine since nodes is not empty
+                let node = conn.least_assigned_node(&nodes)?.unwrap();
+                Ok(node)
+            }
+        }
+    }
+
+    fn place_in_shard(&self, mut shards: Vec<Shard>) -> Result<Shard, StoreError> {
+        Ok(shards.pop().unwrap())
+    }
+
     fn place(
         &self,
         name: &SubgraphName,
@@ -402,16 +428,10 @@ impl SubgraphStoreInner {
 
         match placement {
             None => Ok((PRIMARY_SHARD.clone(), default_node)),
-            Some((_, nodes)) if nodes.is_empty() => {
-                // This is really a configuration error
-                Ok((PRIMARY_SHARD.clone(), default_node))
-            }
-            Some((shard, mut nodes)) if nodes.len() == 1 => Ok((shard, nodes.pop().unwrap())),
-            Some((shard, nodes)) => {
-                let conn = self.primary_conn()?;
+            Some((shards, nodes)) => {
+                let node = self.place_on_node(nodes, default_node)?;
+                let shard = self.place_in_shard(shards)?;
 
-                // unwrap is fine since nodes is not empty
-                let node = conn.least_assigned_node(&nodes)?.unwrap();
                 Ok((shard, node))
             }
         }

store/test-store/src/store.rs

@@ -142,7 +142,7 @@ pub fn remove_subgraphs() {
         .expect("deleting test entities succeeds");
 }
 
-pub fn place(name: &str) -> Result<Option<(Shard, Vec<NodeId>)>, String> {
+pub fn place(name: &str) -> Result<Option<(Vec<Shard>, Vec<NodeId>)>, String> {
     CONFIG.deployment.place(name, NETWORK_NAME)
 }