Skip to content
This repository was archived by the owner on Jul 13, 2026. It is now read-only.

Commit da39bf4

Browse files
FabianMeiswinkelFabianMeiswinkel
authored
Merge pull request #62 from TheovanKraay/throughput-control-sample
Throughput control sample
2 parents a946084 + b54aa0a commit da39bf4

1 file changed

Lines changed: 327 additions & 0 deletions

File tree

‎src/main/java/com/azure/cosmos/examples/throughputcontrol/async/ThroughputControlQuickstartAsync.java‎

Lines changed: 327 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,327 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.azure.cosmos.examples.throughputcontrol.async;
5+
6+
import com.azure.cosmos.ConsistencyLevel;
7+
import com.azure.cosmos.CosmosAsyncClient;
8+
import com.azure.cosmos.CosmosAsyncContainer;
9+
import com.azure.cosmos.CosmosAsyncDatabase;
10+
import com.azure.cosmos.CosmosClientBuilder;
11+
import com.azure.cosmos.GlobalThroughputControlConfig;
12+
import com.azure.cosmos.ThrottlingRetryOptions;
13+
import com.azure.cosmos.ThroughputControlGroupConfig;
14+
import com.azure.cosmos.ThroughputControlGroupConfigBuilder;
15+
import com.azure.cosmos.examples.common.AccountSettings;
16+
import com.azure.cosmos.models.CosmosContainerProperties;
17+
import com.azure.cosmos.models.CosmosItemRequestOptions;
18+
import com.azure.cosmos.models.PriorityLevel;
19+
import com.azure.cosmos.models.ThroughputProperties;
20+
import com.fasterxml.jackson.databind.JsonNode;
21+
import org.slf4j.Logger;
22+
import org.slf4j.LoggerFactory;
23+
import reactor.core.publisher.Flux;
24+
import reactor.core.publisher.Mono;
25+
26+
import java.time.Duration;
27+
import java.util.ArrayList;
28+
import java.util.Arrays;
29+
import java.util.List;
30+
import java.util.concurrent.atomic.AtomicInteger;
31+
32+
import static com.azure.cosmos.examples.common.Profile.generateDocs;
33+
34+
public class ThroughputControlQuickstartAsync {
35+
private static CosmosAsyncClient client1;
36+
private static CosmosAsyncClient client2;
37+
int random_int = (int) Math.floor(Math.random() * (100 - 50 + 1) + 50);
38+
private final String databaseName = "ThroughputControlDemoDB" + random_int;
39+
private final String throughputControlContainerName = "ThroughputControl";
40+
private final String priorityBasedThrottlingContainerName = "priorityBasedThrottlingContainer";
41+
private final String containerName = "CosmosContainer";
42+
private CosmosAsyncDatabase database1;
43+
private CosmosAsyncDatabase database2;
44+
private CosmosAsyncContainer ThroughputControlTestContainerObject1;
45+
private CosmosAsyncContainer ThroughputControlTestContainerObject2;
46+
private CosmosAsyncContainer priorityBasedThrottlingContainerObject1;
47+
private CosmosAsyncContainer priorityBasedThrottlingContainerObject2;
48+
private static AtomicInteger request_count = new AtomicInteger(1);
49+
private static AtomicInteger rate_limit_error_count = new AtomicInteger(0);
50+
//set max number of retries for 429 (throttling) to low number so that we see rate limiting errors
51+
ThrottlingRetryOptions retryOptions = new ThrottlingRetryOptions().setMaxRetryAttemptsOnThrottledRequests(1);
52+
public static final int PROVISIONED_RUS = 10000;
53+
public static final int THROUGHPUT_CONTROL_RUS = 200;
54+
public static final int NUMBER_OF_DOCS = 2000;
55+
public static final int NUMBER_OF_DOCS_PRIORITY_BASED_THROTTLING = 100;
56+
57+
public ArrayList<JsonNode> docs;
58+
CosmosItemRequestOptions options = new CosmosItemRequestOptions();
59+
private final static Logger logger = LoggerFactory.getLogger(ThroughputControlQuickstartAsync.class);
60+
public void closeClient1() {
61+
client1.close();
62+
}
63+
public void closeClient2() {
64+
client2.close();
65+
}
66+
67+
/**
68+
* Demo throughput control...
69+
*/
70+
// <Main>
71+
public static void main(String[] args) {
72+
ThroughputControlQuickstartAsync p = new ThroughputControlQuickstartAsync();
73+
74+
try {
75+
p.throughControlDemo();
76+
logger.info("Demo complete, please hold while resources are released");
77+
} catch (Exception e) {
78+
logger.info("Cosmos getStarted failed with: " + e);
79+
} finally {
80+
logger.info("Closing clients");
81+
p.closeClient1();
82+
p.closeClient2();
83+
}
84+
System.exit(0);
85+
}
86+
// </Main>
87+
88+
private void throughControlDemo() throws Exception {
89+
logger.info("Using Azure Cosmos DB endpoint: " + AccountSettings.HOST);
90+
91+
client1 = new CosmosClientBuilder()
92+
.endpoint(AccountSettings.HOST)
93+
.key(AccountSettings.MASTER_KEY)
94+
.consistencyLevel(ConsistencyLevel.SESSION)
95+
.contentResponseOnWriteEnabled(true)
96+
.directMode()
97+
.throttlingRetryOptions(retryOptions)
98+
.buildAsyncClient();
99+
100+
client2 = new CosmosClientBuilder()
101+
.endpoint(AccountSettings.HOST)
102+
.key(AccountSettings.MASTER_KEY)
103+
.consistencyLevel(ConsistencyLevel.SESSION)
104+
.contentResponseOnWriteEnabled(true)
105+
.directMode()
106+
.throttlingRetryOptions(retryOptions)
107+
.buildAsyncClient();
108+
109+
Mono<Void> databaseContainerIfNotExist = client1.createDatabaseIfNotExists(databaseName)
110+
.flatMap(databaseResponse -> {
111+
database1 = client1.getDatabase(databaseResponse.getProperties().getId());
112+
logger.info("Created database" + databaseName);
113+
CosmosContainerProperties containerProperties = new CosmosContainerProperties(containerName, "/id");
114+
ThroughputProperties throughputProperties = ThroughputProperties.createManualThroughput(PROVISIONED_RUS);
115+
return database1.createContainerIfNotExists(containerProperties, throughputProperties);
116+
}).flatMap(containerResponse -> {
117+
ThroughputControlTestContainerObject1 = database1.getContainer(containerResponse.getProperties().getId());
118+
logger.info("Created container " + containerName);
119+
return Mono.empty();
120+
});
121+
122+
logger.info("Creating database and container asynchronously...");
123+
databaseContainerIfNotExist.block();
124+
125+
logger.info("Creating throughput control container to manage global throughput control metadata...");
126+
//<CreateThroughputControlContainer>
127+
database1 = client1.getDatabase(databaseName);
128+
//NOTE: this container is not subject to throughput control, as it is used to manage throughput control
129+
//NOTE: throughput control container MUST be created with partition key of /groupId and TTL must be set.
130+
CosmosContainerProperties throughputContainerProperties = new CosmosContainerProperties(throughputControlContainerName, "/groupId").setDefaultTimeToLiveInSeconds(-1);
131+
ThroughputProperties throughputProperties = ThroughputProperties.createManualThroughput(PROVISIONED_RUS);
132+
database1.createContainerIfNotExists(throughputContainerProperties, throughputProperties).block();
133+
//</CreateThroughputControlContainer>
134+
135+
logger.info("Creating container for priority based throttling test, with only 400 RUs...");
136+
CosmosContainerProperties priorityBasedThrottlingContainerProperties = new CosmosContainerProperties(priorityBasedThrottlingContainerName, "/id").setDefaultTimeToLiveInSeconds(-1);
137+
ThroughputProperties priorityBasedThrottlingContainerThroughputProperties = ThroughputProperties.createManualThroughput(400);
138+
database1.createContainerIfNotExists(priorityBasedThrottlingContainerProperties, priorityBasedThrottlingContainerThroughputProperties).block();
139+
140+
logger.info("Running a base test which should not produce rate limiting...");
141+
baseTest();
142+
logger.info("Running a local throughput control test with 2 clients loading " + NUMBER_OF_DOCS + " docs and having " + THROUGHPUT_CONTROL_RUS + " RU limit each...");
143+
localThroughputControlTest();
144+
logger.info("Running a global throughput control test with 2 clients loading " + NUMBER_OF_DOCS + " docs and sharing " + THROUGHPUT_CONTROL_RUS + " RU limit...");
145+
//this should produce more rate limiting than the local throughput control above, as both clients share the same RU limit in the group
146+
//search for isThroughputControlRequestRateTooLarge in the logs to see the rate limiting from throughput control.
147+
globalThroughputControlTest();
148+
logger.info("Running priority based throttling test with two clients loading " + NUMBER_OF_DOCS_PRIORITY_BASED_THROTTLING + " docs");
149+
//one client has priority over the other, so it should be throttled less
150+
priorityBasedThrottling();
151+
database1.delete().block();
152+
}
153+
154+
private void baseTest() {
155+
try {
156+
docs = generateDocs(NUMBER_OF_DOCS);
157+
createManyItems("BASE TEST", docs, options);
158+
} catch (Exception e) {
159+
throw new RuntimeException(e);
160+
}
161+
}
162+
163+
/**
164+
* Local throughput control test - we will create a local throughput control group
165+
* for each client, and each client will have a limit of 200 RU/s
166+
*/
167+
private void localThroughputControlTest() {
168+
database2 = client2.getDatabase(databaseName);
169+
ThroughputControlTestContainerObject2 = database2.getContainer(containerName);
170+
options.setThroughputControlGroupName("localControlGroup");
171+
ThroughputControlGroupConfig groupConfig1 =
172+
new ThroughputControlGroupConfigBuilder()
173+
.groupName("localControlGroup")
174+
// use below to set a target throughput threshold as % of provisioned throughput instead of absolute value
175+
//.targetThroughputThreshold(0.25)
176+
.targetThroughput(THROUGHPUT_CONTROL_RUS)
177+
.build();
178+
ThroughputControlTestContainerObject1.enableLocalThroughputControlGroup(groupConfig1);
179+
ThroughputControlGroupConfig groupConfig2 =
180+
new ThroughputControlGroupConfigBuilder()
181+
.groupName("localControlGroup")
182+
// use below to set a target throughput threshold as % of provisioned throughput instead of absolute value
183+
//.targetThroughputThreshold(0.25)
184+
.targetThroughput(THROUGHPUT_CONTROL_RUS)
185+
.build();
186+
ThroughputControlTestContainerObject2.enableLocalThroughputControlGroup(groupConfig2);
187+
try {
188+
//createManyItems("loacl throughput test", docs, options);
189+
createManyItemsWithTwoClients(NUMBER_OF_DOCS, "LOCAL THROUGHPUT CONTROL TEST", options, Arrays.asList(ThroughputControlTestContainerObject1, ThroughputControlTestContainerObject2));
190+
} catch (Exception e) {
191+
logger.info("Exception in localThroughputControlTest: " + e);
192+
}
193+
}
194+
195+
/**
196+
* Global throughput control test - we will create a global throughput control group
197+
* that will be shared by both clients - they will share a limit of 200 RU/s
198+
*/
199+
private void globalThroughputControlTest() {
200+
database2 = client2.getDatabase(databaseName);
201+
ThroughputControlTestContainerObject2 = database2.getContainer(containerName);
202+
CosmosItemRequestOptions options = new CosmosItemRequestOptions();
203+
options.setThroughputControlGroupName("globalControlGroup");
204+
ThroughputControlGroupConfig groupConfig =
205+
new ThroughputControlGroupConfigBuilder()
206+
.groupName("globalControlGroup")
207+
// use below to set a target throughput threshold as % of provisioned throughput instead of absolute value
208+
//.targetThroughputThreshold(0.25)
209+
.targetThroughput(THROUGHPUT_CONTROL_RUS)
210+
.build();
211+
GlobalThroughputControlConfig globalControlConfig1 =
212+
this.client1.createGlobalThroughputControlConfigBuilder(database1.getId(), throughputControlContainerName)
213+
.setControlItemRenewInterval(Duration.ofSeconds(5))
214+
.setControlItemExpireInterval(Duration.ofSeconds(20))
215+
.build();
216+
GlobalThroughputControlConfig globalControlConfig2 =
217+
this.client2.createGlobalThroughputControlConfigBuilder(database1.getId(), throughputControlContainerName)
218+
.setControlItemRenewInterval(Duration.ofSeconds(5))
219+
.setControlItemExpireInterval(Duration.ofSeconds(20))
220+
.build();
221+
ThroughputControlTestContainerObject1.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig1);
222+
ThroughputControlTestContainerObject2.enableGlobalThroughputControlGroup(groupConfig, globalControlConfig2);
223+
try {
224+
createManyItemsWithTwoClients(NUMBER_OF_DOCS, "GLOBAL THROUGHPUT CONTROL TEST", options, Arrays.asList(ThroughputControlTestContainerObject1, ThroughputControlTestContainerObject2));
225+
} catch (Exception e) {
226+
logger.info("Exception in globalThroughputControlTest: " + e);
227+
}
228+
}
229+
230+
/**
231+
* Priority based throttling test - one client will have priority of LOW, the other HIGH
232+
*/
233+
private void priorityBasedThrottling() {
234+
database1 = client1.getDatabase(databaseName);
235+
database2 = client2.getDatabase(databaseName);
236+
priorityBasedThrottlingContainerObject1 = database1.getContainer(priorityBasedThrottlingContainerName);
237+
priorityBasedThrottlingContainerObject2 = database2.getContainer(priorityBasedThrottlingContainerName);
238+
CosmosItemRequestOptions options = new CosmosItemRequestOptions();
239+
options.setThroughputControlGroupName("priorityBasedThrottling");
240+
ThroughputControlGroupConfig groupConfig1 =
241+
new ThroughputControlGroupConfigBuilder()
242+
.groupName("priorityBasedThrottling")
243+
.priorityLevel(PriorityLevel.HIGH)
244+
.build();
245+
ThroughputControlGroupConfig groupConfig2 =
246+
new ThroughputControlGroupConfigBuilder()
247+
.groupName("priorityBasedThrottling")
248+
.priorityLevel(PriorityLevel.LOW)
249+
.build();
250+
ThroughputControlTestContainerObject1.enableLocalThroughputControlGroup(groupConfig1);
251+
ThroughputControlTestContainerObject2.enableLocalThroughputControlGroup(groupConfig2);
252+
try {
253+
createManyItemsWithTwoClients(NUMBER_OF_DOCS_PRIORITY_BASED_THROTTLING, "PRIORITY BASED THROTTLING TEST", options, Arrays.asList(priorityBasedThrottlingContainerObject1, priorityBasedThrottlingContainerObject2));
254+
Thread.sleep(2000);
255+
} catch (Exception e) {
256+
logger.info("Exception in globalThroughputControlTest: " + e);
257+
}
258+
259+
}
260+
261+
262+
private void createManyItems(String test, ArrayList<JsonNode> docs, CosmosItemRequestOptions options) throws Exception {
263+
Flux.fromIterable(docs).flatMap(doc -> ThroughputControlTestContainerObject1.createItem(doc, options)).flatMap(itemResponse -> {
264+
if (itemResponse.getStatusCode() == 201) {
265+
request_count.incrementAndGet();
266+
} else {
267+
logger.info("WARNING insert status code {} != 201" + itemResponse.getStatusCode());
268+
request_count.incrementAndGet();
269+
}
270+
return Mono.empty();
271+
}).doOnError((exception) -> {
272+
request_count.incrementAndGet();
273+
rate_limit_error_count.incrementAndGet();
274+
logger.info(
275+
"error creating item in " + test + " e: {}",
276+
exception.getLocalizedMessage(),
277+
exception);
278+
})
279+
.blockLast(); // block to wait for all items to be created
280+
logger.info("total request count was: " + request_count.get());
281+
logger.info("\n\n\n****************\n***TOTAL NUMBER OF RATE LIMIT ERRORS RECORDED IN " + test + " was: " + rate_limit_error_count.get()+"\n***************\n\n\n");
282+
request_count.set(0);
283+
rate_limit_error_count.set(0);
284+
}
285+
286+
private void createManyItemsWithTwoClients(int noOfDocs, String test, CosmosItemRequestOptions options, List<CosmosAsyncContainer> containers) throws Exception {
287+
int clientId = 1;
288+
for (CosmosAsyncContainer cosmosAsyncContainer : containers) {
289+
logger.info("client " + clientId + " of " + test);
290+
docs = generateDocs(noOfDocs);
291+
int finalClientId = clientId;
292+
Flux.fromIterable(docs).flatMap(doc -> cosmosAsyncContainer.createItem(doc, options)).flatMap(itemResponse -> {
293+
if (itemResponse.getStatusCode() == 201) {
294+
//uncomment below to see diagnostics in logs showing retries and isThroughputControlRequestRateTooLarge value
295+
//logger.info("printing diagnostics to see retries isThroughputControlRequestRateTooLarge value" + itemResponse.getDiagnostics());
296+
} else {
297+
logger.info("WARNING insert status code {} != 201" + itemResponse.getStatusCode());
298+
}
299+
request_count.incrementAndGet();
300+
return Mono.empty();
301+
}).doOnError((exception) -> {
302+
request_count.incrementAndGet();
303+
rate_limit_error_count.incrementAndGet();
304+
logger.info(
305+
"error creating item in " + test + " from client number " + finalClientId + " e: {}",
306+
exception.getLocalizedMessage(),
307+
exception);
308+
}).subscribe();
309+
clientId++;
310+
}
311+
int currentRequestCount = 0;
312+
//when request count stops increasing, we can assume that all requests and retries have been exhausted
313+
while (request_count.get() < noOfDocs * 2) {
314+
if (currentRequestCount == request_count.get()) {
315+
logger.info("request count has not increased in last 1 second, current request count is: " + request_count.get());
316+
break;
317+
}
318+
Thread.sleep(1000);
319+
currentRequestCount = request_count.get();
320+
}
321+
logger.info("total request count was: " + request_count.get());
322+
logger.info("\n\n\n****************\n***TOTAL NUMBER OF RATE LIMIT ERRORS RECORDED IN " + test + " was: " + rate_limit_error_count.get()+"\n***************\n\n\n");
323+
request_count.set(0);
324+
rate_limit_error_count.set(0);
325+
}
326+
327+
}

0 commit comments

Comments
 (0)