Examples with Limiter org.apache.gobblin.util.limiter.Limiter used on opensource projects

Example 1 with Limiter

use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.

the class MRStressTest method createLimiter.

static Limiter createLimiter(Configuration configuration, SharedResourcesBroker<SimpleScopeType> broker) {
    try {
        Limiter limiter = new NoopLimiter();
        long localQps = configuration.getLong(LOCALLY_ENFORCED_QPS, 0);
        if (localQps > 0) {
            log.info("Setting up local qps " + localQps);
            limiter = new MultiLimiter(limiter, new RateBasedLimiter(localQps));
        }
        if (configuration.getBoolean(USE_THROTTLING_SERVER, false)) {
            log.info("Setting up remote throttling.");
            String resourceId = configuration.get(RESOURCE_ID);
            Limiter globalLimiter = broker.getSharedResource(new RestliLimiterFactory<SimpleScopeType>(), new SharedLimiterKey(resourceId));
            limiter = new MultiLimiter(limiter, globalLimiter);
        }
        return limiter;
    } catch (NotConfiguredException nce) {
        throw new RuntimeException(nce);
    }
}

Example 2 with Limiter

use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.

the class LimiterStopEventTest method testGetLimiterStopMetadataCase0.

@Test
public void testGetLimiterStopMetadataCase0() throws InterruptedException {
    Properties properties = new Properties();
    String key1 = "topic";
    String key2 = "partition.id";
    String key3 = "others";
    String keyList = Joiner.on(',').join(key1, key2);
    properties.setProperty(LimiterConfigurationKeys.LIMITER_REPORT_KEY_LIST, keyList);
    properties.setProperty(key1, "1111");
    properties.setProperty(key2, "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(key1), true);
        Assert.assertEquals(metaData.containsKey(key2), true);
        Assert.assertEquals(metaData.containsKey(key3), false);
    } catch (Exception e) {
        Assert.fail();
    }
}

Example 3 with Limiter

use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.

the class LimiterStopEventTest method testGetLimiterStopMetadataCase3.

@Test
public void testGetLimiterStopMetadataCase3() throws InterruptedException {
    Properties properties = new Properties();
    String key1 = "topic";
    String key2 = "partition.id";
    String keyList = Joiner.on(',').join(key1, key2);
    String subKey1 = key2 + "....";
    String subKey2 = key2 + "##fjpaierbng;";
    String subKey3 = key2 + "x[n  sdf";
    String subKey4 = key2 + "";
    properties.setProperty(LimiterConfigurationKeys.LIMITER_REPORT_KEY_LIST, keyList);
    properties.setProperty(subKey1, "1111");
    properties.setProperty(subKey2, "1111");
    properties.setProperty(subKey3, "1111");
    properties.setProperty(subKey4, "1111");
    properties.setProperty(key1, "1111");
    properties.setProperty(key2, "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(key1), true);
        Assert.assertEquals(metaData.containsKey(key2), true);
        Assert.assertEquals(metaData.containsKey(subKey1), true);
        Assert.assertEquals(metaData.containsKey(subKey2), true);
        Assert.assertEquals(metaData.containsKey(subKey3), true);
        Assert.assertEquals(metaData.containsKey(subKey4), true);
    } catch (Exception e) {
        Assert.fail();
    }
}

Example 4 with Limiter

use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.

the class StreamThrottler method doThrottleInputStream.

/**
 * Throttles an {@link InputStream} if throttling is configured.
 * @param inputStream {@link InputStream} to throttle.
 * @param sourceURI used for selecting the throttling policy.
 * @param targetURI used for selecting the throttling policy.
 */
@Builder(builderMethodName = "throttleInputStream", builderClassName = "InputStreamThrottler")
private ThrottledInputStream doThrottleInputStream(InputStream inputStream, URI sourceURI, URI targetURI) {
    Preconditions.checkNotNull(inputStream, "InputStream cannot be null.");
    Limiter limiter = new NoopLimiter();
    if (sourceURI != null && targetURI != null) {
        StreamCopierSharedLimiterKey key = new StreamCopierSharedLimiterKey(sourceURI, targetURI);
        try {
            limiter = new MultiLimiter(limiter, this.broker.getSharedResource(new SharedLimiterFactory<S>(), key));
        } catch (NotConfiguredException nce) {
            log.warn("Could not create a Limiter for key " + key, nce);
        }
    } else {
        log.info("Not throttling input stream because source or target URIs are not defined.");
    }
    Optional<MeteredInputStream> meteredStream = MeteredInputStream.findWrappedMeteredInputStream(inputStream);
    if (!meteredStream.isPresent()) {
        meteredStream = Optional.of(MeteredInputStream.builder().in(inputStream).build());
        inputStream = meteredStream.get();
    }
    return new ThrottledInputStream(inputStream, limiter, meteredStream.get());
}

Example 5 with Limiter

use of org.apache.gobblin.util.limiter.Limiter in project incubator-gobblin by apache.

the class ThrottledFileSystemTest method testListing.

@Test
public void testListing() throws Exception {
    FileSystem fs = Mockito.mock(FileSystem.class);
    Mockito.when(fs.listStatus(Mockito.any(Path.class))).thenAnswer(new Answer<FileStatus[]>() {

        @Override
        public FileStatus[] answer(InvocationOnMock invocation) throws Throwable {
            Path path = (Path) invocation.getArguments()[0];
            int files = Integer.parseInt(path.getName());
            FileStatus status = new FileStatus(0, false, 0, 0, 0, new Path("/"));
            FileStatus[] out = new FileStatus[files];
            for (int i = 0; i < files; i++) {
                out[i] = status;
            }
            return out;
        }
    });
    Limiter limiter = new CountBasedLimiter(5);
    ThrottledFileSystem throttledFileSystem = new ThrottledFileSystem(fs, limiter, "testService");
    Assert.assertEquals(throttledFileSystem.getServiceName(), "testService");
    // use 1 permit
    Assert.assertNotNull(throttledFileSystem.listStatus(new Path("/files/99")));
    // use 3 permits
    Assert.assertNotNull(throttledFileSystem.listStatus(new Path("/files/250")));
    try {
        // requires 2 permits
        throttledFileSystem.listStatus(new Path("/files/150"));
        Assert.fail();
    } catch (NotEnoughPermitsException expected) {
    // Expected
    }
    // requires 1 permit
    Assert.assertNotNull(throttledFileSystem.listStatus(new Path("/files/99")));
}