fix: improve handling for rollout percentage reached (#14870)

airbytehq · Dec 19, 2024 · 674c19e · 674c19e
1 parent 3bb941b
commit 674c19e
Show file tree

Hide file tree

Showing 5 changed files with 50 additions and 18 deletions.
diff --git a/airbyte-api/problems-api/src/main/openapi/api-problems.yaml b/airbyte-api/problems-api/src/main/openapi/api-problems.yaml
@@ -837,6 +837,27 @@ components:
               default: Invalid connector rollout request.
             data:
               $ref: "#/components/schemas/ProblemMessageData"
+    ConnectorRolloutMaximumRolloutPercentageReachedProblemResponse:
+      x-implements: io.airbyte.api.problems.ProblemResponse
+      type: object
+      allOf:
+        - $ref: "#/components/schemas/BaseProblemFields"
+        - type: object
+          properties:
+            status:
+              type: integer
+              default: 400
+            type:
+              type: string
+              default: error:connector-rollout/rollout-percentage-reached
+            title:
+              type: string
+              default: Max rollout percentage already reached
+            detail:
+              type: string
+              default: Cannot pin more actors; the maximum rollout percentage was already reached
+            data:
+              $ref: "#/components/schemas/ProblemMessageData"
     ConnectorRolloutNotEnoughActorsProblemResponse:
       x-implements: io.airbyte.api.problems.ProblemResponse
       type: object

diff --git a/...mons-server/src/main/kotlin/io/airbyte/commons/server/handlers/ConnectorRolloutHandler.kt b/...mons-server/src/main/kotlin/io/airbyte/commons/server/handlers/ConnectorRolloutHandler.kt
@@ -20,6 +20,7 @@ import io.airbyte.api.model.generated.ConnectorRolloutStrategy
 import io.airbyte.api.model.generated.ConnectorRolloutUpdateStateRequestBody
 import io.airbyte.api.problems.model.generated.ProblemMessageData
 import io.airbyte.api.problems.throwable.generated.ConnectorRolloutInvalidRequestProblem
+import io.airbyte.api.problems.throwable.generated.ConnectorRolloutMaximumRolloutPercentageReachedProblem
 import io.airbyte.api.problems.throwable.generated.ConnectorRolloutNotEnoughActorsProblem
 import io.airbyte.config.ConnectorEnumRolloutState
 import io.airbyte.config.ConnectorEnumRolloutStrategy
@@ -393,7 +394,7 @@ open class ConnectorRolloutHandler
       }
 
       if (connectorRollout.currentTargetRolloutPct >= actualTargetRolloutPct) {
-        throw ConnectorRolloutInvalidRequestProblem(
+        throw ConnectorRolloutMaximumRolloutPercentageReachedProblem(
           ProblemMessageData().message(
             "Requested to pin $actualTargetRolloutPct% of actors but already pinned ${connectorRollout.currentTargetRolloutPct}.",
           ),

diff --git a/...-server/src/test/kotlin/io/airbyte/commons/server/handlers/ConnectorRolloutHandlerTest.kt b/...-server/src/test/kotlin/io/airbyte/commons/server/handlers/ConnectorRolloutHandlerTest.kt
@@ -12,6 +12,7 @@ import io.airbyte.api.model.generated.ConnectorRolloutStateTerminal
 import io.airbyte.api.model.generated.ConnectorRolloutStrategy
 import io.airbyte.api.model.generated.ConnectorRolloutUpdateStateRequestBody
 import io.airbyte.api.problems.throwable.generated.ConnectorRolloutInvalidRequestProblem
+import io.airbyte.api.problems.throwable.generated.ConnectorRolloutMaximumRolloutPercentageReachedProblem
 import io.airbyte.api.problems.throwable.generated.ConnectorRolloutNotEnoughActorsProblem
 import io.airbyte.config.ActorDefinitionVersion
 import io.airbyte.config.ConnectorEnumRolloutState
@@ -848,10 +849,10 @@ internal class ConnectorRolloutHandlerTest {
     connectorRollout.withFinalTargetRolloutPct(50)
     connectorRollout.withCurrentTargetRolloutPct(50)
 
-    assertThrows<ConnectorRolloutInvalidRequestProblem> {
+    assertThrows<ConnectorRolloutMaximumRolloutPercentageReachedProblem> {
       connectorRolloutHandler.getValidPercentageToPin(connectorRollout, 50, rolloutStrategy = ConnectorRolloutStrategy.AUTOMATED)
     }
-    assertThrows<ConnectorRolloutInvalidRequestProblem> {
+    assertThrows<ConnectorRolloutMaximumRolloutPercentageReachedProblem> {
       connectorRolloutHandler.getValidPercentageToPin(connectorRollout, 100, rolloutStrategy = ConnectorRolloutStrategy.AUTOMATED)
     }
   }

diff --git a/airbyte-connector-rollout-worker/build.gradle.kts b/airbyte-connector-rollout-worker/build.gradle.kts
@@ -17,6 +17,7 @@ dependencies {
   implementation(project(mapOf("path" to ":oss:airbyte-commons-temporal")))
   implementation(libs.okhttp)
   implementation(project(":oss:airbyte-config:config-models"))
+  implementation(project(":oss:airbyte-api:problems-api"))
   implementation(project(":oss:airbyte-api:server-api"))
   implementation(project(":oss:airbyte-connector-rollout-shared"))
   implementation(project(":oss:airbyte-commons-storage"))

diff --git a/...orker/src/main/kotlin/io/airbyte/connector/rollout/worker/ConnectorRolloutWorkflowImpl.kt b/...orker/src/main/kotlin/io/airbyte/connector/rollout/worker/ConnectorRolloutWorkflowImpl.kt
@@ -5,6 +5,7 @@
 package io.airbyte.connector.rollout.worker
 
 import com.google.common.annotations.VisibleForTesting
+import io.airbyte.api.problems.model.generated.ConnectorRolloutMaximumRolloutPercentageReachedProblemResponse
 import io.airbyte.config.ConnectorEnumRolloutState
 import io.airbyte.config.ConnectorEnumRolloutStrategy
 import io.airbyte.config.ConnectorRolloutFinalState
@@ -210,6 +211,7 @@ class ConnectorRolloutWorkflowImpl : ConnectorRolloutWorkflow {
             " targetPercentage=$targetPercentageToPin" +
             " rolloutStrategy = ${rollout.rolloutStrategy}"
         }
+        nextRolloutStageAt = currentTime.plusSeconds(waitBetweenRolloutsSeconds.toLong())
         try {
           connectorRollout =
             progressRollout(
@@ -224,22 +226,23 @@ class ConnectorRolloutWorkflowImpl : ConnectorRolloutWorkflow {
                 rolloutStrategy = ConnectorEnumRolloutStrategy.AUTOMATED,
               ),
             )
-          nextRolloutStageAt = currentTime.plusSeconds(waitBetweenRolloutsSeconds.toLong())
         } catch (e: Exception) {
-          logger.error { "Failed to advance the rollout. workflowId=$workflowId e=${e.message}" }
-          connectorRollout =
-            pauseRollout(
-              ConnectorRolloutActivityInputPause(
-                input.dockerRepository,
-                input.dockerImageTag,
-                input.actorDefinitionId,
-                input.rolloutId,
-                "Paused due to an exception while pinning connections: $e",
-                null,
-                ConnectorEnumRolloutStrategy.AUTOMATED,
-              ),
-            )
-          break
+          logger.error { "Failed to advance the rollout. workflowId=$workflowId e=${e.message} e.cause=${e.cause} e=$e" }
+          if (!rolloutPercentageReached(e)) {
+            connectorRollout =
+              pauseRollout(
+                ConnectorRolloutActivityInputPause(
+                  input.dockerRepository,
+                  input.dockerImageTag,
+                  input.actorDefinitionId,
+                  input.rolloutId,
+                  "Paused due to an exception while pinning connections: $e",
+                  null,
+                  ConnectorEnumRolloutStrategy.AUTOMATED,
+                ),
+              )
+            break
+          }
         }
       }
 
@@ -281,6 +284,11 @@ class ConnectorRolloutWorkflowImpl : ConnectorRolloutWorkflow {
     return getRolloutState()
   }
 
+  private fun rolloutPercentageReached(e: Exception): Boolean {
+    logger.info { "e.cause.message = ${e.cause?.message}" }
+    return e.cause?.message?.contains(ConnectorRolloutMaximumRolloutPercentageReachedProblemResponse().type) ?: false
+  }
+
   private fun getCurrentTimeMilli(): Instant {
     return Instant.ofEpochMilli(Workflow.currentTimeMillis())
   }