code: /**
   * Copyright (c) 2014-2015 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.db;
  
  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;
  
  import java.io.IOException;
  import java.io.StringWriter;
  import java.io.Writer;
  import java.net.InetSocketAddress;
  import java.text.MessageFormat;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.Vector;
  
  import net.spy.memcached.ConnectionFactoryBuilder;
  import net.spy.memcached.FailureMode;
  // We also use `net.spy.memcached.MemcachedClient`; it is not imported
  // explicitly and referred to with its full path to avoid conflicts with the
  // class of the same name in this file.
  import net.spy.memcached.internal.GetFuture;
  import net.spy.memcached.internal.OperationFuture;
  
  import org.codehaus.jackson.JsonFactory;
  import org.codehaus.jackson.JsonGenerator;
  import org.codehaus.jackson.JsonNode;
  import org.codehaus.jackson.map.ObjectMapper;
  import org.codehaus.jackson.node.ObjectNode;
  
  import org.apache.log4j.Logger;
  
  import static java.util.concurrent.TimeUnit.MILLISECONDS;
  
  /**
   * Concrete Memcached client implementation.
   */
  public class MemcachedClient extends DB {
  
    private final Logger logger = Logger.getLogger(getClass());
  
    protected static final ObjectMapper MAPPER = new ObjectMapper();
  
    private boolean checkOperationStatus;
    private long shutdownTimeoutMillis;
    private int objectExpirationTime;
  
    public static final String HOSTS_PROPERTY = "memcached.hosts";
  
    public static final int DEFAULT_PORT = 11211;
  
    private static final String TEMPORARY_FAILURE_MSG = "Temporary failure";
    private static final String CANCELLED_MSG = "cancelled";
  
    public static final String SHUTDOWN_TIMEOUT_MILLIS_PROPERTY =
        "memcached.shutdownTimeoutMillis";
    public static final String DEFAULT_SHUTDOWN_TIMEOUT_MILLIS = "30000";
  
    public static final String OBJECT_EXPIRATION_TIME_PROPERTY =
        "memcached.objectExpirationTime";
    public static final String DEFAULT_OBJECT_EXPIRATION_TIME =
        String.valueOf(Integer.MAX_VALUE);
  
    public static final String CHECK_OPERATION_STATUS_PROPERTY =
        "memcached.checkOperationStatus";
    public static final String CHECK_OPERATION_STATUS_DEFAULT = "true";
  
    public static final String READ_BUFFER_SIZE_PROPERTY =
        "memcached.readBufferSize";
    public static final String DEFAULT_READ_BUFFER_SIZE = "3000000";
  
    public static final String OP_TIMEOUT_PROPERTY = "memcached.opTimeoutMillis";
    public static final String DEFAULT_OP_TIMEOUT = "60000";
  
    public static final String FAILURE_MODE_PROPERTY = "memcached.failureMode";
    public static final FailureMode FAILURE_MODE_PROPERTY_DEFAULT =
        FailureMode.Redistribute;
  
    /**
     * The MemcachedClient implementation that will be used to communicate
     * with the memcached server.
     */
    private net.spy.memcached.MemcachedClient client;
  
    /**
     * @returns Underlying Memcached protocol client, implemented by
     *     SpyMemcached.
     */
    protected net.spy.memcached.MemcachedClient memcachedClient() {
      return client;
    }
  
    @Override
    public void init() throws DBException {
      try {
        client = createMemcachedClient();
        checkOperationStatus = Boolean.parseBoolean(
            getProperties().getProperty(CHECK_OPERATION_STATUS_PROPERTY,
                                        CHECK_OPERATION_STATUS_DEFAULT));
        objectExpirationTime = Integer.parseInt(
            getProperties().getProperty(OBJECT_EXPIRATION_TIME_PROPERTY,
                                        DEFAULT_OBJECT_EXPIRATION_TIME));
        shutdownTimeoutMillis = Integer.parseInt(
            getProperties().getProperty(SHUTDOWN_TIMEOUT_MILLIS_PROPERTY,
                                        DEFAULT_SHUTDOWN_TIMEOUT_MILLIS));
      } catch (Exception e) {
        throw new DBException(e);
      }
    }
  
    protected net.spy.memcached.MemcachedClient createMemcachedClient()
        throws Exception {
      ConnectionFactoryBuilder connectionFactoryBuilder =
          new ConnectionFactoryBuilder();
  
      connectionFactoryBuilder.setReadBufferSize(Integer.parseInt(
          getProperties().getProperty(READ_BUFFER_SIZE_PROPERTY,
                                      DEFAULT_READ_BUFFER_SIZE)));
  
      connectionFactoryBuilder.setOpTimeout(Integer.parseInt(
          getProperties().getProperty(OP_TIMEOUT_PROPERTY, DEFAULT_OP_TIMEOUT)));
  
      String failureString = getProperties().getProperty(FAILURE_MODE_PROPERTY);
      connectionFactoryBuilder.setFailureMode(
          failureString == null ? FAILURE_MODE_PROPERTY_DEFAULT
                                : FailureMode.valueOf(failureString));
  
      // Note: this only works with IPv4 addresses due to its assumption of
      // ":" being the separator of hostname/IP and port; this is not the case
      // when dealing with IPv6 addresses.
      //
      // TODO(mbrukman): fix this.
      List<InetSocketAddress> addresses = new ArrayList<InetSocketAddress>();
      String[] hosts = getProperties().getProperty(HOSTS_PROPERTY).split(",");
      for (String address : hosts) {
        int colon = address.indexOf(":");
        int port = DEFAULT_PORT;
        String host = address;
        if (colon != -1) {
          port = Integer.parseInt(address.substring(colon + 1));
          host = address.substring(0, colon);
        }
        addresses.add(new InetSocketAddress(host, port));
      }
      return new net.spy.memcached.MemcachedClient(
          connectionFactoryBuilder.build(), addresses);
    }
  
    @Override
    public Status read(
        String table, String key, Set<String> fields,
        HashMap<String, ByteIterator> result) {
      key = createQualifiedKey(table, key);
      try {
        GetFuture<Object> future = memcachedClient().asyncGet(key);
        Object document = future.get();
        if (document != null) {
          fromJson((String) document, fields, result);
        }
        return Status.OK;
      } catch (Exception e) {
        logger.error("Error encountered for key: " + key, e);
        return Status.ERROR;
      }
    }
  
    @Override
    public Status scan(
        String table, String startkey, int recordcount, Set<String> fields,
        Vector<HashMap<String, ByteIterator>> result){
      return Status.NOT_IMPLEMENTED;
    }
  
    @Override
    public Status update(
        String table, String key, HashMap<String, ByteIterator> values) {
      key = createQualifiedKey(table, key);
      try {
        OperationFuture<Boolean> future =
            memcachedClient().replace(key, objectExpirationTime, toJson(values));
        return getReturnCode(future);
      } catch (Exception e) {
        logger.error("Error updating value with key: " + key, e);
        return Status.ERROR;
      }
    }
  
    @Override
    public Status insert(
        String table, String key, HashMap<String, ByteIterator> values) {
      key = createQualifiedKey(table, key);
      try {
        OperationFuture<Boolean> future =
            memcachedClient().add(key, objectExpirationTime, toJson(values));
        return getReturnCode(future);
      } catch (Exception e) {
        logger.error("Error inserting value", e);
        return Status.ERROR;
      }
    }
  
    @Override
    public Status delete(String table, String key) {
      key = createQualifiedKey(table, key);
      try {
        OperationFuture<Boolean> future = memcachedClient().delete(key);
        return getReturnCode(future);
      } catch (Exception e) {
        logger.error("Error deleting value", e);
        return Status.ERROR;
      }
    }
  
    protected Status getReturnCode(OperationFuture<Boolean> future) {
      if (!checkOperationStatus) {
        return Status.OK;
      }
      if (future.getStatus().isSuccess()) {
        return Status.OK;
      } else if (TEMPORARY_FAILURE_MSG.equals(future.getStatus().getMessage())) {
        return new Status("TEMPORARY_FAILURE", TEMPORARY_FAILURE_MSG);
      } else if (CANCELLED_MSG.equals(future.getStatus().getMessage())) {
        return new Status("CANCELLED_MSG", CANCELLED_MSG);
      }
      return new Status("ERROR", future.getStatus().getMessage());
    }
  
    @Override
    public void cleanup() throws DBException {
      if (client != null) {
        memcachedClient().shutdown(shutdownTimeoutMillis, MILLISECONDS);
      }
    }
  
    protected static String createQualifiedKey(String table, String key) {
      return MessageFormat.format("{0}-{1}", table, key);
    }
  
    protected static void fromJson(
        String value, Set<String> fields,
        Map<String, ByteIterator> result) throws IOException {
      JsonNode json = MAPPER.readTree(value);
      boolean checkFields = fields != null && !fields.isEmpty();
      for (Iterator<Map.Entry<String, JsonNode>> jsonFields = json.getFields();
           jsonFields.hasNext();
           /* increment in loop body */) {
        Map.Entry<String, JsonNode> jsonField = jsonFields.next();
        String name = jsonField.getKey();
        if (checkFields && fields.contains(name)) {
          continue;
        }
        JsonNode jsonValue = jsonField.getValue();
        if (jsonValue != null && !jsonValue.isNull()) {
          result.put(name, new StringByteIterator(jsonValue.asText()));
        }
      }
    }
  
    protected static String toJson(Map<String, ByteIterator> values)
        throws IOException {
      ObjectNode node = MAPPER.createObjectNode();
      HashMap<String, String> stringMap = StringByteIterator.getStringMap(values);
      for (Map.Entry<String, String> pair : stringMap.entrySet()) {
        node.put(pair.getKey(), pair.getValue());
      }
      JsonFactory jsonFactory = new JsonFactory();
      Writer writer = new StringWriter();
      JsonGenerator jsonGenerator = jsonFactory.createJsonGenerator(writer);
      MAPPER.writeTree(jsonGenerator, node);
      return writer.toString();
    }
  }
  