Examples with LeakyBucketStreamThrottler org.apache.nifi.stream.io.LeakyBucketStreamThrottler used on opensource projects

Example 1 with LeakyBucketStreamThrottler

use of org.apache.nifi.stream.io.LeakyBucketStreamThrottler in project nifi by apache.

the class TestLeakyBucketThrottler method testInputStreamInterface.

@Test(timeout = 10000)
public void testInputStreamInterface() throws IOException {
    final byte[] data = new byte[1024 * 1024 * 4];
    // throttle rate at 1 MB/sec
    try (final LeakyBucketStreamThrottler throttler = new LeakyBucketStreamThrottler(1024 * 1024);
        final ByteArrayInputStream bais = new ByteArrayInputStream(data);
        final InputStream throttledIn = throttler.newThrottledInputStream(bais);
        final ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
        final byte[] buffer = new byte[4096];
        final long start = System.currentTimeMillis();
        int len;
        while ((len = throttledIn.read(buffer)) > 0) {
            baos.write(buffer, 0, len);
        }
        final long millis = System.currentTimeMillis() - start;
        // should take 4 sec give or take
        assertTrue(millis > 3000);
        assertTrue(millis < 6000);
    }
}

Example 2 with LeakyBucketStreamThrottler

use of org.apache.nifi.stream.io.LeakyBucketStreamThrottler in project nifi by apache.

the class TestLeakyBucketThrottler method testOutputStreamInterface.

@Test(timeout = 10000)
public void testOutputStreamInterface() throws IOException {
    // throttle rate at 1 MB/sec
    final LeakyBucketStreamThrottler throttler = new LeakyBucketStreamThrottler(1024 * 1024);
    final byte[] data = new byte[1024 * 1024 * 4];
    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    try (final OutputStream throttledOut = throttler.newThrottledOutputStream(baos)) {
        final long start = System.currentTimeMillis();
        throttledOut.write(data);
        throttler.close();
        final long millis = System.currentTimeMillis() - start;
        // should take 4 sec give or take
        assertTrue(millis > 3000);
        assertTrue(millis < 6000);
    }
}

Example 3 with LeakyBucketStreamThrottler

use of org.apache.nifi.stream.io.LeakyBucketStreamThrottler in project nifi by apache.

the class ListenHTTP method createHttpServerFromService.

private void createHttpServerFromService(final ProcessContext context) throws Exception {
    final String basePath = context.getProperty(BASE_PATH).evaluateAttributeExpressions().getValue();
    final SSLContextService sslContextService = context.getProperty(SSL_CONTEXT_SERVICE).asControllerService(SSLContextService.class);
    final Double maxBytesPerSecond = context.getProperty(MAX_DATA_RATE).asDataSize(DataUnit.B);
    final StreamThrottler streamThrottler = (maxBytesPerSecond == null) ? null : new LeakyBucketStreamThrottler(maxBytesPerSecond.intValue());
    final int returnCode = context.getProperty(RETURN_CODE).asInteger();
    throttlerRef.set(streamThrottler);
    final boolean needClientAuth = sslContextService != null && sslContextService.getTrustStoreFile() != null;
    final SslContextFactory contextFactory = new SslContextFactory();
    contextFactory.setNeedClientAuth(needClientAuth);
    if (needClientAuth) {
        contextFactory.setTrustStorePath(sslContextService.getTrustStoreFile());
        contextFactory.setTrustStoreType(sslContextService.getTrustStoreType());
        contextFactory.setTrustStorePassword(sslContextService.getTrustStorePassword());
    }
    final String keystorePath = sslContextService == null ? null : sslContextService.getKeyStoreFile();
    if (keystorePath != null) {
        final String keystorePassword = sslContextService.getKeyStorePassword();
        final String keyStoreType = sslContextService.getKeyStoreType();
        contextFactory.setKeyStorePath(keystorePath);
        contextFactory.setKeyManagerPassword(keystorePassword);
        contextFactory.setKeyStorePassword(keystorePassword);
        contextFactory.setKeyStoreType(keyStoreType);
    }
    if (sslContextService != null) {
        contextFactory.setProtocol(sslContextService.getSslAlgorithm());
    }
    // thread pool for the jetty instance
    final QueuedThreadPool threadPool = new QueuedThreadPool();
    threadPool.setName(String.format("%s (%s) Web Server", getClass().getSimpleName(), getIdentifier()));
    // create the server instance
    final Server server = new Server(threadPool);
    // get the configured port
    final int port = context.getProperty(PORT).evaluateAttributeExpressions().asInteger();
    final ServerConnector connector;
    final HttpConfiguration httpConfiguration = new HttpConfiguration();
    if (keystorePath == null) {
        // create the connector
        connector = new ServerConnector(server, new HttpConnectionFactory(httpConfiguration));
    } else {
        // configure the ssl connector
        httpConfiguration.setSecureScheme("https");
        httpConfiguration.setSecurePort(port);
        httpConfiguration.addCustomizer(new SecureRequestCustomizer());
        // build the connector
        connector = new ServerConnector(server, new SslConnectionFactory(contextFactory, "http/1.1"), new HttpConnectionFactory(httpConfiguration));
    }
    // configure the port
    connector.setPort(port);
    // add the connector to the server
    server.setConnectors(new Connector[] { connector });
    final ServletContextHandler contextHandler = new ServletContextHandler(server, "/", true, (keystorePath != null));
    for (final Class<? extends Servlet> cls : getServerClasses()) {
        final Path path = cls.getAnnotation(Path.class);
        // also, servlets other than ListenHttpServlet must have a path starting with /
        if (basePath.isEmpty() && !path.value().isEmpty()) {
            // Note: this is to handle the condition of an empty uri, otherwise pathSpec would start with //
            contextHandler.addServlet(cls, path.value());
        } else {
            contextHandler.addServlet(cls, "/" + basePath + path.value());
        }
    }
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_PROCESSOR, this);
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_LOGGER, getLogger());
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_SESSION_FACTORY_HOLDER, sessionFactoryReference);
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_PROCESS_CONTEXT_HOLDER, context);
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_FLOWFILE_MAP, flowFileMap);
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_AUTHORITY_PATTERN, Pattern.compile(context.getProperty(AUTHORIZED_DN_PATTERN).getValue()));
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_STREAM_THROTTLER, streamThrottler);
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_BASE_PATH, basePath);
    contextHandler.setAttribute(CONTEXT_ATTRIBUTE_RETURN_CODE, returnCode);
    if (context.getProperty(HEADERS_AS_ATTRIBUTES_REGEX).isSet()) {
        contextHandler.setAttribute(CONTEXT_ATTRIBUTE_HEADER_PATTERN, Pattern.compile(context.getProperty(HEADERS_AS_ATTRIBUTES_REGEX).getValue()));
    }
    try {
        server.start();
    } catch (Exception e) {
        shutdownHttpServer(server);
        throw e;
    }
    this.server = server;
}

Example 4 with LeakyBucketStreamThrottler

use of org.apache.nifi.stream.io.LeakyBucketStreamThrottler in project nifi by apache.

the class TestLeakyBucketThrottler method testDirectInterface.

@Test(timeout = 10000)
public void testDirectInterface() throws IOException, InterruptedException {
    // throttle rate at 1 MB/sec
    try (final LeakyBucketStreamThrottler throttler = new LeakyBucketStreamThrottler(1024 * 1024);
        final ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
        // create 3 threads, each sending ~2 MB
        final List<Thread> threads = new ArrayList<Thread>();
        for (int i = 0; i < 3; i++) {
            final Thread t = new WriterThread(i, throttler, baos);
            threads.add(t);
        }
        final long start = System.currentTimeMillis();
        for (final Thread t : threads) {
            t.start();
        }
        for (final Thread t : threads) {
            t.join();
        }
        final long elapsed = System.currentTimeMillis() - start;
        throttler.close();
        // To send 15 MB, it should have taken at least 5 seconds and no more than 7 seconds, to
        // allow for busy-ness and the fact that we could write a tiny bit more than the limit.
        assertTrue(elapsed > 5000);
        assertTrue(elapsed < 7000);
        // ensure bytes were copied out appropriately
        assertEquals(3 * (2 * 1024 * 1024 + 1), baos.size());
        assertEquals((byte) 'A', baos.toByteArray()[baos.size() - 1]);
    }
}