Skip to content

Commit d80be09

Browse files
sumitasrsumitasr
committed
Perform code refactor
Signed-off-by: Sumit Bansal <[email protected]>
1 parent 27b68c5 commit d80be09

File tree

2 files changed

+103
-38
lines changed

2 files changed

+103
-38
lines changed

‎server/src/main/java/org/opensearch/cluster/service/ClusterManagerTaskThrottler.java

Lines changed: 34 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
* <p>
3434
* Set specific setting to for setting the threshold of throttling of particular task type.
3535
* e.g : Set "cluster_manager.throttling.thresholds.put_mapping" to set throttling limit of "put mapping" tasks,
36-
* Set it to default value(-1) to disable the throttling for this task type.
36+
* Set it to default value(-1) to disable the throttling for this task type.
3737
*/
3838
public class ClusterManagerTaskThrottler implements TaskBatcherListener {
3939
private static final Logger logger = LogManager.getLogger(ClusterManagerTaskThrottler.class);
@@ -69,7 +69,7 @@ public class ClusterManagerTaskThrottler implements TaskBatcherListener {
6969
private final int MIN_THRESHOLD_VALUE = -1; // Disabled throttling
7070
private final ClusterManagerTaskThrottlerListener clusterManagerTaskThrottlerListener;
7171

72-
private final ConcurrentMap<String, Long> tasksCount;
72+
final ConcurrentMap<String, Long> tasksCount;
7373
private final ConcurrentMap<String, Long> tasksThreshold;
7474
private final Supplier<Version> minNodeVersionSupplier;
7575

@@ -210,23 +210,14 @@ Long getThrottlingLimit(final String taskKey) {
210210
}
211211

212212
private void checkForClusterManagerThrottling(
213-
final ThrottlingKey clusterManagerThrottlingKey,
214-
final String taskThrottlingKey,
213+
final boolean throttlingEnabledWithThreshold,
214+
final Long threshold,
215215
final long taskCount,
216-
final int tasksSize
216+
final int tasksSize,
217+
final String taskThrottlingKey
217218
) {
218-
if (clusterManagerThrottlingKey.isThrottlingEnabled()) {
219-
Long threshold = tasksThreshold.get(taskThrottlingKey);
220-
if (threshold != null && shouldThrottle(threshold, taskCount, tasksSize)) {
221-
clusterManagerTaskThrottlerListener.onThrottle(taskThrottlingKey, tasksSize);
222-
logger.warn(
223-
"Throwing Throttling Exception for [{}]. Trying to add [{}] tasks to queue, limit is set to [{}]",
224-
taskThrottlingKey,
225-
tasksSize,
226-
threshold
227-
);
228-
throw new ClusterManagerThrottlingException("Throttling Exception : Limit exceeded for " + taskThrottlingKey);
229-
}
219+
if (throttlingEnabledWithThreshold && shouldThrottle(threshold, taskCount, tasksSize)) {
220+
throw new ClusterManagerThrottlingException("Throttling Exception : Limit exceeded for " + taskThrottlingKey);
230221
}
231222
}
232223

@@ -235,16 +226,33 @@ public void onBeginSubmit(List<? extends TaskBatcher.BatchedTask> tasks) {
235226
final ThrottlingKey clusterManagerThrottlingKey = ((ClusterStateTaskExecutor<Object>) tasks.get(0).batchingKey)
236227
.getClusterManagerThrottlingKey();
237228
final String taskThrottlingKey = clusterManagerThrottlingKey.getTaskThrottlingKey();
229+
final Long threshold = getThrottlingLimit(taskThrottlingKey);
230+
final boolean isThrottlingEnabledWithThreshold = clusterManagerThrottlingKey.isThrottlingEnabled() && threshold != null;
238231
tasksCount.putIfAbsent(taskThrottlingKey, 0L);
239-
240-
// Performing shallow check before taking lock, performing throttle check and computing new count
241-
checkForClusterManagerThrottling(clusterManagerThrottlingKey, taskThrottlingKey, tasksCount.get(taskThrottlingKey), tasks.size());
242-
243-
tasksCount.computeIfPresent(taskThrottlingKey, (key, count) -> {
244-
int size = tasks.size();
245-
checkForClusterManagerThrottling(clusterManagerThrottlingKey, taskThrottlingKey, count, size);
246-
return count + size;
247-
});
232+
int tasksSize = tasks.size();
233+
234+
try {
235+
checkForClusterManagerThrottling(
236+
isThrottlingEnabledWithThreshold,
237+
threshold,
238+
tasksCount.get(taskThrottlingKey),
239+
tasksSize,
240+
taskThrottlingKey
241+
);
242+
tasksCount.computeIfPresent(taskThrottlingKey, (key, count) -> {
243+
checkForClusterManagerThrottling(isThrottlingEnabledWithThreshold, threshold, count, tasksSize, taskThrottlingKey);
244+
return count + tasksSize;
245+
});
246+
} catch (final ClusterManagerThrottlingException e) {
247+
clusterManagerTaskThrottlerListener.onThrottle(taskThrottlingKey, tasksSize);
248+
logger.trace(
249+
"Throwing Throttling Exception for [{}]. Trying to add [{}] tasks to queue, limit is set to [{}]",
250+
taskThrottlingKey,
251+
tasksSize,
252+
threshold
253+
);
254+
throw e;
255+
}
248256
}
249257

250258
/**

‎server/src/test/java/org/opensearch/cluster/service/ClusterManagerTaskThrottlerTests.java

Lines changed: 69 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030
import java.util.ArrayList;
3131
import java.util.Collections;
3232
import java.util.List;
33+
import java.util.Optional;
34+
import java.util.concurrent.CountDownLatch;
3335
import java.util.concurrent.TimeUnit;
3436

3537
import static org.opensearch.test.ClusterServiceUtils.setState;
@@ -69,7 +71,7 @@ public static void afterClass() {
6971
public void testDefaults() {
7072
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
7173
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
72-
return clusterService.getMasterService().getMinNodeVersion();
74+
return clusterService.getClusterManagerService().getMinNodeVersion();
7375
}, new ClusterManagerThrottlingStats());
7476
throttler.registerClusterManagerTask("put-mapping", true);
7577
throttler.registerClusterManagerTask("create-index", true);
@@ -108,7 +110,7 @@ public void testValidateSettingsForDifferentVersion() {
108110
}
109111
}
110112

111-
public void testValidateSettingsForTaskWihtoutRetryOnDataNode() {
113+
public void testValidateSettingsForTaskWithoutRetryOnDataNode() {
112114
DiscoveryNode clusterManagerNode = getClusterManagerNode(Version.V_2_5_0);
113115
DiscoveryNode dataNode = getDataNode(Version.V_2_5_0);
114116
setState(
@@ -139,7 +141,7 @@ public void testUpdateSettingsForNullValue() {
139141

140142
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
141143
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
142-
return clusterService.getMasterService().getMinNodeVersion();
144+
return clusterService.getClusterManagerService().getMinNodeVersion();
143145
}, new ClusterManagerThrottlingStats());
144146
throttler.registerClusterManagerTask("put-mapping", true);
145147

@@ -173,7 +175,7 @@ public void testSettingsOnBootstrap() {
173175
.put("cluster_manager.throttling.retry.max.delay", maxDelay + "s")
174176
.build();
175177
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(initialSettings, clusterSettings, () -> {
176-
return clusterService.getMasterService().getMinNodeVersion();
178+
return clusterService.getClusterManagerService().getMinNodeVersion();
177179
}, new ClusterManagerThrottlingStats());
178180
throttler.registerClusterManagerTask("put-mapping", true);
179181

@@ -187,7 +189,7 @@ public void testSettingsOnBootstrap() {
187189
public void testUpdateRetryDelaySetting() {
188190
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
189191
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
190-
return clusterService.getMasterService().getMinNodeVersion();
192+
return clusterService.getClusterManagerService().getMinNodeVersion();
191193
}, new ClusterManagerThrottlingStats());
192194

193195
// verify defaults
@@ -217,7 +219,7 @@ public void testValidateSettingsForUnknownTask() {
217219

218220
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
219221
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
220-
return clusterService.getMasterService().getMinNodeVersion();
222+
return clusterService.getClusterManagerService().getMinNodeVersion();
221223
}, new ClusterManagerThrottlingStats());
222224

223225
// set some limit for update snapshot tasks
@@ -236,7 +238,7 @@ public void testUpdateThrottlingLimitForBasicSanity() {
236238

237239
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
238240
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
239-
return clusterService.getMasterService().getMinNodeVersion();
241+
return clusterService.getClusterManagerService().getMinNodeVersion();
240242
}, new ClusterManagerThrottlingStats());
241243
throttler.registerClusterManagerTask("put-mapping", true);
242244

@@ -263,7 +265,7 @@ public void testValidateSettingForLimit() {
263265

264266
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
265267
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
266-
return clusterService.getMasterService().getMinNodeVersion();
268+
return clusterService.getClusterManagerService().getMinNodeVersion();
267269
}, new ClusterManagerThrottlingStats());
268270
throttler.registerClusterManagerTask("put-mapping", true);
269271

@@ -274,7 +276,7 @@ public void testValidateSettingForLimit() {
274276
public void testUpdateLimit() {
275277
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
276278
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
277-
return clusterService.getMasterService().getMinNodeVersion();
279+
return clusterService.getClusterManagerService().getMinNodeVersion();
278280
}, new ClusterManagerThrottlingStats());
279281
throttler.registerClusterManagerTask("put-mapping", true);
280282

@@ -309,7 +311,7 @@ public void testThrottlingForDisabledThrottlingTask() {
309311
String taskKey = "test";
310312
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
311313
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
312-
return clusterService.getMasterService().getMinNodeVersion();
314+
return clusterService.getClusterManagerService().getMinNodeVersion();
313315
}, throttlingStats);
314316
ClusterManagerTaskThrottler.ThrottlingKey throttlingKey = throttler.registerClusterManagerTask(taskKey, false);
315317

@@ -321,6 +323,9 @@ public void testThrottlingForDisabledThrottlingTask() {
321323

322324
// Asserting that there was not any throttling for it
323325
assertEquals(0L, throttlingStats.getThrottlingCount(taskKey));
326+
327+
// Asserting value in tasksCount map to make sure it gets updated even when throttling is disabled
328+
assertEquals(Optional.of(10L).get(), throttler.tasksCount.get(taskKey));
324329
}
325330

326331
public void testThrottlingForInitialStaticSettingAndVersionCheck() {
@@ -339,7 +344,7 @@ public void testThrottlingForInitialStaticSettingAndVersionCheck() {
339344
.put("cluster_manager.throttling.thresholds.put-mapping.value", put_mapping_threshold_value)
340345
.build();
341346
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(initialSettings, clusterSettings, () -> {
342-
return clusterService.getMasterService().getMinNodeVersion();
347+
return clusterService.getClusterManagerService().getMinNodeVersion();
343348
}, throttlingStats);
344349
ClusterManagerTaskThrottler.ThrottlingKey throttlingKey = throttler.registerClusterManagerTask("put-mapping", true);
345350

@@ -367,7 +372,7 @@ public void testThrottling() {
367372
String taskKey = "test";
368373
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
369374
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
370-
return clusterService.getMasterService().getMinNodeVersion();
375+
return clusterService.getClusterManagerService().getMinNodeVersion();
371376
}, throttlingStats);
372377
ClusterManagerTaskThrottler.ThrottlingKey throttlingKey = throttler.registerClusterManagerTask(taskKey, true);
373378

@@ -406,6 +411,58 @@ public void testThrottling() {
406411
throttler.onBeginSubmit(getMockUpdateTaskList(taskKey, throttlingKey, 1));
407412
}
408413

414+
public void testThrottlingWithLock() {
415+
ClusterManagerThrottlingStats throttlingStats = new ClusterManagerThrottlingStats();
416+
String taskKey = "test";
417+
ClusterSettings clusterSettings = new ClusterSettings(Settings.builder().build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS);
418+
ClusterManagerTaskThrottler throttler = new ClusterManagerTaskThrottler(Settings.EMPTY, clusterSettings, () -> {
419+
return clusterService.getClusterManagerService().getMinNodeVersion();
420+
}, throttlingStats);
421+
ClusterManagerTaskThrottler.ThrottlingKey throttlingKey = throttler.registerClusterManagerTask(taskKey, true);
422+
423+
throttler.updateLimit(taskKey, 5);
424+
425+
// adding 3 tasks
426+
throttler.onBeginSubmit(getMockUpdateTaskList(taskKey, throttlingKey, 3));
427+
428+
// adding 3 more tasks, these tasks should be throttled
429+
// taskCount in Queue: 3 Threshold: 5
430+
assertThrows(
431+
ClusterManagerThrottlingException.class,
432+
() -> throttler.onBeginSubmit(getMockUpdateTaskList(taskKey, throttlingKey, 3))
433+
);
434+
assertEquals(3L, throttlingStats.getThrottlingCount(taskKey));
435+
436+
// remove one task
437+
throttler.onBeginProcessing(getMockUpdateTaskList(taskKey, throttlingKey, 1));
438+
439+
// add 3 tasks should pass now.
440+
// taskCount in Queue: 2 Threshold: 5
441+
throttler.onBeginSubmit(getMockUpdateTaskList(taskKey, throttlingKey, 3));
442+
443+
final CountDownLatch latch = new CountDownLatch(1);
444+
445+
// Taking lock on tasksCount will not impact throttling behaviour now.
446+
new Thread(() -> {
447+
throttler.tasksCount.computeIfPresent(taskKey, (key, count) -> {
448+
try {
449+
latch.await();
450+
} catch (InterruptedException e) {
451+
throw new RuntimeException(e);
452+
}
453+
return 10L;
454+
});
455+
}).start();
456+
457+
// adding one task will throttle
458+
// taskCount in Queue: 5 Threshold: 5
459+
assertThrows(
460+
ClusterManagerThrottlingException.class,
461+
() -> throttler.onBeginSubmit(getMockUpdateTaskList(taskKey, throttlingKey, 1))
462+
);
463+
latch.countDown();
464+
}
465+
409466
private List<TaskBatcherTests.TestTaskBatcher.UpdateTask> getMockUpdateTaskList(
410467
String taskKey,
411468
ClusterManagerTaskThrottler.ThrottlingKey throttlingKey,

0 commit comments

Comments
 (0)