From 5761d9e7909ffb393e70d96869254d7f59bdb887 Mon Sep 17 00:00:00 2001
From: Kamil Kisiela <kamil.kisiela@gmail.com>
Date: Tue, 23 Aug 2022 14:31:02 +0200
Subject: [PATCH] Fix sample_rate in Router (#297)

* Fix sample_rate
* Sort definitions
---
 packages/libraries/router/src/graphql.rs | 38 ++++++++++++++++++++----
 packages/libraries/router/src/usage.rs   |  2 +-
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/packages/libraries/router/src/graphql.rs b/packages/libraries/router/src/graphql.rs
index 20fc6a4f31..67e9c80c51 100644
--- a/packages/libraries/router/src/graphql.rs
+++ b/packages/libraries/router/src/graphql.rs
@@ -270,6 +270,19 @@ impl<'s, T: Text<'s> + Clone> SortSelectionsTransform<'s, T> {
 }
 
 impl<'s, T: Text<'s> + Clone> OperationTransformer<'s, T> for SortSelectionsTransform<'s, T> {
+    fn transform_document(
+        &mut self,
+        document: &Document<'s, T>,
+    ) -> TransformedValue<Document<'s, T>> {
+        let mut next_definitions = self
+            .transform_list(&document.definitions, Self::transform_definition)
+            .replace_or_else(|| document.definitions.to_vec());
+        next_definitions.sort_unstable_by(|a, b| self.compare_definitions(a, b));
+        TransformedValue::Replace(Document {
+            definitions: next_definitions,
+        })
+    }
+
     fn transform_selection_set(
         &mut self,
         selections: &SelectionSet<'s, T>,
@@ -364,6 +377,17 @@ impl<'s, T: Text<'s> + Clone> OperationTransformer<'s, T> for SortSelectionsTran
 }
 
 impl<'s, T: Text<'s> + Clone> SortSelectionsTransform<'s, T> {
+    fn compare_definitions(&self, a: &Definition<'s, T>, b: &Definition<'s, T>) -> Ordering {
+        match (a, b) {
+            // Keep operations as they are
+            (Definition::Operation(_), Definition::Operation(_)) => Ordering::Equal,
+            // Sort fragments by name
+            (Definition::Fragment(a), Definition::Fragment(b)) => a.name.cmp(&b.name),
+            // Operation -> Fragment
+            _ => definition_kind_ordering(a).cmp(&definition_kind_ordering(b)),
+        }
+    }
+
     fn compare_selections(&self, a: &Selection<'s, T>, b: &Selection<'s, T>) -> Ordering {
         match (a, b) {
             (Selection::Field(a), Selection::Field(b)) => a.name.cmp(&b.name),
@@ -373,12 +397,6 @@ impl<'s, T: Text<'s> + Clone> SortSelectionsTransform<'s, T> {
             _ => {
                 let a_ordering = selection_kind_ordering(a);
                 let b_ordering = selection_kind_ordering(b);
-                assert!(
-                    a_ordering != b_ordering,
-                    "expected different ordering, got {} == {}",
-                    a_ordering,
-                    b_ordering
-                );
                 a_ordering.cmp(&b_ordering)
             }
         }
@@ -411,6 +429,14 @@ fn selection_kind_ordering<'s, T: Text<'s>>(selection: &Selection<'s, T>) -> u8
     }
 }
 
+/// Assigns an order to different variants of Definition
+fn definition_kind_ordering<'a, T: Text<'a>>(definition: &Definition<'a, T>) -> u8 {
+    match definition {
+        Definition::Operation(_) => 1,
+        Definition::Fragment(_) => 2,
+    }
+}
+
 #[derive(Clone)]
 pub struct ProcessedOperation {
     pub operation: String,
diff --git a/packages/libraries/router/src/usage.rs b/packages/libraries/router/src/usage.rs
index d6ee62a53d..621a92c917 100644
--- a/packages/libraries/router/src/usage.rs
+++ b/packages/libraries/router/src/usage.rs
@@ -116,7 +116,7 @@ impl UsagePlugin {
         let mut dropped = match sample_rate {
             Some(rate) => {
                 let num: f64 = rng.gen();
-                num <= rate
+                num > rate
             }
             None => false,
         };