use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.
the class CountBasedLimiterTest method testThrottling.
@Test
public void testThrottling() throws InterruptedException {
Limiter limiter = new CountBasedLimiter(10);
limiter.start();
for (int i = 0; i < 10; i++) {
Assert.assertTrue(limiter.acquirePermits(1) != null);
}
Assert.assertTrue(limiter.acquirePermits(1) == null);
limiter.stop();
}
use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.
the class SharedLimiterFactoryTest method testMultiLevelLimiter.
@Test
public void testMultiLevelLimiter() throws Exception {
SharedResourcesBrokerImpl<SimpleScopeType> broker = getBrokerForConfigMap(ImmutableMap.of(JOINER.join(BrokerConstants.GOBBLIN_BROKER_CONFIG_PREFIX, SharedLimiterFactory.NAME, SimpleScopeType.GLOBAL, SharedLimiterFactory.LIMITER_CLASS_KEY), "CountBasedLimiter", JOINER.join(BrokerConstants.GOBBLIN_BROKER_CONFIG_PREFIX, SharedLimiterFactory.NAME, SimpleScopeType.GLOBAL, CountBasedLimiter.Factory.COUNT_KEY), "10", JOINER.join(BrokerConstants.GOBBLIN_BROKER_CONFIG_PREFIX, SharedLimiterFactory.NAME, SimpleScopeType.LOCAL, SharedLimiterFactory.LIMITER_CLASS_KEY), "CountBasedLimiter", JOINER.join(BrokerConstants.GOBBLIN_BROKER_CONFIG_PREFIX, SharedLimiterFactory.NAME, SimpleScopeType.LOCAL, CountBasedLimiter.Factory.COUNT_KEY), "5"));
SharedResourcesBroker<SimpleScopeType> localBroker1 = broker.newSubscopedBuilder(new SimpleScope<>(SimpleScopeType.LOCAL, "local1")).build();
SharedResourcesBroker<SimpleScopeType> localBroker2 = broker.newSubscopedBuilder(new SimpleScope<>(SimpleScopeType.LOCAL, "local2")).build();
SharedLimiterFactory<SimpleScopeType> factory = new SharedLimiterFactory<>();
Limiter limiter1 = ((ResourceInstance<Limiter>) factory.createResource(localBroker1, broker.getConfigView(SimpleScopeType.LOCAL, new SharedLimiterKey("resource"), factory.getName()))).getResource();
Limiter limiter2 = ((ResourceInstance<Limiter>) factory.createResource(localBroker2, broker.getConfigView(SimpleScopeType.LOCAL, new SharedLimiterKey("resource"), factory.getName()))).getResource();
Assert.assertTrue(limiter1 instanceof MultiLimiter);
Assert.assertTrue(limiter2 instanceof MultiLimiter);
Assert.assertEquals(((CountBasedLimiter) ((MultiLimiter) limiter1).getUnderlyingLimiters().get(0)).getCountLimit(), 5);
Assert.assertEquals(((CountBasedLimiter) ((MultiLimiter) limiter1).getUnderlyingLimiters().get(1)).getCountLimit(), 10);
Assert.assertEquals(((CountBasedLimiter) ((MultiLimiter) limiter2).getUnderlyingLimiters().get(0)).getCountLimit(), 5);
Assert.assertEquals(((CountBasedLimiter) ((MultiLimiter) limiter2).getUnderlyingLimiters().get(1)).getCountLimit(), 10);
Assert.assertNotEquals(((MultiLimiter) limiter1).getUnderlyingLimiters().get(0), ((MultiLimiter) limiter2).getUnderlyingLimiters().get(0));
Assert.assertEquals(((MultiLimiter) limiter1).getUnderlyingLimiters().get(1), ((MultiLimiter) limiter2).getUnderlyingLimiters().get(1));
}
use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.
the class LimiterStopEventTest method testGetLimiterStopMetadataCase1.
@Test
public void testGetLimiterStopMetadataCase1() throws InterruptedException {
Properties properties = new Properties();
String key1 = "topic";
String key2 = "partition.id";
String keyList = Joiner.on(',').join(key1, key2);
String subKey1 = key2 + ".0";
String subKey2 = key2 + ".1";
String subKey3 = key2 + ".2";
String subKey4 = key2 + ".3";
String subKey5 = "partition";
properties.setProperty(LimiterConfigurationKeys.LIMITER_REPORT_KEY_LIST, keyList);
properties.setProperty(subKey1, "1111");
properties.setProperty(subKey2, "1111");
properties.setProperty(subKey3, "1111");
properties.setProperty(subKey4, "1111");
Extractor extractor = mock(Extractor.class);
Limiter limiter = mock(Limiter.class);
TaskState taskState = mock(TaskState.class);
WorkUnit workUnit = mock(WorkUnit.class);
Mockito.when(taskState.getWorkunit()).thenReturn(workUnit);
Mockito.when(taskState.getJobId()).thenReturn("123");
Mockito.when(taskState.getTaskAttemptId()).thenReturn(Optional.of("555"));
Mockito.when(taskState.getTaskId()).thenReturn("888");
Mockito.when(limiter.acquirePermits(1)).thenReturn(null);
Mockito.when(taskState.getProp(ConfigurationKeys.DATASET_URN_KEY, ConfigurationKeys.DEFAULT_DATASET_URN)).thenReturn("file://xyz");
Mockito.when(workUnit.getProperties()).thenReturn(properties);
LimitingExtractorDecorator<String, String> decorator = new LimitingExtractorDecorator<>(extractor, limiter, taskState);
try {
Method method = LimitingExtractorDecorator.class.getDeclaredMethod("getLimiterStopMetadata");
method.setAccessible(true);
ImmutableMap<String, String> metaData = (ImmutableMap<String, String>) method.invoke(decorator);
Assert.assertEquals(metaData.containsKey(subKey1), true);
Assert.assertEquals(metaData.containsKey(subKey2), true);
Assert.assertEquals(metaData.containsKey(subKey3), true);
Assert.assertEquals(metaData.containsKey(subKey4), true);
Assert.assertEquals(metaData.containsKey(subKey5), false);
} catch (Exception e) {
Assert.fail();
}
}
use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.
the class TaskContext method getExtractor.
/**
* Get a {@link Extractor} instance.
*
* @return a {@link Extractor} instance
*/
public Extractor getExtractor() {
try {
this.rawSourceExtractor = getSource().getExtractor(this.taskState);
boolean throttlingEnabled = this.taskState.getPropAsBoolean(ConfigurationKeys.EXTRACT_LIMIT_ENABLED_KEY, ConfigurationKeys.DEFAULT_EXTRACT_LIMIT_ENABLED);
if (throttlingEnabled) {
Limiter limiter = DefaultLimiterFactory.newLimiter(this.taskState);
if (!(limiter instanceof NonRefillableLimiter)) {
throw new IllegalArgumentException("The Limiter used with an Extractor should be an instance of " + NonRefillableLimiter.class.getSimpleName());
}
return new LimitingExtractorDecorator<>(this.rawSourceExtractor, limiter, this.taskState);
}
return this.rawSourceExtractor;
} catch (IOException ioe) {
throw new RuntimeException(ioe);
}
}
use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.
the class SharedLimiterFactory method createResource.
@Override
public SharedResourceFactoryResponse<Limiter> createResource(SharedResourcesBroker<S> broker, ScopedConfigView<S, SharedLimiterKey> configView) throws NotConfiguredException {
Config config = configView.getConfig();
SharedLimiterKey.GlobalLimiterPolicy globalLimiterPolicy = configView.getKey().getGlobalLimiterPolicy();
if (ConfigUtils.getBoolean(config, SKIP_GLOBAL_LIMITER_KEY, false)) {
if (globalLimiterPolicy != SharedLimiterKey.GlobalLimiterPolicy.LOCAL_ONLY) {
SharedLimiterKey modifiedKey = new SharedLimiterKey(configView.getKey().getResourceLimitedPath(), SharedLimiterKey.GlobalLimiterPolicy.LOCAL_ONLY);
return new ResourceCoordinate<>(this, modifiedKey, (S) configView.getScope());
}
} else if (config.hasPath(FAIL_IF_NO_GLOBAL_LIMITER_KEY) && config.getBoolean(FAIL_IF_NO_GLOBAL_LIMITER_KEY) && globalLimiterPolicy != SharedLimiterKey.GlobalLimiterPolicy.USE_GLOBAL) {
// if user has specified FAIL_IF_NO_GLOBAL_LIMITER_KEY, promote the policy from USE_GLOBAL_IF_CONFIGURED to USE_GLOBAL
// e.g. fail if no GLOBAL configuration is present
SharedLimiterKey modifiedKey = new SharedLimiterKey(configView.getKey().getResourceLimitedPath(), SharedLimiterKey.GlobalLimiterPolicy.USE_GLOBAL);
return new ResourceCoordinate<>(this, modifiedKey, (S) configView.getScope());
}
Limiter limiter;
if (!configView.getScope().isLocal() && !globalLimiterPolicy.equals(SharedLimiterKey.GlobalLimiterPolicy.LOCAL_ONLY)) {
try {
Class<?> klazz = Class.forName("org.apache.gobblin.util.limiter.RestliLimiterFactory");
return new ResourceCoordinate<>((SharedResourceFactory<Limiter, SharedLimiterKey, S>) klazz.newInstance(), configView.getKey(), (S) configView.getScope());
} catch (ReflectiveOperationException roe) {
if (globalLimiterPolicy.equals(SharedLimiterKey.GlobalLimiterPolicy.USE_GLOBAL)) {
throw new RuntimeException("There is no Global limiter factory in the classpath.");
}
}
}
if (config.hasPath(LIMITER_CLASS_KEY)) {
try {
LimiterFactory factory = RESOLVER.resolveClass(config.getString(LIMITER_CLASS_KEY)).newInstance();
limiter = factory.buildLimiter(config);
} catch (ReflectiveOperationException roe) {
throw new RuntimeException(roe);
}
} else {
if (config.hasPath(FAIL_ON_UNKNOWN_RESOURCE_ID) && config.getBoolean(FAIL_ON_UNKNOWN_RESOURCE_ID)) {
throw new NotConfiguredException();
}
limiter = new NoopLimiter();
}
ScopeType<S> scope = configView.getScope();
Collection<S> parentScopes = scope.parentScopes();
if (parentScopes != null) {
try {
for (S parentScope : parentScopes) {
limiter = new MultiLimiter(limiter, broker.getSharedResourceAtScope(this, configView.getKey(), parentScope));
}
} catch (NoSuchScopeException nsse) {
throw new RuntimeException("Could not get higher scope limiter. This is an error in code.", nsse);
}
}
return new ResourceInstance<>(limiter);
}
Aggregations