Package com.pardot.rhombus.functional

Source Code of com.pardot.rhombus.functional.TableScannerITCase

package com.pardot.rhombus.functional;


import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.pardot.rhombus.ConnectionManager;
import com.pardot.rhombus.ObjectMapper;
import com.pardot.rhombus.RhombusException;
import com.pardot.rhombus.TableScanner;
import com.pardot.rhombus.cobject.CKeyspaceDefinition;
import com.pardot.rhombus.cobject.CObjectTokenVisitor;
import com.pardot.rhombus.cobject.CObjectTokenVisitorFactory;
import com.pardot.rhombus.cobject.CQLGenerationException;
import com.pardot.rhombus.util.JsonUtil;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.commons.io.input.ReversedLinesFileReader;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileOutputStream;
import java.io.PrintWriter;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;

import static org.junit.Assert.*;

public class TableScannerITCase extends RhombusFunctionalTest {

  private static Logger logger = LoggerFactory.getLogger(TableScannerITCase.class);


  @Test
  public void testOnePartition() throws Exception {
    String objectType = "simple";
    String index1Value = "value1";
    String index2Value = "value2";

    //Build the connection manager
    ConnectionManager cm = getConnectionManager();
    cm.setLogCql(true);

    //Build our keyspace definition object
    CKeyspaceDefinition definition = JsonUtil.objectFromJsonResource(CKeyspaceDefinition.class, this.getClass().getClassLoader(), "SimpleKeyspace.js");

    //Rebuild the keyspace and get the object mapper
    cm.buildKeyspace(definition, true);
    logger.debug("Built keyspace: {}", definition.getName());
    cm.setDefaultKeyspace(definition);
    ObjectMapper om = cm.getObjectMapper();
    om.setLogCql(true);

    // Insert 20 values
    List<Map<String, Object>> values = this.getNValues(20, index1Value, index2Value);
    for(Map<String, Object> insertValue : values) {
      om.insert(objectType, insertValue);
    }

    VisitorFactoryTester visitorFactory = new VisitorFactoryTester();
    TableScanner scanner = new TableScanner(om, objectType, 1, visitorFactory, null);
    scanner.scan();

    long totalCount = 0;
    for(VisitorTester visitor : visitorFactory.getInstances()) {
      totalCount += visitor.getObjectCount();
    }

    assertEquals(20l, totalCount);
    cm.teardown();
  }

  @Test
  public void testFourPartitions() throws Exception {
    String objectType = "simple";
    String index1Value = "value1";
    String index2Value = "value2";
    String savepointDirectoryName = "savepoint-test-dir";
    File savepointDirectory = new File(savepointDirectoryName);
    if (savepointDirectory.exists()) {
      FileUtils.deleteRecursive(savepointDirectory);
    }
    Integer numPartitions = 4;

    //Build the connection manager
    ConnectionManager cm = getConnectionManager();
    cm.setLogCql(true);

    //Build our keyspace definition object
    CKeyspaceDefinition definition = JsonUtil.objectFromJsonResource(CKeyspaceDefinition.class, this.getClass().getClassLoader(), "SimpleKeyspace.js");

    //Rebuild the keyspace and get the object mapper
    cm.buildKeyspace(definition, true);
    logger.debug("Built keyspace: {}", definition.getName());
    cm.setDefaultKeyspace(definition);
    ObjectMapper om = cm.getObjectMapper();
    om.setLogCql(true);

    // Insert 83 values
    Long valueCount = 83L;
    List<Map<String, Object>> values = this.getNValues(valueCount, index1Value, index2Value);
    for(Map<String, Object> insertValue : values) {
      om.insert(objectType, insertValue);
    }

    VisitorFactoryTester visitorFactory = new VisitorFactoryTester();
    TableScanner scanner = new TableScanner(om, objectType, numPartitions, visitorFactory, savepointDirectoryName);
    List<Map.Entry<Long, Long>> ranges = scanner.makeRanges();
    scanner.scan();

    Long totalCount = 0L;
    for(VisitorTester visitor : visitorFactory.getInstances()) {
      totalCount += visitor.getObjectCount();
    }

    assertEquals(valueCount, totalCount);

    this.verifySavepoints(numPartitions, savepointDirectoryName, objectType, om, ranges);

    cm.teardown();
  }

  @Test
  public void testStartingFromSavepoint() throws Exception {
    String objectType = "simple";
    String index1Value = "value1";
    String index2Value = "value2";
    String savepointDirectoryName = "savepoint-test-dir";
    File savepointDirectory = new File(savepointDirectoryName);
    if (savepointDirectory.exists()) {
      FileUtils.deleteRecursive(savepointDirectory);
    }
    Integer numPartitions = 2;

    //Build the connection manager
    ConnectionManager cm = getConnectionManager();
    cm.setLogCql(true);

    //Build our keyspace definition object
    CKeyspaceDefinition definition = JsonUtil.objectFromJsonResource(CKeyspaceDefinition.class, this.getClass().getClassLoader(), "SimpleKeyspace.js");

    //Rebuild the keyspace and get the object mapper
    cm.buildKeyspace(definition, true);
    logger.debug("Built keyspace: {}", definition.getName());
    cm.setDefaultKeyspace(definition);
    ObjectMapper om = cm.getObjectMapper();
    om.setLogCql(true);

    // Insert 83 values
    Long valueCount = 83L;
    List<Map<String, Object>> values = this.getNValues(valueCount, index1Value, index2Value);
    for(Map<String, Object> insertValue : values) {
      om.insert(objectType, insertValue);
    }

    // Calculate halfway through each range
    UUID[] rangeHalfwayPoints = new UUID[numPartitions];
    BigInteger fullRange = BigInteger.valueOf(TableScanner.maxToken).subtract(BigInteger.valueOf(TableScanner.minToken));
    BigInteger rangeLength = fullRange.divide(BigInteger.valueOf(numPartitions));
    BigInteger rangeStart = BigInteger.valueOf(TableScanner.minToken);
    for(int i = 0 ; i < numPartitions - 1 ; i++) {
      BigInteger rangeEnd = rangeStart.add(rangeLength).subtract(BigInteger.ONE);
      Long halfwayToken = rangeLength.divide(BigInteger.valueOf(2)).add(rangeStart).longValue();
      rangeHalfwayPoints[i] = (UUID)om.scanTableWithStartToken(objectType, halfwayToken, TableScanner.maxToken, 1L).get(0).get("id");

      rangeStart = rangeEnd.add(BigInteger.ONE);
    }
    Long halfwayToken = rangeLength.divide(BigInteger.valueOf(2)).add(rangeStart).longValue();
    rangeHalfwayPoints[numPartitions - 1] = (UUID)om.scanTableWithStartToken(objectType, halfwayToken, TableScanner.maxToken, 1L).get(0).get("id");

    // Write halfway points to the savepoint directory
    savepointDirectory = new File(savepointDirectoryName);
    savepointDirectory.mkdir();

    for (int i = 0; i < numPartitions; i++) {
      String filename = TableScanner.getSavepointFilename(i);
      PrintWriter writer = new PrintWriter(new FileOutputStream(savepointDirectory.getName() + "/" + filename, false));
      writer.write(rangeHalfwayPoints[i].toString() + "\n");
      writer.close();
    }

    // Create and run scanner
    VisitorFactoryTester visitorFactory = new VisitorFactoryTester();
    TableScanner scanner = new TableScanner(om, objectType, numPartitions, visitorFactory, savepointDirectoryName);
    List<Map.Entry<Long, Long>> ranges = scanner.makeRanges();
    scanner.scan();

    Long actualCount = 0L;
    for(VisitorTester visitor : visitorFactory.getInstances()) {
      actualCount += visitor.getObjectCount();
    }

    // The total number of visits we have will fluctuate depending on where the ids land among the token ranges, but we should be close to half
    assertTrue(Math.abs(actualCount - (valueCount/2)) < (valueCount/(numPartitions * 2)));

    this.verifySavepoints(numPartitions, savepointDirectoryName, objectType, om, ranges);

    cm.teardown();
  }

  private void verifySavepoints(Integer numPartitions, String savepointDirectoryName, String objectType, ObjectMapper om, List<Map.Entry<Long, Long>> ranges) throws Exception {
    // Open up the savepoint directory
    File savepointDirectory = new File(savepointDirectoryName);
    assertTrue(savepointDirectory.exists());
    assertTrue(savepointDirectory.isDirectory());

    File[] fileArray = savepointDirectory.listFiles();
    assertNotNull(fileArray);

    // Make sure we have all the savepoint files we expect
    Map<String, File> files = Maps.newHashMap();
    for (File savepoint : fileArray) {
      files.put(savepoint.getName(), savepoint);
    }

    boolean foundAtLeastOneLine = false;
    UUID highestUuid = null;
    for (int i = 0; i < numPartitions; i++) {
      String filename = TableScanner.getSavepointFilename(i);
      assertTrue(files.containsKey(filename));

      File file = files.get(filename);
      ReversedLinesFileReader reader = new ReversedLinesFileReader(file);
      String line = reader.readLine();
      // A null line is actually ok, that just means there were no results in that partition's token range this time around
      if (line != null) {
        foundAtLeastOneLine = true;
        UUID savedUuid = UUID.fromString(line);
        assertNotNull(savedUuid);
        List<Map<String, Object>> values = om.scanTableWithStartId(objectType, savedUuid.toString(), TableScanner.maxToken, 1L);

        // This means there is no next id from this uuid so our normal check doesn't work.
        // This also means this is the highest uuid in the total range (might not be in the last partition if the last partition didn't end up with any objects)
        if (values.size() == 0) {
          // Make sure this only happens once
          assertNull(highestUuid);
          highestUuid = savedUuid;
        } else {
          // Otherwise there is a next uuid from the saved point, so compare that to the next uuid from the end of the partition's range and make sure they match
          UUID nextSavedUuid = (UUID) values.get(0).get("id");
          UUID nextExpectedUuid = (UUID) om.scanTableWithStartToken(objectType, ranges.get(i).getValue(), TableScanner.maxToken, 1L).get(0).get("id");
          assertEquals(nextExpectedUuid, nextSavedUuid);
        }
      }
    }
    assertTrue(foundAtLeastOneLine);
  }

  @Test
  public void testBigScan() throws Exception {
    String objectType = "simple";

    //Build the connection manager
    ConnectionManager cm = getConnectionManager();
    cm.setLogCql(true);

    //Build our keyspace definition object
    CKeyspaceDefinition definition = JsonUtil.objectFromJsonResource(CKeyspaceDefinition.class, this.getClass().getClassLoader(), "SimpleKeyspace.js");

    //Rebuild the keyspace and get the object mapper
    cm.buildKeyspace(definition, true);
    cm.setDefaultKeyspace(definition);
    ObjectMapper om = cm.getObjectMapper();
    om.setLogCql(true);

    long insertNum = 2000l;
    long batchSize = 200;
    insertNObjects(om, insertNum, batchSize);

    VisitorFactoryTester visitorFactory = new VisitorFactoryTester();
    TableScanner scanner = new TableScanner(om, objectType, 1, visitorFactory, null);
    scanner.scan();

    long totalCount = 0;
    for(VisitorTester visitor : visitorFactory.getInstances()) {
      totalCount += visitor.getObjectCount();
    }

    assertEquals(insertNum, totalCount);
    cm.teardown();
  }

  public void insertNObjects(ObjectMapper om, long number, long batchSize) throws RhombusException, CQLGenerationException {
    if(number < batchSize) {
      batchSize = number;
    }
    long numInserted = 0;
    while(numInserted < number) {
      long numLeft = number - numInserted;
      long numToInsert = batchSize;
      if(numLeft < batchSize) {
        numToInsert = numLeft;
      }
      List<Map<String, Object>> objects = getNValues(numToInsert, "value1", "value2");
      Map<String, List<Map<String, Object>>> batch = Maps.newHashMap();
      batch.put("simple", objects);
      om.insertBatchMixed(batch);
      numInserted += numToInsert;
    }
  }

  private List<Map<String, Object>> getNValues(long number, String index1Value, String index2Value) {
    List<Map<String, Object>> values = new ArrayList<Map<String, Object>>();
    for(int i = 0 ; i < number ; i++) {
      Map<String, Object> value = Maps.newHashMap();
      value.put("index_1", index1Value);
      value.put("index_2", index2Value);
      value.put("value", String.valueOf(i));
      values.add(value);
    }
    return values;
  }

  private class VisitorTester extends CObjectTokenVisitor {

    private long objectCount = 0;

    public VisitorTester(String name) {
      super(name);
    }

    @Override
    public void applyAction(Map<String, Object> object) {
      objectCount++;
    }

    @Override
    public boolean shouldInclude(Map<String, Object> object) {
      return true;
    }

    @Override
    public void setUp() {

    }

    @Override
    public void cleanUp() {

    }

    public long getObjectCount() {
      return objectCount;
    }
  }

  private class VisitorFactoryTester implements CObjectTokenVisitorFactory {

    private List<VisitorTester> instances;

    public VisitorFactoryTester() {
      instances = Lists.newArrayList();
    }

    @Override
    public CObjectTokenVisitor getInstance(String name) {
      VisitorTester tester = new VisitorTester(name);
      instances.add(tester);
      return tester;
    }

    public List<VisitorTester> getInstances() {
      return instances;
    }
  }
}
TOP

Related Classes of com.pardot.rhombus.functional.TableScannerITCase

TOP
Copyright © 2018 www.massapi.com. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.