Java tutorial
/** * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.hadoop.gateway.websockets; import java.io.File; import java.io.FileOutputStream; import java.io.IOException; import java.lang.management.ManagementFactory; import java.lang.management.MemoryMXBean; import java.net.URI; import java.net.URL; import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import javax.websocket.ContainerProvider; import javax.websocket.Endpoint; import javax.websocket.EndpointConfig; import javax.websocket.MessageHandler; import javax.websocket.Session; import javax.websocket.WebSocketContainer; import org.apache.commons.io.FileUtils; import org.apache.hadoop.gateway.config.GatewayConfig; import org.apache.hadoop.gateway.config.impl.GatewayConfigImpl; import org.apache.hadoop.gateway.deploy.DeploymentFactory; import org.apache.hadoop.gateway.services.DefaultGatewayServices; import org.apache.hadoop.gateway.services.GatewayServices; import org.apache.hadoop.gateway.services.ServiceLifecycleException; import org.apache.hadoop.gateway.services.topology.TopologyService; import org.apache.hadoop.gateway.topology.TopologyEvent; import org.apache.hadoop.gateway.topology.TopologyListener; import org.apache.hadoop.test.TestUtils; import org.easymock.EasyMock; import org.eclipse.jetty.server.Server; import org.eclipse.jetty.server.ServerConnector; import org.eclipse.jetty.server.handler.ContextHandler; import org.eclipse.jetty.server.handler.HandlerCollection; import org.eclipse.jetty.util.thread.QueuedThreadPool; import org.junit.AfterClass; import org.junit.BeforeClass; import org.junit.Test; import com.mycila.xmltool.XMLDoc; import com.mycila.xmltool.XMLTag; /** * Test how Knox holds up under multiple concurrent connections. * */ public class WebsocketMultipleConnectionTest { /** * Simulate backend websocket */ private static Server backendServer; /** * URI for backend websocket server */ private static URI backendServerUri; /** * Mock Gateway server */ private static Server gatewayServer; /** * Mock gateway config */ private static GatewayConfig gatewayConfig; private static GatewayServices services; /** * URI for gateway server */ private static URI serverUri; private static File topoDir; /** * Maximum number of open connections to test. */ private static int MAX_CONNECTIONS = 100; public WebsocketMultipleConnectionTest() { super(); } @BeforeClass public static void startServers() throws Exception { startWebsocketServer(); startGatewayServer(); } @AfterClass public static void stopServers() { try { gatewayServer.stop(); backendServer.stop(); } catch (final Exception e) { e.printStackTrace(System.err); } /* Cleanup the created files */ FileUtils.deleteQuietly(topoDir); } /** * Test websocket proxying through gateway. * * @throws Exception */ @Test public void testMultipleConnections() throws Exception { WebSocketContainer container = ContainerProvider.getWebSocketContainer(); final CountDownLatch latch = new CountDownLatch(MAX_CONNECTIONS); Session[] sessions = new Session[MAX_CONNECTIONS]; MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean(); System.gc(); final long heapt1 = memoryMXBean.getHeapMemoryUsage().getUsed(); final long nonHeapt1 = memoryMXBean.getNonHeapMemoryUsage().getUsed(); for (int i = 0; i < MAX_CONNECTIONS; i++) { sessions[i] = container.connectToServer(new WebsocketClient() { @Override public void onMessage(String message) { latch.countDown(); } }, new URI(serverUri.toString() + "gateway/websocket/ws")); } for (int i = 0; i < MAX_CONNECTIONS; i++) { /* make sure the session is active and valid before trying to connect */ if (sessions[i].isOpen() && sessions[i].getBasicRemote() != null) { sessions[i].getBasicRemote().sendText("OK"); } } latch.await(5 * MAX_CONNECTIONS, TimeUnit.MILLISECONDS); System.gc(); final long heapUsed = memoryMXBean.getHeapMemoryUsage().getUsed() - heapt1; final long nonHeapUsed = memoryMXBean.getNonHeapMemoryUsage().getUsed() - nonHeapt1; System.out.println("heapUsed = " + heapUsed); System.out.println("nonHeapUsed = " + nonHeapUsed); /* 90 KB per connection */ /* long expected = 90 * 1024 * MAX_CONNECTIONS; assertThat("heap used", heapUsed, lessThan(expected)); */ } /** * Start Mock Websocket server that acts as backend. * * @throws Exception */ private static void startWebsocketServer() throws Exception { backendServer = new Server(new QueuedThreadPool(254)); ServerConnector connector = new ServerConnector(backendServer); backendServer.addConnector(connector); final WebsocketEchoHandler handler = new WebsocketEchoHandler(); ContextHandler context = new ContextHandler(); context.setContextPath("/"); context.setHandler(handler); backendServer.setHandler(context); // Start Server backendServer.start(); String host = connector.getHost(); if (host == null) { host = "localhost"; } int port = connector.getLocalPort(); backendServerUri = new URI(String.format("ws://%s:%d/ws", host, port)); } /** * Start Gateway Server. * * @throws Exception */ private static void startGatewayServer() throws Exception { /* use default Max threads */ gatewayServer = new Server(new QueuedThreadPool(254)); final ServerConnector connector = new ServerConnector(gatewayServer); gatewayServer.addConnector(connector); /* workaround so we can add our handler later at runtime */ HandlerCollection handlers = new HandlerCollection(true); /* add some initial handlers */ ContextHandler context = new ContextHandler(); context.setContextPath("/"); handlers.addHandler(context); gatewayServer.setHandler(handlers); // Start Server gatewayServer.start(); String host = connector.getHost(); if (host == null) { host = "localhost"; } int port = connector.getLocalPort(); serverUri = new URI(String.format("ws://%s:%d/", host, port)); /* Setup websocket handler */ setupGatewayConfig(backendServerUri.toString()); final GatewayWebsocketHandler gatewayWebsocketHandler = new GatewayWebsocketHandler(gatewayConfig, services); handlers.addHandler(gatewayWebsocketHandler); gatewayWebsocketHandler.start(); } /** * Initialize the configs and components required for this test. * * @param backend * @throws IOException */ private static void setupGatewayConfig(final String backend) throws IOException { services = new DefaultGatewayServices(); topoDir = createDir(); URL serviceUrl = ClassLoader.getSystemResource("websocket-services"); final File descriptor = new File(topoDir, "websocket.xml"); final FileOutputStream stream = new FileOutputStream(descriptor); createKnoxTopology(backend).toStream(stream); stream.close(); final TestTopologyListener topoListener = new TestTopologyListener(); final Map<String, String> options = new HashMap<>(); options.put("persist-master", "false"); options.put("master", "password"); gatewayConfig = EasyMock.createNiceMock(GatewayConfig.class); EasyMock.expect(gatewayConfig.getGatewayTopologyDir()).andReturn(topoDir.toString()).anyTimes(); EasyMock.expect(gatewayConfig.getGatewayServicesDir()).andReturn(serviceUrl.getFile()).anyTimes(); EasyMock.expect(gatewayConfig.getEphemeralDHKeySize()).andReturn("2048").anyTimes(); EasyMock.expect(gatewayConfig.getGatewaySecurityDir()).andReturn(topoDir.toString()).anyTimes(); /* Websocket configs */ EasyMock.expect(gatewayConfig.isWebsocketEnabled()).andReturn(true).anyTimes(); EasyMock.expect(gatewayConfig.getWebsocketMaxTextMessageSize()) .andReturn(GatewayConfigImpl.DEFAULT_WEBSOCKET_MAX_TEXT_MESSAGE_SIZE).anyTimes(); EasyMock.expect(gatewayConfig.getWebsocketMaxBinaryMessageSize()) .andReturn(GatewayConfigImpl.DEFAULT_WEBSOCKET_MAX_BINARY_MESSAGE_SIZE).anyTimes(); EasyMock.expect(gatewayConfig.getWebsocketMaxTextMessageBufferSize()) .andReturn(GatewayConfigImpl.DEFAULT_WEBSOCKET_MAX_TEXT_MESSAGE_BUFFER_SIZE).anyTimes(); EasyMock.expect(gatewayConfig.getWebsocketMaxBinaryMessageBufferSize()) .andReturn(GatewayConfigImpl.DEFAULT_WEBSOCKET_MAX_BINARY_MESSAGE_BUFFER_SIZE).anyTimes(); EasyMock.expect(gatewayConfig.getWebsocketInputBufferSize()) .andReturn(GatewayConfigImpl.DEFAULT_WEBSOCKET_INPUT_BUFFER_SIZE).anyTimes(); EasyMock.expect(gatewayConfig.getWebsocketAsyncWriteTimeout()) .andReturn(GatewayConfigImpl.DEFAULT_WEBSOCKET_ASYNC_WRITE_TIMEOUT).anyTimes(); EasyMock.expect(gatewayConfig.getWebsocketIdleTimeout()) .andReturn(GatewayConfigImpl.DEFAULT_WEBSOCKET_IDLE_TIMEOUT).anyTimes(); EasyMock.replay(gatewayConfig); try { services.init(gatewayConfig, options); } catch (ServiceLifecycleException e) { e.printStackTrace(); } DeploymentFactory.setGatewayServices(services); final TopologyService monitor = services.getService(GatewayServices.TOPOLOGY_SERVICE); monitor.addTopologyChangeListener(topoListener); monitor.reloadTopologies(); } private static File createDir() throws IOException { return TestUtils.createTempDir(WebsocketEchoTest.class.getSimpleName() + "-"); } private static XMLTag createKnoxTopology(final String backend) { XMLTag xml = XMLDoc.newDocument(true).addRoot("topology").addTag("service").addTag("role") .addText("WEBSOCKET").addTag("url").addText(backend).gotoParent().gotoRoot(); // System.out.println( "GATEWAY=" + xml.toString() ); return xml; } private static class TestTopologyListener implements TopologyListener { public ArrayList<List<TopologyEvent>> events = new ArrayList<List<TopologyEvent>>(); @Override public void handleTopologyEvent(List<TopologyEvent> events) { this.events.add(events); synchronized (this) { for (TopologyEvent event : events) { if (!event.getType().equals(TopologyEvent.Type.DELETED)) { /* for this test we only care about this part */ DeploymentFactory.createDeployment(gatewayConfig, event.getTopology()); } } } } } private static abstract class WebsocketClient extends Endpoint implements MessageHandler.Whole<String> { @Override public void onOpen(Session session, EndpointConfig config) { session.addMessageHandler(this); } } }