forked from opensearch-project/opensearch-spark
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add rate limiter for bulk request (opensearch-project#567)
Signed-off-by: Tomoyuki Morita <[email protected]>
- Loading branch information
Showing
10 changed files
with
159 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
30 changes: 30 additions & 0 deletions
30
flint-core/src/main/scala/org/opensearch/flint/core/storage/BulkRequestRateLimiter.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,30 @@ | ||
package org.opensearch.flint.core.storage; | ||
|
||
import dev.failsafe.RateLimiter; | ||
import java.time.Duration; | ||
import java.util.logging.Logger; | ||
import org.opensearch.flint.core.FlintOptions; | ||
|
||
public class BulkRequestRateLimiter { | ||
private static final Logger LOG = Logger.getLogger(BulkRequestRateLimiter.class.getName()); | ||
private RateLimiter<Void> rateLimiter; | ||
|
||
public BulkRequestRateLimiter(FlintOptions flintOptions) { | ||
long bulkRequestRateLimitPerNode = flintOptions.getBulkRequestRateLimitPerNode(); | ||
if (bulkRequestRateLimitPerNode > 0) { | ||
LOG.info("Setting rate limit for bulk request to " + bulkRequestRateLimitPerNode + "/sec"); | ||
this.rateLimiter = RateLimiter.<Void>smoothBuilder( | ||
flintOptions.getBulkRequestRateLimitPerNode(), | ||
Duration.ofSeconds(1)).build(); | ||
} else { | ||
LOG.info("Rate limit for bulk request was not set."); | ||
} | ||
} | ||
|
||
// Wait so it won't exceed rate limit. Does nothing if rate limit is not set. | ||
public void acquirePermit() throws InterruptedException { | ||
if (rateLimiter != null) { | ||
this.rateLimiter.acquirePermit(); | ||
} | ||
} | ||
} |
22 changes: 22 additions & 0 deletions
22
...t-core/src/main/scala/org/opensearch/flint/core/storage/BulkRequestRateLimiterHolder.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,22 @@ | ||
package org.opensearch.flint.core.storage; | ||
|
||
import org.opensearch.flint.core.FlintOptions; | ||
|
||
/** | ||
* Hold shared instance of BulkRequestRateLimiter. This class is introduced to make | ||
* BulkRequestRateLimiter testable and share single instance. | ||
*/ | ||
public class BulkRequestRateLimiterHolder { | ||
|
||
private static BulkRequestRateLimiter instance; | ||
|
||
private BulkRequestRateLimiterHolder() {} | ||
|
||
public synchronized static BulkRequestRateLimiter getBulkRequestRateLimiter( | ||
FlintOptions flintOptions) { | ||
if (instance == null) { | ||
instance = new BulkRequestRateLimiter(flintOptions); | ||
} | ||
return instance; | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
20 changes: 20 additions & 0 deletions
20
...re/src/test/scala/org/opensearch/flint/core/storage/BulkRequestRateLimiterHolderTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,20 @@ | ||
package org.opensearch.flint.core.storage; | ||
|
||
import static org.junit.jupiter.api.Assertions.assertEquals; | ||
import static org.junit.jupiter.api.Assertions.assertNotNull; | ||
|
||
import com.google.common.collect.ImmutableMap; | ||
import org.junit.jupiter.api.Test; | ||
import org.opensearch.flint.core.FlintOptions; | ||
|
||
class BulkRequestRateLimiterHolderTest { | ||
FlintOptions flintOptions = new FlintOptions(ImmutableMap.of()); | ||
@Test | ||
public void getBulkRequestRateLimiter() { | ||
BulkRequestRateLimiter instance0 = BulkRequestRateLimiterHolder.getBulkRequestRateLimiter(flintOptions); | ||
BulkRequestRateLimiter instance1 = BulkRequestRateLimiterHolder.getBulkRequestRateLimiter(flintOptions); | ||
|
||
assertNotNull(instance0); | ||
assertEquals(instance0, instance1); | ||
} | ||
} |
46 changes: 46 additions & 0 deletions
46
flint-core/src/test/scala/org/opensearch/flint/core/storage/BulkRequestRateLimiterTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,46 @@ | ||
package org.opensearch.flint.core.storage; | ||
|
||
|
||
import static org.junit.jupiter.api.Assertions.assertEquals; | ||
import static org.junit.jupiter.api.Assertions.assertNotNull; | ||
import static org.junit.jupiter.api.Assertions.assertTrue; | ||
|
||
import com.google.common.collect.ImmutableMap; | ||
import org.junit.jupiter.api.Test; | ||
import org.opensearch.flint.core.FlintOptions; | ||
|
||
class BulkRequestRateLimiterTest { | ||
FlintOptions flintOptionsWithRateLimit = new FlintOptions(ImmutableMap.of(FlintOptions.BULK_REQUEST_RATE_LIMIT_PER_NODE, "1")); | ||
FlintOptions flintOptionsWithoutRateLimit = new FlintOptions(ImmutableMap.of(FlintOptions.BULK_REQUEST_RATE_LIMIT_PER_NODE, "0")); | ||
|
||
@Test | ||
void acquirePermitWithRateConfig() throws Exception { | ||
BulkRequestRateLimiter limiter = new BulkRequestRateLimiter(flintOptionsWithRateLimit); | ||
|
||
assertTrue(timer(() -> { | ||
limiter.acquirePermit(); | ||
limiter.acquirePermit(); | ||
}) >= 1000); | ||
} | ||
|
||
@Test | ||
void acquirePermitWithoutRateConfig() throws Exception { | ||
BulkRequestRateLimiter limiter = new BulkRequestRateLimiter(flintOptionsWithoutRateLimit); | ||
|
||
assertTrue(timer(() -> { | ||
limiter.acquirePermit(); | ||
limiter.acquirePermit(); | ||
}) < 100); | ||
} | ||
|
||
private interface Procedure { | ||
void run() throws Exception; | ||
} | ||
|
||
private long timer(Procedure procedure) throws Exception { | ||
long start = System.currentTimeMillis(); | ||
procedure.run(); | ||
long end = System.currentTimeMillis(); | ||
return end - start; | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters