package com.example.d2.client; import com.linkedin.common.callback.Callback; import com.linkedin.common.util.None; import com.linkedin.d2.balancer.D2Client; import com.linkedin.d2.balancer.D2ClientBuilder; import com.linkedin.r2.message.rest.RestRequest; import com.linkedin.r2.message.rest.RestRequestBuilder; import com.linkedin.r2.message.rest.RestResponse; import org.json.simple.JSONObject; import org.json.simple.parser.JSONParser; import org.json.simple.parser.ParseException; import java.io.File; import java.io.FileReader; import java.io.IOException; import java.net.URI; import java.util.Map; import java.util.Random; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; /** * In this example we will demonstrate the effects of assigning weights in partition. * As explained in the README.me we have 2 partitions: partition 0 and partition 1. * We have 2 servers. Both servers belong to partition 0 and 1 but the difference is * we set weight for server 1 to be 10x more than server 2. * * We use HASH partitioning scheme for this example. So when we send a request, D2 will * parse the partition key using this regex "user=(.*)". * * This regex is defined in d2Config.json * * So we can extract the key "oby" from a URI like d2://inbox?user=oby. If we cannot get * the key from the URI, we'll throw PartitionAccessException. * * Because we specified MD5 algorithm in d2Config.json, after we get the key we'll md5 * hash the key and determine the partition from the hash mod number of partitions. * * So in this example you will see server1 will receive 10x more requests * than server2. * * @author Oby Sumampouw (osumampo@linkedin.com) */ public class EmailClientExample { public static void main(String[] args) throws Exception { //get client configuration JSONObject json = parseConfig(); String zkConnectString = (String) json.get("zkConnectString"); Long zkSessionTimeout = (Long) json.get("zkSessionTimeout"); String zkBasePath = (String) json.get("zkBasePath"); Long zkStartupTimeout = (Long) json.get("zkStartupTimeout"); Long zkLoadBalancerNotificationTimeout = (Long) json.get("zkLoadBalancerNotificationTimeout"); String zkFlagFile = (String) json.get("zkFlagFile"); String fsBasePath = (String) json.get("fsBasePath"); final Map<String, Long> trafficProportion = (Map<String, Long>) json.get("trafficProportion"); final Long clientShutdownTimeout = (Long) json.get("clientShutdownTimeout"); final Long clientStartTimeout = (Long) json.get("clientStartTimeout"); System.out.println("Finished parsing client config"); //create d2 client final D2Client d2Client = new D2ClientBuilder().setZkHosts(zkConnectString) .setZkSessionTimeout( zkSessionTimeout, TimeUnit.MILLISECONDS) .setZkStartupTimeout( zkStartupTimeout, TimeUnit.MILLISECONDS) .setLbWaitTimeout( zkLoadBalancerNotificationTimeout, TimeUnit.MILLISECONDS) .setFlagFile(zkFlagFile) .setBasePath(zkBasePath) .setFsBasePath(fsBasePath) .build(); System.out.println("Finished creating d2 client, starting d2 client..."); ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor(); final CountDownLatch latch = new CountDownLatch(1); //start d2 client by connecting to zookeeper startClient(d2Client, executorService, clientStartTimeout, new Callback<None>() { @Override public void onError (Throwable e) { System.exit(1); } @Override public void onSuccess (None result) { latch.countDown(); } }); latch.await(); System.out.println("D2 client is sending traffic to both " + "partition 0 and partition1."); System.out.println("Note that traffic for server1 will be 22x more than server2"); ScheduledFuture task = executorService.scheduleAtFixedRate(new Runnable() { @Override public void run () { try { sendTraffic(trafficProportion, d2Client); } catch (Exception e) { e.printStackTrace(); } } }, 0, 1000, TimeUnit.MILLISECONDS); System.out.println("Press enter to shutdown"); System.out.println("===========================================================\n\n"); System.in.read(); task.cancel(false); System.out.println("Shutting down..."); shutdown(d2Client, executorService, clientShutdownTimeout); } private static void startClient(final D2Client d2Client, ExecutorService executorService, Long timeout, final Callback<None> callback) { try { executorService.submit(new Runnable() { @Override public void run () { d2Client.start(new Callback<None>() { @Override public void onError (Throwable e) { System.err.println("Error starting d2Client. Aborting... "); e.printStackTrace(); System.exit(1); } @Override public void onSuccess (None result) { System.out.println("D2 client started"); callback.onSuccess(None.none()); } }); } }).get(timeout, TimeUnit.MILLISECONDS); } catch (Exception e) { System.err.println("Cannot start d2 client. Timeout is set to " + timeout + " ms"); e.printStackTrace(); } } private static void shutdown(final D2Client d2Client, ExecutorService executorService, Long timeout) { try { executorService.submit(new Runnable() { @Override public void run () { d2Client.shutdown(new Callback<None>() { @Override public void onError (Throwable e) { System.err.println("Error shutting down d2Client."); e.printStackTrace(); } @Override public void onSuccess (None result) { System.out.println("D2 client stopped"); } }); } }).get(timeout, TimeUnit.MILLISECONDS); } catch (Exception e) { System.err.println("Cannot stop d2 client. Timeout is set to " + timeout + " ms"); e.printStackTrace(); } finally { executorService.shutdown(); } } private static JSONObject parseConfig() throws IOException, ParseException { String path = new File(new File(".").getAbsolutePath()).getCanonicalPath() + "/src/main/config/emailExample.json"; JSONParser parser = new JSONParser(); Object object = parser.parse(new FileReader(path)); return (JSONObject) object; } private static void sendTraffic (Map<String, Long> trafficProportion, D2Client d2Client) throws Exception { for (Map.Entry<String, Long> proportion : trafficProportion.entrySet()) { long queryPerSecond = proportion.getValue(); String serviceName = proportion.getKey(); Random random = new Random(); for (int i = 0; i < queryPerSecond; i++) { final URI uri = new URI("d2://" + serviceName + "?user=" + random.nextInt()); RestRequestBuilder requestBuilder = new RestRequestBuilder(uri).setMethod("get"); RestRequest request = requestBuilder.build(); //we don't care about the result from the server after all, //you can see the traffic hits the echo server from stdout d2Client.restRequest(request, new Callback<RestResponse>() { @Override public void onError (Throwable e) { System.err.println("URI = " + uri.toString() + " didn't get any response"); } @Override public void onSuccess (RestResponse result) { System.out.println("URI = " + uri.toString() + " was served by " + result.getEntity().asString("UTF-8")); } }); } } } }