Java tutorial
/** * Copyright 2016 Yahoo Inc. * * Licensed 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 com.yahoo.pulsar.discovery.service; import static com.yahoo.pulsar.discovery.service.web.ZookeeperCacheLoader.LOADBALANCE_BROKERS_ROOT; import static org.testng.Assert.assertNotEquals; import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; import java.net.InetSocketAddress; import java.net.URI; import java.net.URISyntaxException; import java.security.PrivateKey; import java.security.cert.X509Certificate; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import org.apache.bookkeeper.util.ZkUtils; import org.apache.zookeeper.CreateMode; import org.apache.zookeeper.ZooDefs; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import com.yahoo.pulsar.common.api.Commands; import com.yahoo.pulsar.common.policies.data.loadbalancer.LoadReport; import com.yahoo.pulsar.common.util.ObjectMapperFactory; import com.yahoo.pulsar.common.util.SecurityUtility; import com.yahoo.pulsar.zookeeper.ZookeeperClientFactoryImpl; import io.netty.bootstrap.Bootstrap; import io.netty.buffer.ByteBuf; import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelInboundHandlerAdapter; import io.netty.channel.ChannelInitializer; import io.netty.channel.nio.NioEventLoopGroup; import io.netty.channel.socket.SocketChannel; import io.netty.channel.socket.nio.NioSocketChannel; import io.netty.handler.ssl.SslContext; import io.netty.handler.ssl.SslContextBuilder; import io.netty.handler.ssl.util.InsecureTrustManagerFactory; public class DiscoveryServiceTest extends BaseDiscoveryTestSetup { private final static String TLS_CLIENT_CERT_FILE_PATH = "./src/test/resources/certificate/client.crt"; private final static String TLS_CLIENT_KEY_FILE_PATH = "./src/test/resources/certificate/client.key"; @BeforeMethod private void init() throws Exception { super.setup(); } @AfterMethod private void clean() throws Exception { super.cleanup(); } /** * Verifies: Discovery-service returns broker is round-robin manner * * @throws Exception */ @Test public void testBrokerDiscoveryRoundRobin() throws Exception { addBrokerToZk(5); String prevUrl = null; for (int i = 0; i < 10; i++) { String current = service.getDiscoveryProvider().nextBroker().getPulsarServiceUrl(); assertNotEquals(prevUrl, current); prevUrl = current; } } /** * It verifies: client connects to Discovery-service and receives discovery response successfully. * * @throws Exception */ @Test public void testClientServerConnection() throws Exception { addBrokerToZk(2); // 1. client connects to DiscoveryService, 2. Client receive service-lookup response final int messageTransfer = 2; final CountDownLatch latch = new CountDownLatch(messageTransfer); NioEventLoopGroup workerGroup = connectToService(service.getServiceUrl(), latch, false); try { assertTrue(latch.await(1, TimeUnit.SECONDS)); } catch (InterruptedException e) { fail("should have received lookup response message from server", e); } workerGroup.shutdownGracefully(); } @Test(enabled = true) public void testClientServerConnectionTls() throws Exception { addBrokerToZk(2); // 1. client connects to DiscoveryService, 2. Client receive service-lookup response final int messageTransfer = 2; final CountDownLatch latch = new CountDownLatch(messageTransfer); NioEventLoopGroup workerGroup = connectToService(service.getServiceUrlTls(), latch, true); try { assertTrue(latch.await(1, TimeUnit.SECONDS)); } catch (InterruptedException e) { fail("should have received lookup response message from server", e); } workerGroup.shutdownGracefully(); } /** * creates ClientHandler channel to connect and communicate with server * * @param serviceUrl * @param latch * @return * @throws URISyntaxException */ public static NioEventLoopGroup connectToService(String serviceUrl, CountDownLatch latch, boolean tls) throws URISyntaxException { NioEventLoopGroup workerGroup = new NioEventLoopGroup(); Bootstrap b = new Bootstrap(); b.group(workerGroup); b.channel(NioSocketChannel.class); b.handler(new ChannelInitializer<SocketChannel>() { @Override public void initChannel(SocketChannel ch) throws Exception { if (tls) { SslContextBuilder builder = SslContextBuilder.forClient(); builder.trustManager(InsecureTrustManagerFactory.INSTANCE); X509Certificate[] certificates = SecurityUtility .loadCertificatesFromPemFile(TLS_CLIENT_CERT_FILE_PATH); PrivateKey privateKey = SecurityUtility.loadPrivateKeyFromPemFile(TLS_CLIENT_KEY_FILE_PATH); builder.keyManager(privateKey, (X509Certificate[]) certificates); SslContext sslCtx = builder.build(); ch.pipeline().addLast("tls", sslCtx.newHandler(ch.alloc())); } ch.pipeline().addLast(new ClientHandler(latch)); } }); URI uri = new URI(serviceUrl); InetSocketAddress serviceAddress = new InetSocketAddress(uri.getHost(), uri.getPort()); b.connect(serviceAddress).addListener((ChannelFuture future) -> { if (!future.isSuccess()) { throw new IllegalStateException(future.cause()); } }); return workerGroup; } static class ClientHandler extends ChannelInboundHandlerAdapter { final CountDownLatch latch; public ClientHandler(CountDownLatch latch) { this.latch = latch; } @Override public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception { ByteBuf buffer = (ByteBuf) msg; buffer.release(); latch.countDown(); ctx.close(); } @Override public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { // Close the connection when an exception is raised. cause.printStackTrace(); ctx.close(); } @Override public void channelActive(ChannelHandlerContext ctx) throws Exception { super.channelActive(ctx); ctx.writeAndFlush(Commands.newConnect("", "")); latch.countDown(); } } private void addBrokerToZk(int number) throws Exception { for (int i = 0; i < number; i++) { LoadReport report = new LoadReport(null, null, "pulsar://broker-:15000" + i, null); String reportData = ObjectMapperFactory.getThreadLocal().writeValueAsString(report); ZkUtils.createFullPathOptimistic(mockZookKeeper, LOADBALANCE_BROKERS_ROOT + "/" + "broker-" + i, reportData.getBytes(ZookeeperClientFactoryImpl.ENCODING_SCHEME), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); } Thread.sleep(100); // wait to get cache updated } }