Examples with IOContext io.questdb.network.IOContext used on opensource projects

Example 1 with IOContext

use of io.questdb.network.IOContext in project questdb by bluestreak01.

the class IODispatcherTest method testQueuedConnectionTimeout.

@Test
public void testQueuedConnectionTimeout() throws Exception {
    LOG.info().$("started testQueuedConnectionTimeout").$();
    assertMemoryLeak(() -> {
        final int listenBackLog = 10;
        final int activeConnectionLimit = 5;
        final long queuedConnectionTimeoutInMs = 250;
        final IODispatcherConfiguration configuration = new DefaultIODispatcherConfiguration() {

            @Override
            public int getActiveConnectionLimit() {
                return activeConnectionLimit;
            }

            @Override
            public int getListenBacklog() {
                return listenBackLog;
            }

            @Override
            public long getQueuedConnectionTimeout() {
                return queuedConnectionTimeoutInMs;
            }
        };
        final AtomicInteger nConnected = new AtomicInteger();
        final LongHashSet serverConnectedFds = new LongHashSet();
        final LongHashSet clientActiveFds = new LongHashSet();
        IOContextFactory<IOContext> contextFactory = (fd, dispatcher) -> {
            LOG.info().$(fd).$(" connected").$();
            serverConnectedFds.add(fd);
            nConnected.incrementAndGet();
            return new IOContext() {

                @Override
                public boolean invalid() {
                    return !serverConnectedFds.contains(fd);
                }

                @Override
                public long getFd() {
                    return fd;
                }

                @Override
                public IODispatcher<?> getDispatcher() {
                    return dispatcher;
                }

                @Override
                public void close() {
                    LOG.info().$(fd).$(" disconnected").$();
                    serverConnectedFds.remove(fd);
                }
            };
        };
        final String request = "\n";
        long mem = TestUtils.toMemory(request);
        final long sockAddr = Net.sockaddr("127.0.0.1", 9001);
        Thread serverThread;
        final CountDownLatch serverLatch = new CountDownLatch(1);
        try (IODispatcher<IOContext> dispatcher = IODispatchers.create(configuration, contextFactory)) {
            serverThread = new Thread("test-io-dispatcher") {

                @Override
                public void run() {
                    long smem = Unsafe.malloc(1, MemoryTag.NATIVE_DEFAULT);
                    try {
                        IORequestProcessor<IOContext> requestProcessor = (operation, context) -> {
                            long fd = context.getFd();
                            int rc;
                            switch(operation) {
                                case IOOperation.READ:
                                    rc = Net.recv(fd, smem, 1);
                                    if (rc == 1) {
                                        dispatcher.registerChannel(context, IOOperation.WRITE);
                                    } else {
                                        dispatcher.disconnect(context, IODispatcher.DISCONNECT_REASON_TEST);
                                    }
                                    break;
                                case IOOperation.WRITE:
                                    rc = Net.send(fd, smem, 1);
                                    if (rc == 1) {
                                        dispatcher.registerChannel(context, IOOperation.READ);
                                    } else {
                                        dispatcher.disconnect(context, IODispatcher.DISCONNECT_REASON_TEST);
                                    }
                                    break;
                                default:
                                    dispatcher.disconnect(context, IODispatcher.DISCONNECT_REASON_TEST);
                            }
                        };
                        do {
                            dispatcher.run(0);
                            dispatcher.processIOQueue(requestProcessor);
                            Thread.yield();
                        } while (!isInterrupted());
                    } finally {
                        Unsafe.free(smem, 1, MemoryTag.NATIVE_DEFAULT);
                        serverLatch.countDown();
                    }
                }
            };
            serverThread.setDaemon(true);
            serverThread.start();
            // Connect exactly the right amount of clients to fill the active connection and connection backlog, after the
            // queuedConnectionTimeoutInMs the connections in the backlog should get refused
            int nClientConnects = 0;
            int nClientConnectRefused = 0;
            for (int i = 0; i < listenBackLog + activeConnectionLimit; i++) {
                long fd = Net.socketTcp(true);
                Assert.assertTrue(fd > -1);
                clientActiveFds.add(fd);
                if (Net.connect(fd, sockAddr) != 0) {
                    nClientConnectRefused++;
                    continue;
                }
                int rc = Net.send(fd, mem, request.length());
                if (rc < 0) {
                    nClientConnectRefused++;
                    continue;
                }
                rc = Net.recv(fd, mem, request.length());
                if (rc < 0) {
                    nClientConnectRefused++;
                } else {
                    nClientConnects++;
                }
            }
            Assert.assertEquals(activeConnectionLimit, nClientConnects);
            Assert.assertEquals(listenBackLog, nClientConnectRefused);
            Assert.assertFalse(dispatcher.isListening());
            // Close all connections and wait for server to resume listening
            while (clientActiveFds.size() > 0) {
                long fd = clientActiveFds.get(0);
                clientActiveFds.remove(fd);
                Net.close(fd);
            }
            long timeoutMs = System.currentTimeMillis() + TimeUnit.MINUTES.toMillis(1);
            while (!dispatcher.isListening()) {
                if (System.currentTimeMillis() > timeoutMs) {
                    Assert.fail("Timeout waiting for server to start listening again");
                }
            }
            // Try connections again to make sure server is listening
            nClientConnects = 0;
            nClientConnectRefused = 0;
            for (int i = 0; i < listenBackLog + activeConnectionLimit; i++) {
                long fd = Net.socketTcp(true);
                Assert.assertTrue(fd > -1);
                clientActiveFds.add(fd);
                if (Net.connect(fd, sockAddr) != 0) {
                    nClientConnectRefused++;
                    continue;
                }
                int rc = Net.send(fd, mem, request.length());
                if (rc < 0) {
                    nClientConnectRefused++;
                    continue;
                }
                rc = Net.recv(fd, mem, request.length());
                if (rc < 0) {
                    nClientConnectRefused++;
                } else {
                    nClientConnects++;
                }
            }
            Assert.assertEquals(activeConnectionLimit, nClientConnects);
            Assert.assertEquals(listenBackLog, nClientConnectRefused);
            Assert.assertFalse(dispatcher.isListening());
            // Close all remaining client connections
            for (int n = 0; n < clientActiveFds.size(); n++) {
                long fd = clientActiveFds.get(n);
                Net.close(fd);
            }
            serverThread.interrupt();
            if (!serverLatch.await(1, TimeUnit.MINUTES)) {
                Assert.fail("Timeout waiting for server to end");
            }
        } finally {
            Net.freeSockAddr(sockAddr);
            Unsafe.free(mem, request.length(), MemoryTag.NATIVE_DEFAULT);
        }
    });
}