LinkStoreRocksDb.java

/*
 * Copyright 2012, Facebook, 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.facebook.LinkBench;

import com.facebook.rocks.swift.*;
import com.facebook.swift.service.ThriftClientManager;
import com.facebook.nifty.client.FramedClientConnector;
import com.google.common.net.HostAndPort;
import org.apache.thrift.transport.TTransportException;

import java.io.IOException;
import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.Charset;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Collections;
import java.util.Date;
import java.util.HashSet;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;

import org.apache.log4j.Level;
import org.apache.log4j.Logger;
import org.apache.commons.codec.binary.Hex;

import static com.google.common.net.HostAndPort.fromParts;

/*
 * This file implements Linkbench methods for loading/requesting data to rocksDb
 * database server by calling thrift apis after creating a java thrift client
 * through swift : rocksClient for the link and the node operations.
 */

public class LinkStoreRocksDb extends GraphStore {
  private static final ThriftClientManager clientManager =
    new ThriftClientManager();
  private ThreadLocal<RocksService> rocksClient =
      new ThreadLocal<RocksService>();

  /* RocksDb database server configuration keys */
  public static final String CONFIG_HOST = "host";
  public static final String CONFIG_PORT = "port";
  public static final String CONFIG_WRITE_SYNC = "write_options_sync";
  public static final String CONFIG_WRITE_DISABLE_WAL =
    "write_options_disableWAL";

  public static final String CONFIG_USER = "user";
  public static final String CONFIG_PASSWORD = "password";

  public static final int DEFAULT_BULKINSERT_SIZE = 1024;
  private static final boolean INTERNAL_TESTING = false;

  private static int totalThreads = 0;

  String host;
  int port;
  WriteOptions writeOptions;
  String user;
  String pwd;

  Level debuglevel;

  int bulkInsertSize = DEFAULT_BULKINSERT_SIZE;

  private final Logger logger = Logger.getLogger(ConfigUtil.LINKBENCH_LOGGER);

  private RocksService getRocksClient() throws Exception {
    if (rocksClient.get() == null) {
      try {
        rocksClient.set(clientManager.createClient(
          new FramedClientConnector(fromParts(host, port)),
          RocksService.class).get());
        logger.info("Opened Rocksdb connection to " + host
                    + ":" + port);
      } catch (Exception e) {
        logger.error("Error in open rocksdb to " + host
                     + ":" + port + " " + e);
        throw e;
      }
    }
    return rocksClient.get();
  }

  static synchronized void incrThreads() {
     totalThreads++;
  }

  static synchronized boolean isLastThread() {
    if (--totalThreads == 0) {
      return true;
    } else {
      return false;
    }
  }

  @Override
  public void close() {
    try {
      if (!isLastThread()) {
        return;
      }
      if (clientManager != null) {
        clientManager.close();
      }
    } catch (Exception ioex) {
      logger.error("Error while closing client connection: " + ioex);
    }
  }

  @Override
  public void initialize(Properties p, Phase currentPhase, int threadId)
      throws IOException, Exception {
    incrThreads();
    host = ConfigUtil.getPropertyRequired(p, CONFIG_HOST);
    port = ConfigUtil.getInt(p, CONFIG_PORT);
    writeOptions = new WriteOptions();
    writeOptions.setSync(ConfigUtil.getBool(p, CONFIG_WRITE_SYNC, false));
    writeOptions.setDisableWAL(
      ConfigUtil.getBool(p, CONFIG_WRITE_DISABLE_WAL, false));
    debuglevel = ConfigUtil.getDebugLevel(p);
  }

  public LinkStoreRocksDb() {
    super();
  }

  public LinkStoreRocksDb(Properties props) throws IOException, Exception {
    super();
    initialize(props, Phase.LOAD, 0);
  }

  public void clearErrors(int threadID) {
    logger.warn("Closing Rocksdb connection in threadID " + threadID);
    try {
      if (rocksClient.get() != null) {
        rocksClient.get().close();
        rocksClient.remove();
      }
    } catch (Throwable e) {
      logger.error("Error in Reopen!" + e);
      e.printStackTrace();
    }
  }

  @Override
  public boolean addLink(String dbid, Link l, boolean noinverse)
      throws Exception{
    try {
      return addLinkImpl(dbid, l, noinverse);
    } catch (Exception ex) {
      logger.error("addlink failed! " + ex);
      throw ex;
    }
  }

  private boolean addLinkImpl(String dbid, Link l, boolean noinverse)
      throws Exception {

    if (Level.DEBUG.isGreaterOrEqual(debuglevel)) {
      logger.debug("addLink " + l.id1 +
                         "." + l.id2 +
                         "." + l.link_type);
    }
    AssocVisibility av = AssocVisibility.values()[l.visibility];
    String s = "wormhole...";
    dbid += "assocs";
    long result = getRocksClient().TaoAssocPut(
        dbid.getBytes(), l.link_type, l.id1, l.id2, l.time,
        av, true, Long.valueOf(l.version), l.data, s.getBytes(),
        writeOptions);

    return result == 1;
  }

  /**
   * Internal method: add links without updating the count
   */
  private boolean addLinksNoCount(String dbid, List<Link> links)
      throws Exception {
    if (links.size() == 0)
      return false;

    dbid += "assocs";
    for (Link l:links) {
      AssocVisibility av = AssocVisibility.values()[l.visibility];
      String s = "wormhole...";
      long result =
      getRocksClient().TaoAssocPut(dbid.getBytes(), l.link_type, l.id1,
         l.id2, l.time, av, false, Long.valueOf(l.version), l.data,
         s.getBytes(), writeOptions);
    }
    return true;
}

  @Override
  public boolean deleteLink(String dbid, long id1, long link_type, long id2,
                         boolean noinverse, boolean expunge)
    throws Exception {
    try {
      return deleteLinkImpl(dbid, id1, link_type, id2, noinverse, expunge);
    } catch (Exception ex) {
      logger.error("deletelink failed! " + ex);
      throw ex;
    }
  }

  private boolean deleteLinkImpl(String dbid, long id1, long link_type,
    long id2, boolean noinverse, boolean expunge) throws Exception {
    if (Level.DEBUG.isGreaterOrEqual(debuglevel)) {
      logger.debug("deleteLink " + id1 +
                         "." + id2 +
                         "." + link_type);
    }
    String s = "wormhole...";
    dbid += "assocs";
    long result = getRocksClient().TaoAssocDelete(
      dbid.getBytes() , link_type, id1, id2,
      -1 /*version ignored*/, AssocVisibility.HARD_DELETE, true,
      s.getBytes(), writeOptions);
    return result == 1;
  }

  @Override
  public boolean updateLink(String dbid, Link l, boolean noinverse)
    throws Exception {
    // Retry logic is in addLink
    boolean added = addLink(dbid, l, noinverse);
    return !added; // return true if updated instead of added
  }


  // lookup using id1, type, id2
  @Override
  public Link getLink(String dbid, long id1, long link_type, long id2)
    throws Exception {
    try {
      return getLinkImpl(dbid, id1, link_type, id2);
    } catch (Exception ex) {
      logger.error("getLink failed! " + ex);
      throw ex;
    }
  }

  private Link getLinkImpl(String dbid, long id1, long link_type, long id2)
    throws Exception {
    Link res[] = multigetLinks(dbid, id1, link_type, new long[] {id2});
    if (res == null)
      return null;
    assert(res.length <= 1);
    return res.length == 0 ? null : res[0];
  }


  @Override
  public Link[] multigetLinks(String dbid, long id1, long link_type,
    long[] id2s) throws Exception {
    try {
      return multigetLinksImpl(dbid, id1, link_type, id2s);
    } catch (Exception ex) {
      logger.error("multigetlinks failed! " + ex);
      throw ex;
    }
  }

  private Link[] multigetLinksImpl(String dbid, long id1, long link_type,
    long[] id2s) throws Exception {
    List<Long> l = new ArrayList<Long>();
    for (int i = 0; i < id2s.length; i++) {
      l.add(new Long(id2s[i]));
    }
    dbid += "assocs";
    List<TaoAssocGetEntry> tr = getRocksClient().TaoAssocGetID2s(
        dbid.getBytes(),
      link_type, id1, l);
    Link results[] = new Link[tr.size()];
    int i = 0;
    for (TaoAssocGetEntry tar : tr) {
      results[i] = new Link(id1, link_type, tar.getId2(),
          LinkStore.VISIBILITY_DEFAULT, tar.getData(),
          (int)(tar.getVersion()), tar.getTime());
    }
    return results;
  }

  // lookup using just id1, type
  @Override
  public Link[] getLinkList(String dbid, long id1, long link_type)
    throws Exception {
    return getLinkListImpl(
        dbid, id1, link_type, 0, Long.MAX_VALUE, 0, rangeLimit);
  }

  @Override
  public Link[] getLinkList(String dbid, long id1, long link_type,
    long minTimestamp, long maxTimestamp, int offset, int limit)
    throws Exception {
    try {
      return getLinkListImpl(dbid, id1, link_type, minTimestamp,
                             maxTimestamp, offset, limit);
    } catch (Exception ex) {
      logger.error("getLinkList failed! " + ex);
      throw ex;
    }
  }

  private Link[] getLinkListImpl(String dbid, long id1, long link_type,
    long minTimestamp, long maxTimestamp, int offset, int limit)
    throws Exception {
    dbid += "assocs";
    List<TaoAssocGetEntry> tr = getRocksClient().TaoAssocGetTimeRange(
        dbid.getBytes(), link_type, id1, minTimestamp, maxTimestamp,
        Long.valueOf(offset), Long.valueOf(limit));
    Link results[] = new Link[tr.size()];
    int i = 0;
    for (TaoAssocGetEntry tar : tr) {
      results[i] = new Link(id1, link_type, tar.getId2(),
          LinkStore.VISIBILITY_DEFAULT, tar.getData(),
          (int)(tar.getVersion()), tar.getTime());
      i++;
    }
    return results;
  }

  // count the #links
  @Override
  public long countLinks(String dbid, long id1, long link_type)
    throws Exception {
    try {
      return countLinksImpl(dbid, id1, link_type);
    } catch (Exception ex) {
      logger.error("countLinks failed! " + ex);
      throw ex;
    }
  }

  private long countLinksImpl(String dbid, long id1, long link_type)
    throws Exception {
    dbid += "assocs";
    long count = getRocksClient().TaoAssocCount(
      dbid.getBytes(), link_type, id1);
    if (Level.TRACE.isGreaterOrEqual(debuglevel)) {
      logger.trace("Count result: " + id1 + "," + link_type +
                         " is " + count);
    }
    return count;
  }

  @Override
  public int bulkLoadBatchSize() {
    return bulkInsertSize;
  }

  @Override
  public void addBulkLinks(String dbid, List<Link> links, boolean noinverse)
    throws Exception {
    try {
      addBulkLinksImpl(dbid, links, noinverse);
    } catch (Exception ex) {
      logger.error("addBulkLinks failed! " + ex);
      throw ex;
    }
  }

  private void addBulkLinksImpl(String dbid, List<Link> links,
    boolean noinverse) throws Exception {
    if (Level.TRACE.isGreaterOrEqual(debuglevel)) {
      logger.trace("addBulkLinks: " + links.size() + " links");
    }
    addLinksNoCount(dbid, links);
  }

  @Override
  public void addBulkCounts(String dbid, List<LinkCount> counts)
    throws Exception {
    try {
      addBulkCountsImpl(dbid, counts);
    } catch (Exception ex) {
      logger.error("addbulkCounts failed! " + ex);
      throw ex;
    }
  }

  private void addBulkCountsImpl(String dbid, List<LinkCount> counts)
    throws Exception {
    if (Level.TRACE.isGreaterOrEqual(debuglevel)) {
      logger.trace("addBulkCounts: " + counts.size() + " link counts");
    }
    if (counts.size() == 0)
      return;

    dbid += "assocs";
    for (LinkCount count: counts) {
      getRocksClient().TaoAssocCountPut(
        dbid.getBytes(), count.link_type, count.id1, count.count, null, writeOptions);
    }
  }

  @Override
  public void resetNodeStore(String dbid, long startID) throws Exception {
    //doesn't have a defined utility for Rocksdb
  }

  @Override
  public long addNode(String dbid, Node node) throws Exception {
    try {
      return addNodeImpl(dbid, node);
    } catch (Exception ex) {
      logger.error("addNode failed! " + ex);
      throw ex;
    }
  }

  private long addNodeImpl(String dbid, Node node) throws Exception {
    long ids[] = bulkAddNodes(dbid, Collections.singletonList(node));
    assert(ids.length == 1);
    return ids[0];
  }

  @Override
  public long[] bulkAddNodes(String dbid, List<Node> nodes) throws Exception {
    try {
      return bulkAddNodesImpl(dbid, nodes);
    } catch (Exception ex) {
      logger.error("bulkAddNodes failed! " + ex);
      throw ex;
    }
  }

  private long[] bulkAddNodesImpl(String dbid, List<Node> nodes)
    throws Exception {
    long newIds[] = new long[nodes.size()];
    int i = 0;
    for (Node n : nodes) {
      getRocksClient().TaoFBObjectPut(
        n.id, n.type, (int) n.version, (int) n.version,
        (long) n.time, n.data, true, null, writeOptions);
      newIds[i++] = n.id;
    }
    return newIds;
  }

  @Override
  public Node getNode(String dbid, int type, long id) throws Exception {
    try {
      return getNodeImpl(dbid, type, id);
    } catch (Exception ex) {
      logger.error("getnode failed! " + ex);
      throw ex;
    }
  }

  private Node getNodeImpl(String dbid, int type, long id) throws Exception {
    ReadOptions ropts = new ReadOptions();
    ropts.setVerifyChecksums(true);
    ropts.setFillCache(true);

    TaoFBObjectGetResult rgr = getRocksClient().TaoFBObjectGet(id, type);
    if (!rgr.isFound()) {
      return null; //Node was not found
    } else {
      return new Node(
        id, type, rgr.getVersion(), (int) rgr.getUpdateTime(), rgr.getData());
    }
  }

  @Override
  public boolean updateNode(String dbid, Node node) throws Exception {
    try {
      return updateNodeImpl(dbid, node);
    } catch (Exception ex) {
      logger.error("updateNode failed! " + ex);
      throw ex;
    }
  }

  private boolean updateNodeImpl(String dbid, Node node) throws Exception {
    return addNode(dbid, node) == 1;
  }

  @Override
  public boolean deleteNode(String dbid, int type, long id) throws Exception {
    try {
      return deleteNodeImpl(dbid, type, id);
    } catch (Exception ex) {
      logger.error("deleteNode failed! " + ex);
      throw ex;
    }
  }

  private boolean deleteNodeImpl(String dbid, int type, long id)
    throws Exception {
    getRocksClient().
      TaoFBObjectDel(id, type, null, writeOptions);
    return true;
  }
}