Parallelize dumping schemas to disk

Reviewed By: arseniy-panfilov

Differential Revision: D61515272

fbshipit-source-id: e67feb735db2911acafe7a68133ca1a0c5f56d55

Author: Deepak Singh

compiler/crates/schema-print/Cargo.toml

@@ -1,4 +1,4 @@
-# @generated by autocargo from //relay/oss/crates/schema-print:[schema-print,schema-print_test]
+# @generated by autocargo from //relay/oss/crates/schema-print:[schema-print,schema-print-consistency_test,schema-print-par_test,schema-print_test]
 
 [package]
 name = "schema-print"
@@ -8,6 +8,14 @@ edition = "2021"
 repository = "https://github.com/facebook/relay"
 license = "MIT"
 
+[[test]]
+name = "schema_print_consistency_test"
+path = "tests/print_schema_consistency_test.rs"
+
+[[test]]
+name = "schema_print_par_test"
+path = "tests/print_schema_in_parallel_test.rs"
+
 [[test]]
 name = "schema_print_test"
 path = "tests/print_schema_test.rs"
@@ -16,8 +24,10 @@ path = "tests/print_schema_test.rs"
 fnv = "1.0"
 intern = { path = "../intern" }
 itertools = "0.11.0"
+rayon = "1.9.0"
 schema = { path = "../schema" }
 
 [dev-dependencies]
+diff = "0.1"
 fixture-tests = { path = "../fixture-tests" }
 tokio = { version = "1.37.0", features = ["full", "test-util", "tracing"] }

compiler/crates/schema-print/src/lib.rs

@@ -11,6 +11,18 @@
 
 mod print_schema;
 mod printer;
+mod shard_printer;
+
+use std::hash::DefaultHasher;
+use std::hash::Hash;
+use std::hash::Hasher;
 
 pub use print_schema::*;
 pub use printer::*;
+pub use shard_printer::*;
+
+pub fn calculate_hash<T: Hash>(t: &T) -> u64 {
+    let mut s = DefaultHasher::new();
+    t.hash(&mut s);
+    s.finish()
+}

compiler/crates/schema-print/src/print_schema.rs

@@ -12,11 +12,17 @@ use fnv::FnvHashMap;
 use intern::string_key::StringKey;
 use intern::Lookup;
 use itertools::Itertools;
+use rayon::iter::IntoParallelIterator;
+use rayon::iter::ParallelIterator;
 use schema::*;
 
+use crate::generate_shard_map;
+use crate::generate_typed_shard_map;
 use crate::has_schema_definition_types;
 use crate::is_schema_of_common_name;
 use crate::Printer;
+use crate::ShardPrinter;
+use crate::TypedShardPrinter;
 use crate::DEAULT_SHARD_COUNT;
 
 pub fn print(schema: &SDLSchema) -> String {
@@ -61,6 +67,8 @@ pub fn print_types(schema: &SDLSchema) -> String {
 
 /// Returns a sharded GraphQL SDLSchema
 ///
+/// Prints shards in sequence. No parallelism is used.
+///
 /// # Arguments
 ///
 /// * `schema` - GraphQL SDLSchema
@@ -108,6 +116,73 @@ pub fn print_types_directives_as_shards(
     shards
 }
 
+/**
+ * Prints a sharded GraphQL SDLSchema in parallel
+ */
+pub fn print_types_directives_as_shards_in_parallel(
+    schema: &SDLSchema,
+    shard_count: usize,
+    type_shard_count_map: FnvHashMap<StringKey, usize>,
+) -> FnvHashMap<usize, String> {
+    let typeshard_count: usize = type_shard_count_map.values().sum();
+    let base_shard_count = shard_count - typeshard_count;
+    let shard_map = generate_shard_map(
+        schema,
+        base_shard_count,
+        type_shard_count_map.keys().copied().collect(),
+    );
+    let typed_shard_map = generate_typed_shard_map(schema, type_shard_count_map);
+
+    // Base shards
+    let mut shards = shard_map
+        .into_par_iter()
+        .map(|(shard_key, shard)| {
+            let printer = ShardPrinter::new(schema, shard);
+            // Print directives
+            let content = if shard_key == 0 {
+                printer.print_shard_with_directives()
+            } else {
+                printer.print_shard()
+            };
+            (shard_key, content)
+        })
+        .collect::<FnvHashMap<usize, String>>();
+
+    // Typed shards (e.g. Query, Mutation, etc)
+    let mut typed_shards = typed_shard_map
+        .into_par_iter()
+        .map(|(type_name, shard_map)| {
+            let type_shard_count = shard_map.len();
+            let typed_shard = shard_map
+                .into_par_iter()
+                .map(|(shard_key, fields)| {
+                    let printer = TypedShardPrinter::new(
+                        schema,
+                        type_name,
+                        fields,
+                        shard_key == 0,
+                        shard_key == type_shard_count - 1,
+                    );
+                    (shard_key, printer.print_shard())
+                })
+                .collect();
+            (type_name, typed_shard)
+        })
+        .collect::<FnvHashMap<StringKey, FnvHashMap<usize, String>>>();
+
+    // Drain typed_shards into shards
+    typed_shards
+        .iter_mut()
+        .sorted_by_key(|(k, _)| k.lookup())
+        .for_each(|(_type_name, type_shards)| {
+            let curr_shard_count = shards.len();
+            type_shards.drain().for_each(|(shard_key, type_shard)| {
+                shards.insert(curr_shard_count + shard_key, type_shard);
+            });
+        });
+    shards
+}
+
 pub fn print_type(schema: &SDLSchema, type_: Type) -> String {
     let mut result = vec![String::new(); DEAULT_SHARD_COUNT];
     write_type(schema, &mut result, type_).unwrap();
@@ -135,7 +210,7 @@ fn write_directive(
 
 fn write_types(schema: &SDLSchema, result: &mut Vec<String>) -> FmtResult {
     let mut printer = Printer::new(schema, result);
-    printer.print_types()
+    printer.print_all_types()
 }
 
 fn write_types_as_shards<'a>(
@@ -145,7 +220,7 @@ fn write_types_as_shards<'a>(
     type_shards: &'a mut FnvHashMap<StringKey, Vec<String>>,
 ) -> FmtResult {
     let mut printer = Printer::new_with_shards(schema, shards, shard_count, Some(type_shards));
-    printer.print_types()
+    printer.print_all_types()
 }
 
 fn write_type(schema: &SDLSchema, result: &mut Vec<String>, type_: Type) -> FmtResult {

compiler/crates/schema-print/src/printer.rs

@@ -1,9 +1,13 @@
-use std::collections::hash_map::DefaultHasher;
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 use std::collections::BTreeMap;
 use std::fmt::Result as FmtResult;
 use std::fmt::Write;
-use std::hash::Hash;
-use std::hash::Hasher;
 
 use fnv::FnvHashMap;
 use intern::string_key::Intern;
@@ -12,9 +16,11 @@ use intern::Lookup;
 use itertools::Itertools;
 use schema::*;
 
+use crate::calculate_hash;
+
 pub const DEAULT_SHARD_COUNT: usize = 1;
 
-pub(crate) struct Printer<'schema, 'writer> {
+pub struct Printer<'schema, 'writer> {
     schema: &'schema SDLSchema,
     // Each writer here represents a shard
     writers: &'writer mut Vec<String>,
@@ -111,14 +117,21 @@ impl<'schema, 'writer> Printer<'schema, 'writer> {
         self.print_definition_end()
     }
 
-    pub fn print_types(&mut self) -> FmtResult {
+    pub fn print_all_types(&mut self) -> FmtResult {
         let ordered_type_map = self.schema.get_type_map().collect::<BTreeMap<_, _>>();
         for (_key, value) in ordered_type_map.iter() {
             self.print_type(**value)?;
         }
         Ok(())
     }
 
+    pub fn print_types(&mut self, types_to_print: &[(StringKey, Type)]) -> FmtResult {
+        for (_key, value) in types_to_print.iter() {
+            self.print_type(*value)?;
+        }
+        Ok(())
+    }
+
     pub fn print_type(&mut self, type_: Type) -> FmtResult {
         self.update_writer_index_for_type_start(&type_);
         match type_ {
@@ -195,7 +208,7 @@ impl<'schema, 'writer> Printer<'schema, 'writer> {
         self.print_definition_end()
     }
 
-    fn print_fields(&mut self, fields: &[FieldID], type_name: impl Lookup) -> FmtResult {
+    pub fn print_fields(&mut self, fields: &[FieldID], type_name: impl Lookup) -> FmtResult {
         if fields.is_empty() {
             return Ok(());
         }
@@ -218,6 +231,26 @@ impl<'schema, 'writer> Printer<'schema, 'writer> {
         write!(self.writer(), "}}")
     }
 
+    // Prints a set of fields indented by 2 spaces
+    // Does not handle writing the type definition and open/closing brackets
+    pub fn print_fields_without_brackets(&mut self, fields: &[FieldID]) -> FmtResult {
+        if fields.is_empty() {
+            return Ok(());
+        }
+        for field_id in fields {
+            let field = &self.schema.field(*field_id);
+            self.print_space()?;
+            self.print_space()?;
+            write!(self.writer(), "{}", field.name.item)?;
+            self.print_args(&field.arguments)?;
+            let type_string = self.schema.get_type_string(&field.type_);
+            write!(self.writer(), ": {}", type_string)?;
+            self.print_directive_values(&field.directives)?;
+            self.print_new_line()?;
+        }
+        Ok(())
+    }
+
     fn print_args(&mut self, args: &ArgumentDefinitions) -> FmtResult {
         if args.is_empty() {
             return Ok(());
@@ -311,7 +344,15 @@ impl<'schema, 'writer> Printer<'schema, 'writer> {
         Ok(())
     }
 
-    fn print_definition_end(&mut self) -> FmtResult {
+    pub fn print_type_declaration(&mut self, type_name: StringKey) -> FmtResult {
+        writeln!(self.writer(), "type {} {{", type_name)
+    }
+
+    pub fn print_definition_closure(&mut self) -> FmtResult {
+        write!(self.writer(), "}}")
+    }
+
+    pub fn print_definition_end(&mut self) -> FmtResult {
         self.print_new_line()?;
         self.print_new_line()
     }
@@ -381,9 +422,3 @@ pub fn has_schema_definition_types(schema: &SDLSchema) -> bool {
         || schema.mutation_type().is_some()
         || schema.subscription_type().is_some()
 }
-
-pub fn calculate_hash<T: Hash>(t: &T) -> u64 {
-    let mut s = DefaultHasher::new();
-    t.hash(&mut s);
-    s.finish()
-}