From 2ef7cd3f2dd9275c7cca155ae880b26f9f35cab2 Mon Sep 17 00:00:00 2001
From: Grzegorz Zagata <grzegorz.zagata2@ing.com>
Date: Mon, 16 Oct 2023 08:18:18 +0200
Subject: [PATCH] propagate slowdown errors to client when checking bucket
 location

---
 src/main/resources/application.conf                        | 3 ++-
 src/main/resources/reference.conf                          | 3 ++-
 .../ing/wbaa/rokku/proxy/config/StorageS3Settings.scala    | 1 +
 .../rokku/proxy/handler/namespace/NamespacesHandler.scala  | 7 +++++++
 4 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/main/resources/application.conf b/src/main/resources/application.conf
index 539611fc..3f7c0df3 100644
--- a/src/main/resources/application.conf
+++ b/src/main/resources/application.conf
@@ -24,7 +24,8 @@ rokku {
         }
         region = ${?ROKKU_STORAGE_S3_AWS_REGION}
         v2SignatureEnabled = ${?ROKKU_STORAGE_S3_V2_ENABLED}
-
+        # To add more then one code, seperate them with comma
+        slowdownCodes = ${?ROKKU_STORAGE_S3_SLOWDOWN_CODES}
         healthCheck {
            # can be one of:
            #   s3ListBucket  - uses AWS S3 client to list single bucket
diff --git a/src/main/resources/reference.conf b/src/main/resources/reference.conf
index 1ac61a2f..0cd720f8 100644
--- a/src/main/resources/reference.conf
+++ b/src/main/resources/reference.conf
@@ -31,7 +31,8 @@ rokku {
         }
         region = "us-east-1"
         v2SignatureEnabled = false
-
+        # To add more then one code, seperate them with comma
+        slowdownCodes = "503"
         healthCheck {
            # can be one of:
            #   s3ListBucket  - uses AWS S3 client to list single bucket
diff --git a/src/main/scala/com/ing/wbaa/rokku/proxy/config/StorageS3Settings.scala b/src/main/scala/com/ing/wbaa/rokku/proxy/config/StorageS3Settings.scala
index cebcf4e1..5d9507c5 100644
--- a/src/main/scala/com/ing/wbaa/rokku/proxy/config/StorageS3Settings.scala
+++ b/src/main/scala/com/ing/wbaa/rokku/proxy/config/StorageS3Settings.scala
@@ -14,6 +14,7 @@ class StorageS3Settings(config: Config) extends Extension {
   val storageS3AdminSecretkey: String = config.getString("rokku.storage.s3.admin.secretkey")
   val awsRegion: String = config.getString("rokku.storage.s3.region")
   val v2SignatureEnabled: Boolean = config.getBoolean("rokku.storage.s3.v2SignatureEnabled")
+  val slowdownCodes: Array[Int] = config.getString("rokku.storage.s3.slowdownCodes").split(",").map(o => o.toInt)
   val isRequestUserQueueEnabled: Boolean = config.getBoolean("rokku.storage.s3.request.queue.enable")
   private val hcMethodString = config.getString("rokku.storage.s3.healthCheck.method")
   val hcMethod: HCMethod = hcMethodString match {
diff --git a/src/main/scala/com/ing/wbaa/rokku/proxy/handler/namespace/NamespacesHandler.scala b/src/main/scala/com/ing/wbaa/rokku/proxy/handler/namespace/NamespacesHandler.scala
index c309054a..ac5eda1a 100644
--- a/src/main/scala/com/ing/wbaa/rokku/proxy/handler/namespace/NamespacesHandler.scala
+++ b/src/main/scala/com/ing/wbaa/rokku/proxy/handler/namespace/NamespacesHandler.scala
@@ -4,8 +4,10 @@ import akka.http.scaladsl.model.HttpRequest
 import com.amazonaws.auth.BasicAWSCredentials
 import com.amazonaws.services.s3.model.AmazonS3Exception
 import com.ing.wbaa.rokku.proxy.config.NamespaceSettings
+import com.ing.wbaa.rokku.proxy.config.StorageS3Settings
 import com.ing.wbaa.rokku.proxy.data.RequestId
 import com.ing.wbaa.rokku.proxy.handler.LoggerHandlerWithId
+import com.ing.wbaa.rokku.proxy.handler.exception.RokkuThrottlingException
 import com.ing.wbaa.rokku.proxy.metrics.MetricsFactory.{ incrementBucketNamespaceCacheHit, incrementBucketNamespacesNotFound, incrementBucketNamespacesSearch }
 import com.ing.wbaa.rokku.proxy.util.S3Utils
 
@@ -21,6 +23,7 @@ trait NamespacesHandler {
   private val bucketCredentials: scala.collection.concurrent.Map[BucketName, BasicAWSCredentials] = scala.collection.concurrent.TrieMap[BucketName, BasicAWSCredentials]()
 
   protected[this] val namespaceSettings: NamespaceSettings
+  protected[this] def storageS3Settings: StorageS3Settings
 
   private def namespaceCredentials: ListMap[NamespaceName, BasicAWSCredentials] = namespaceSettings.namespaceCredentialsMap
 
@@ -72,6 +75,10 @@ trait NamespacesHandler {
               logger.info("bucket {} in namespace {} return 403 for credentials {} so bucket exist but the credentials cannot see location ", bucketName.name, namespaceName, ex, credentials.getAWSAccessKeyId)
               return true
             }
+            if (storageS3Settings.slowdownCodes contains ex.asInstanceOf[AmazonS3Exception].getStatusCode) {
+              logger.info("throttling, cannot check bucket location")
+              throw new RokkuThrottlingException("cannot check bucket location")
+            }
             if (ex.asInstanceOf[AmazonS3Exception].getStatusCode != 404) {
               logger.error("namespace {} returned exception {} for credentials {} but should only status code 404", namespaceName, ex, credentials.getAWSAccessKeyId)
             }