Java tutorial
/** * Copyright (c) 2016 YCSB contributors. All rights reserved. * * 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. See accompanying * LICENSE file. */ package com.yahoo.ycsb.webservice.rest; import java.io.BufferedReader; import java.io.ByteArrayInputStream; import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; import java.util.HashMap; import java.util.Properties; import java.util.Set; import java.util.Vector; import java.util.zip.GZIPInputStream; import javax.ws.rs.HttpMethod; import org.apache.http.HttpEntity; import org.apache.http.client.ClientProtocolException; import org.apache.http.client.config.RequestConfig; import org.apache.http.client.methods.CloseableHttpResponse; import org.apache.http.client.methods.HttpDelete; import org.apache.http.client.methods.HttpEntityEnclosingRequestBase; import org.apache.http.client.methods.HttpGet; import org.apache.http.client.methods.HttpPost; import org.apache.http.client.methods.HttpPut; import org.apache.http.entity.ContentType; import org.apache.http.entity.InputStreamEntity; import org.apache.http.impl.client.CloseableHttpClient; import org.apache.http.impl.client.HttpClientBuilder; import org.apache.http.util.EntityUtils; import com.yahoo.ycsb.ByteIterator; import com.yahoo.ycsb.DB; import com.yahoo.ycsb.DBException; import com.yahoo.ycsb.Status; import com.yahoo.ycsb.StringByteIterator; /** * Class responsible for making web service requests for benchmarking purpose. * Using Apache HttpClient over standard Java HTTP API as this is more flexible * and provides better functionality. For example HttpClient can automatically * handle redirects and proxy authentication which the standard Java API can't. */ public class RestClient extends DB { private static final String URL_PREFIX = "url.prefix"; private static final String CON_TIMEOUT = "timeout.con"; private static final String READ_TIMEOUT = "timeout.read"; private static final String EXEC_TIMEOUT = "timeout.exec"; private static final String LOG_ENABLED = "log.enable"; private static final String HEADERS = "headers"; private static final String COMPRESSED_RESPONSE = "response.compression"; private boolean compressedResponse; private boolean logEnabled; private String urlPrefix; private Properties props; private String[] headers; private CloseableHttpClient client; private int conTimeout = 10000; private int readTimeout = 10000; private int execTimeout = 10000; private volatile Criteria requestTimedout = new Criteria(false); @Override public void init() throws DBException { props = getProperties(); urlPrefix = props.getProperty(URL_PREFIX, "http://127.0.0.1:8080"); conTimeout = Integer.valueOf(props.getProperty(CON_TIMEOUT, "10")) * 1000; readTimeout = Integer.valueOf(props.getProperty(READ_TIMEOUT, "10")) * 1000; execTimeout = Integer.valueOf(props.getProperty(EXEC_TIMEOUT, "10")) * 1000; logEnabled = Boolean.valueOf(props.getProperty(LOG_ENABLED, "false").trim()); compressedResponse = Boolean.valueOf(props.getProperty(COMPRESSED_RESPONSE, "false").trim()); headers = props.getProperty(HEADERS, "Accept */* Content-Type application/xml user-agent Mozilla/5.0 ") .trim().split(" "); setupClient(); } private void setupClient() { RequestConfig.Builder requestBuilder = RequestConfig.custom(); requestBuilder = requestBuilder.setConnectTimeout(conTimeout); requestBuilder = requestBuilder.setConnectionRequestTimeout(readTimeout); requestBuilder = requestBuilder.setSocketTimeout(readTimeout); HttpClientBuilder clientBuilder = HttpClientBuilder.create() .setDefaultRequestConfig(requestBuilder.build()); this.client = clientBuilder.setConnectionManagerShared(true).build(); } @Override public Status read(String table, String endpoint, Set<String> fields, HashMap<String, ByteIterator> result) { int responseCode; try { responseCode = httpGet(urlPrefix + endpoint, result); } catch (Exception e) { responseCode = handleExceptions(e, urlPrefix + endpoint, HttpMethod.GET); } if (logEnabled) { System.err.println(new StringBuilder("GET Request: ").append(urlPrefix).append(endpoint) .append(" | Response Code: ").append(responseCode).toString()); } return getStatus(responseCode); } @Override public Status insert(String table, String endpoint, HashMap<String, ByteIterator> values) { int responseCode; try { responseCode = httpExecute(new HttpPost(urlPrefix + endpoint), values.get("data").toString()); } catch (Exception e) { responseCode = handleExceptions(e, urlPrefix + endpoint, HttpMethod.POST); } if (logEnabled) { System.err.println(new StringBuilder("POST Request: ").append(urlPrefix).append(endpoint) .append(" | Response Code: ").append(responseCode).toString()); } return getStatus(responseCode); } @Override public Status delete(String table, String endpoint) { int responseCode; try { responseCode = httpDelete(urlPrefix + endpoint); } catch (Exception e) { responseCode = handleExceptions(e, urlPrefix + endpoint, HttpMethod.DELETE); } if (logEnabled) { System.err.println(new StringBuilder("DELETE Request: ").append(urlPrefix).append(endpoint) .append(" | Response Code: ").append(responseCode).toString()); } return getStatus(responseCode); } @Override public Status update(String table, String endpoint, HashMap<String, ByteIterator> values) { int responseCode; try { responseCode = httpExecute(new HttpPut(urlPrefix + endpoint), values.get("data").toString()); } catch (Exception e) { responseCode = handleExceptions(e, urlPrefix + endpoint, HttpMethod.PUT); } if (logEnabled) { System.err.println(new StringBuilder("PUT Request: ").append(urlPrefix).append(endpoint) .append(" | Response Code: ").append(responseCode).toString()); } return getStatus(responseCode); } @Override public Status scan(String table, String startkey, int recordcount, Set<String> fields, Vector<HashMap<String, ByteIterator>> result) { return Status.NOT_IMPLEMENTED; } // Maps HTTP status codes to YCSB status codes. private Status getStatus(int responseCode) { int rc = responseCode / 100; if (responseCode == 400) { return Status.BAD_REQUEST; } else if (responseCode == 403) { return Status.FORBIDDEN; } else if (responseCode == 404) { return Status.NOT_FOUND; } else if (responseCode == 501) { return Status.NOT_IMPLEMENTED; } else if (responseCode == 503) { return Status.SERVICE_UNAVAILABLE; } else if (rc == 5) { return Status.ERROR; } return Status.OK; } private int handleExceptions(Exception e, String url, String method) { if (logEnabled) { System.err.println(new StringBuilder(method).append(" Request: ").append(url).append(" | ") .append(e.getClass().getName()).append(" occured | Error message: ").append(e.getMessage()) .toString()); } if (e instanceof ClientProtocolException) { return 400; } return 500; } // Connection is automatically released back in case of an exception. private int httpGet(String endpoint, HashMap<String, ByteIterator> result) throws IOException { requestTimedout.setIsSatisfied(false); Thread timer = new Thread(new Timer(execTimeout, requestTimedout)); timer.start(); int responseCode = 200; HttpGet request = new HttpGet(endpoint); for (int i = 0; i < headers.length; i = i + 2) { request.setHeader(headers[i], headers[i + 1]); } CloseableHttpResponse response = client.execute(request); responseCode = response.getStatusLine().getStatusCode(); HttpEntity responseEntity = response.getEntity(); // If null entity don't bother about connection release. if (responseEntity != null) { InputStream stream = responseEntity.getContent(); /* * TODO: Gzip Compression must be supported in the future. Header[] * header = response.getAllHeaders(); * if(response.getHeaders("Content-Encoding")[0].getValue().contains * ("gzip")) stream = new GZIPInputStream(stream); */ BufferedReader reader = new BufferedReader(new InputStreamReader(stream, "UTF-8")); StringBuffer responseContent = new StringBuffer(); String line = ""; while ((line = reader.readLine()) != null) { if (requestTimedout.isSatisfied()) { // Must avoid memory leak. reader.close(); stream.close(); EntityUtils.consumeQuietly(responseEntity); response.close(); client.close(); throw new TimeoutException(); } responseContent.append(line); } timer.interrupt(); result.put("response", new StringByteIterator(responseContent.toString())); // Closing the input stream will trigger connection release. stream.close(); } EntityUtils.consumeQuietly(responseEntity); response.close(); client.close(); return responseCode; } private int httpExecute(HttpEntityEnclosingRequestBase request, String data) throws IOException { requestTimedout.setIsSatisfied(false); Thread timer = new Thread(new Timer(execTimeout, requestTimedout)); timer.start(); int responseCode = 200; for (int i = 0; i < headers.length; i = i + 2) { request.setHeader(headers[i], headers[i + 1]); } InputStreamEntity reqEntity = new InputStreamEntity(new ByteArrayInputStream(data.getBytes()), ContentType.APPLICATION_FORM_URLENCODED); reqEntity.setChunked(true); request.setEntity(reqEntity); CloseableHttpResponse response = client.execute(request); responseCode = response.getStatusLine().getStatusCode(); HttpEntity responseEntity = response.getEntity(); // If null entity don't bother about connection release. if (responseEntity != null) { InputStream stream = responseEntity.getContent(); if (compressedResponse) { stream = new GZIPInputStream(stream); } BufferedReader reader = new BufferedReader(new InputStreamReader(stream, "UTF-8")); StringBuffer responseContent = new StringBuffer(); String line = ""; while ((line = reader.readLine()) != null) { if (requestTimedout.isSatisfied()) { // Must avoid memory leak. reader.close(); stream.close(); EntityUtils.consumeQuietly(responseEntity); response.close(); client.close(); throw new TimeoutException(); } responseContent.append(line); } timer.interrupt(); // Closing the input stream will trigger connection release. stream.close(); } EntityUtils.consumeQuietly(responseEntity); response.close(); client.close(); return responseCode; } private int httpDelete(String endpoint) throws IOException { requestTimedout.setIsSatisfied(false); Thread timer = new Thread(new Timer(execTimeout, requestTimedout)); timer.start(); int responseCode = 200; HttpDelete request = new HttpDelete(endpoint); for (int i = 0; i < headers.length; i = i + 2) { request.setHeader(headers[i], headers[i + 1]); } CloseableHttpResponse response = client.execute(request); responseCode = response.getStatusLine().getStatusCode(); response.close(); client.close(); return responseCode; } /** * Marks the input {@link Criteria} as satisfied when the input time has elapsed. */ class Timer implements Runnable { private long timeout; private Criteria timedout; public Timer(long timeout, Criteria timedout) { this.timedout = timedout; this.timeout = timeout; } @Override public void run() { try { Thread.sleep(timeout); this.timedout.setIsSatisfied(true); } catch (InterruptedException e) { // Do nothing. } } } /** * Sets the flag when a criteria is fulfilled. */ class Criteria { private boolean isSatisfied; public Criteria(boolean isSatisfied) { this.isSatisfied = isSatisfied; } public boolean isSatisfied() { return isSatisfied; } public void setIsSatisfied(boolean satisfied) { this.isSatisfied = satisfied; } } /** * Private exception class for execution timeout. */ class TimeoutException extends RuntimeException { private static final long serialVersionUID = 1L; public TimeoutException() { super("HTTP Request exceeded execution time limit."); } } }