From 51a68766097d76f12799133c174d05c7289facca Mon Sep 17 00:00:00 2001 From: Wojciech Baszczyk Date: Wed, 23 Oct 2024 15:29:05 +0200 Subject: [PATCH] fix(api): mongo upgraders batch updates to ensure no timeouts --- .../MissingEnvironmentUpgrader.java | 23 ++++++++++++++----- .../plan/PlanDefinitionVersionUpgrader.java | 14 +++++++++-- 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/gravitee-apim-repository/gravitee-apim-repository-mongodb/src/main/java/io/gravitee/repository/mongodb/management/upgrade/upgrader/environment/MissingEnvironmentUpgrader.java b/gravitee-apim-repository/gravitee-apim-repository-mongodb/src/main/java/io/gravitee/repository/mongodb/management/upgrade/upgrader/environment/MissingEnvironmentUpgrader.java index eb6326fd860..c9d1493710a 100644 --- a/gravitee-apim-repository/gravitee-apim-repository-mongodb/src/main/java/io/gravitee/repository/mongodb/management/upgrade/upgrader/environment/MissingEnvironmentUpgrader.java +++ b/gravitee-apim-repository/gravitee-apim-repository-mongodb/src/main/java/io/gravitee/repository/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 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 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/gravitee-apim-repository/gravitee-apim-repository-mongodb/src/main/java/io/gravitee/repository/mongodb/management/upgrade/upgrader/plan/PlanDefinitionVersionUpgrader.java b/gravitee-apim-repository/gravitee-apim-repository-mongodb/src/main/java/io/gravitee/repository/mongodb/management/upgrade/upgrader/plan/PlanDefinitionVersionUpgrader.java index 08270f6194e..01c652c6f7e 100644 --- a/gravitee-apim-repository/gravitee-apim-repository-mongodb/src/main/java/io/gravitee/repository/mongodb/management/upgrade/upgrader/plan/PlanDefinitionVersionUpgrader.java +++ b/gravitee-apim-repository/gravitee-apim-repository-mongodb/src/main/java/io/gravitee/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 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; }