fix(api): mongo upgraders batch updates to ensure no timeouts

gravitee-io · Oct 24, 2024 · 51a6876 · 51a6876
1 parent 2444294
commit 51a6876
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 8 deletions.
diff --git a/...epository/mongodb/management/upgrade/upgrader/environment/MissingEnvironmentUpgrader.java b/...epository/mongodb/management/upgrade/upgrader/environment/MissingEnvironmentUpgrader.java
@@ -33,11 +33,12 @@
 @Component
 public class MissingEnvironmentUpgrader extends MongoUpgrader {
 
+    private static final int UPGRADER_BATCH_SIZE = 1000;
     public static final int MISSING_ENVIRONMENT_UPGRADER_ORDER = ThemeTypeUpgrader.THEME_TYPE_UPGRADER_ORDER + 1;
 
     @Override
     public String version() {
-        return "v1";
+        return "v2";
     }
 
     @Override
@@ -64,10 +65,15 @@ private void updateEnvironmentFromApi(final Set<Boolean> upgradeStatus) {
             );
 
         if (!bulkActions.isEmpty()) {
-            // This upgrade is only done on data created before 3.17.0 as ApiKey#api as been deprecated
-            upgradeStatus.add(this.getCollection("keys").bulkWrite(bulkActions).wasAcknowledged());
-            upgradeStatus.add(this.getCollection("plans").bulkWrite(bulkActions).wasAcknowledged());
-            upgradeStatus.add(this.getCollection("subscriptions").bulkWrite(bulkActions).wasAcknowledged());
+            for (int i = 0; i < bulkActions.size(); i += UPGRADER_BATCH_SIZE) {
+                // Get the end index for the sublist, ensuring it doesn't go out of bounds
+                int end = Math.min(i + UPGRADER_BATCH_SIZE, bulkActions.size());
+                var batchActions = bulkActions.subList(i, end);
+                // This upgrade is only done on data created before 3.17.0 as ApiKey#api as been deprecated
+                upgradeStatus.add(this.getCollection("keys").bulkWrite(batchActions).wasAcknowledged());
+                upgradeStatus.add(this.getCollection("plans").bulkWrite(batchActions).wasAcknowledged());
+                upgradeStatus.add(this.getCollection("subscriptions").bulkWrite(batchActions).wasAcknowledged());
+            }
         }
     }
 
@@ -87,7 +93,12 @@ private void updateEnvironmentFromSubscription(final Set<Boolean> upgradeStatus)
             );
 
         if (!bulkActions.isEmpty()) {
-            upgradeStatus.add(this.getCollection("keys").bulkWrite(bulkActions).wasAcknowledged());
+            for (int i = 0; i < bulkActions.size(); i += UPGRADER_BATCH_SIZE) {
+                // Get the end index for the sublist, ensuring it doesn't go out of bounds
+                int end = Math.min(i + UPGRADER_BATCH_SIZE, bulkActions.size());
+                var batchActions = bulkActions.subList(i, end);
+                upgradeStatus.add(this.getCollection("keys").bulkWrite(batchActions).wasAcknowledged());
+            }
         }
     }
 

diff --git a/...ee/repository/mongodb/management/upgrade/upgrader/plan/PlanDefinitionVersionUpgrader.java b/...ee/repository/mongodb/management/upgrade/upgrader/plan/PlanDefinitionVersionUpgrader.java
@@ -23,6 +23,8 @@
 import io.gravitee.repository.mongodb.management.upgrade.upgrader.common.MongoUpgrader;
 import io.gravitee.repository.mongodb.management.upgrade.upgrader.dashboards.DashboardTypeUpgrader;
 import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.Set;
 import org.bson.Document;
 import org.springframework.stereotype.Component;
 
@@ -32,15 +34,17 @@
 @Component
 public class PlanDefinitionVersionUpgrader extends MongoUpgrader {
 
+    private static final int UPGRADER_BATCH_SIZE = 1000;
     public static final int PLAN_DEFINITION_VERSION_UPGRADER_ORDER = DashboardTypeUpgrader.DASHBOARD_TYPE_UPGRADER_ORDER + 1;
 
     @Override
     public String version() {
-        return "v1";
+        return "v2";
     }
 
     @Override
     public boolean upgrade() {
+        Set<Boolean> upgradeStatus = new HashSet<>();
         var query = new Document("definitionVersion", DefinitionVersion.V4.name());
         var projection = Projections.fields(Projections.include("_id", "definitionVersion"));
 
@@ -58,7 +62,13 @@ public boolean upgrade() {
             });
 
         if (!bulkActions.isEmpty()) {
-            return this.getCollection("plans").bulkWrite(bulkActions).wasAcknowledged();
+            for (int i = 0; i < bulkActions.size(); i += UPGRADER_BATCH_SIZE) {
+                // Get the end index for the sublist, ensuring it doesn't go out of bounds
+                int end = Math.min(i + UPGRADER_BATCH_SIZE, bulkActions.size());
+                var batchActions = bulkActions.subList(i, end);
+                upgradeStatus.add(this.getCollection("plans").bulkWrite(batchActions).wasAcknowledged());
+            }
+            return upgradeStatus.stream().allMatch(Boolean.TRUE::equals);
         }
         return true;
     }