com.yugabyte.sample.apps.CassandraPersonalization Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of yb-sample-apps Show documentation
Sample applications for benchmarking YugaByte DB functionality on various workload types.
The newest version!
// Copyright (c) YugaByte, 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.yugabyte.sample.apps;

import java.util.Arrays;
import java.util.Calendar;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Random;
import java.util.Vector;

import org.apache.log4j.Logger;

import com.datastax.driver.core.BatchStatement;
import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.BoundStatement;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.yugabyte.sample.common.SimpleLoadGenerator.Key;

/**
 * This workload writes and reads some coupons for some random customers repeatedly. The numbers of
 * concurrent writers and readers are configurable.
 */
public class CassandraPersonalization extends AppBase {
  private static final Logger LOG = Logger.getLogger(CassandraPersonalization.class);

  // Static initialization of this workload's config. These are good defaults for getting a decent
  // read dominated workload on a reasonably powered machine. Exact IOPS will of course vary
  // depending on the machine and what resources it has to spare.
  static {
    // Disable the read-write percentage.
    appConfig.readIOPSPercentage = -1;
    // Set the read and write threads to 1 each.
    appConfig.numReaderThreads = 24;
    appConfig.numWriterThreads = 2;
    // The number of keys to read.
    appConfig.numKeysToRead = -1;
    // The number of keys to write. This is the combined total number of inserts and updates.
    appConfig.numKeysToWrite = -1;
    // The number of unique keys to write. This determines the number of inserts (as opposed to
    // updates).
    appConfig.numUniqueKeysToWrite = NUM_UNIQUE_KEYS;
  }

  // The default table name to create and use for CRUD ops.
  private static final String DEFAULT_TABLE_NAME = "coupon_recommendation";

  // The shared prepared select statement for fetching the data.
  private static volatile PreparedStatement preparedSelect;

  // The shared prepared statement for inserting into the table.
  private static volatile PreparedStatement preparedInsert;

  // Lock for initializing prepared statement objects.
  private static final Object prepareInitLock = new Object();

  static Random rand = new Random();

  Vector coupons;

  static class Coupon {
    public final String code;
    public final Date beginDate;
    public final Date endDate;

    public Coupon(String code) {
      this.code = code;
      this.beginDate = generateRandomDate(true);
      this.endDate = generateRandomDate(false);
    }

    private Date generateRandomDate(boolean past) {
      Calendar cal = Calendar.getInstance();
      cal.set(Calendar.HOUR, 0);
      cal.set(Calendar.MINUTE, 0);
      cal.set(Calendar.SECOND, 0);
      cal.set(Calendar.MILLISECOND, 0);
      cal.add(Calendar.DATE, (past ? -1 : 1) * rand.nextInt(7));
      return cal.getTime();
    }
  }

  public CassandraPersonalization() {
    if (appConfig.readOnly) {
      return;
    }
    coupons = new Vector(appConfig.maxCouponsPerCustomer);
    for (int i = 0; i < appConfig.maxCouponsPerCustomer; i++) {
      coupons.add(new Coupon(Long.toString(800000000000L + (long)i)));
    }
    if (appConfig.numNewCouponsPerCustomer < appConfig.maxCouponsPerCustomer) {
      Collections.shuffle(coupons);
    }
  }

  /**
   * Drop the table created by this app.
   */
  @Override
  public void dropTable() {
    dropCassandraTable(getTableName());
  }

  @Override
  public List getCreateTableStatements() {
    String create_stmt = String.format(
      "CREATE TABLE IF NOT EXISTS %s (" +
      "  customer_id text," +
      "  store_id text," +
      "  coupon_code text," +
      "  begin_date timestamp," +
      "  end_date timestamp," +
      "  relevance_score double," +
      "  PRIMARY KEY ((customer_id, store_id), coupon_code));",
      getTableName());
    return Arrays.asList(create_stmt);
  }

  public String getTableName() {
    return appConfig.tableName != null ? appConfig.tableName : DEFAULT_TABLE_NAME;
  }

  protected PreparedStatement getPreparedSelect(String selectStmt)  {
    if (preparedSelect == null) {
      synchronized (prepareInitLock) {
        if (preparedSelect == null) {
          // Create the prepared statement object.
          preparedSelect = getCassandraClient().prepare(selectStmt);
        }
      }
    }
    return preparedSelect;
  }

  private PreparedStatement getPreparedSelect()  {
    return getPreparedSelect(
      String.format("SELECT * FROM %s WHERE customer_id = ? AND store_id = ?;", getTableName()));
  }

  @Override
  public long doRead() {
    Key key = getSimpleLoadGenerator().getKeyToRead();
    if (key == null) {
      // There are no keys to read yet.
      return 0;
    }
    // Do the read from Cassandra.
    // Bind the select statement.
    String customerId = key.asString();
    String storeId = Integer.toString(rand.nextInt(appConfig.numStores));
    BoundStatement select = getPreparedSelect().bind(customerId, storeId);
    ResultSet rs = getCassandraClient().execute(select);
    List rows = rs.all();
    LOG.debug("Read coupon count: " + rows.size());
    return 1;
  }

  protected PreparedStatement getPreparedInsert(String insertStmt)  {
    if (preparedInsert == null) {
      synchronized (prepareInitLock) {
        if (preparedInsert == null) {
          // Create the prepared statement object.
          preparedInsert = getCassandraClient().prepare(insertStmt);
        }
      }
    }
    return preparedInsert;
  }

  protected PreparedStatement getPreparedInsert()  {
    return getPreparedInsert(
      String.format(
        "INSERT INTO %s " +
        "(customer_id, store_id, coupon_code, begin_date, end_date, relevance_score) " +
        "VALUES (?, ?, ?, ?, ?, ?);",
        getTableName()));
  }

  private double generateRandomRelevanceScore() {
    return rand.nextDouble();
  }

  private int generateRandomCustomerScore() {
    return rand.nextInt(10);
  }

  @Override
  public long doWrite(int threadIdx) {
    BatchStatement batch = new BatchStatement();
    PreparedStatement insert = getPreparedInsert();
    Key key = getSimpleLoadGenerator().getKeyToWrite();
    try {
      int totalCouponCount = 0;
      for (int i = 0; i < appConfig.numStores; i++) {
        String customerId = key.asString();
        String storeId = Integer.toString(i);
        int couponCount = appConfig.numNewCouponsPerCustomer / appConfig.numStores;
        for (int j = 0; j < couponCount; j++) {
          Coupon coupon = coupons.elementAt(j);
          batch.add(insert.bind(customerId, storeId, coupon.code, coupon.beginDate, coupon.endDate,
                                Double.valueOf(generateRandomRelevanceScore())));
        }
        totalCouponCount += couponCount;
      }
      ResultSet resultSet = getCassandraClient().execute(batch);
      LOG.debug("Wrote coupon count: " + totalCouponCount + ", return code: " +
                resultSet.toString());
      getSimpleLoadGenerator().recordWriteSuccess(key);
      return 1;
    } catch (Exception e) {
      getSimpleLoadGenerator().recordWriteFailure(key);
      throw e;
    }
  }

  @Override
  public synchronized void resetClients() {
    synchronized (prepareInitLock) {
      preparedInsert = null;
      preparedSelect = null;
    }
    super.resetClients();
  }

  @Override
  public synchronized void destroyClients() {
    synchronized (prepareInitLock) {
      preparedInsert = null;
      preparedSelect = null;
    }
    super.destroyClients();
  }

  @Override
  public void appendMessage(StringBuilder sb) {
    super.appendMessage(sb);
    sb.append("maxWrittenKey: " + getSimpleLoadGenerator().getMaxWrittenKey() +  " | ");
    sb.append("maxGeneratedKey: " + getSimpleLoadGenerator().getMaxGeneratedKey() +  " | ");
  }

  public void appendParentMessage(StringBuilder sb) {
    super.appendMessage(sb);
  }

  @Override
  public List getWorkloadDescription() {
    return Arrays.asList(
      "User personalization app. Writes unique customer ids, each with a set of coupons for different stores.",
      " There are multiple readers and writers that update these keys and read them indefinitely. ",
      "Note that the number of reads and writes to perform can be specified as a parameter.");
  }

  @Override
  public List getWorkloadOptionalArguments() {
    return Arrays.asList(
      "--num_unique_keys " + appConfig.numUniqueKeysToWrite,
      "--num_reads " + appConfig.numKeysToRead,
      "--num_writes " + appConfig.numKeysToWrite,
      "--num_threads_read " + appConfig.numReaderThreads,
      "--num_threads_write " + appConfig.numWriterThreads,
      "--num_stores " + appConfig.numStores,
      "--num_new_coupons_per_customer " + appConfig.numNewCouponsPerCustomer,
      "--max_coupons_per_customer " + appConfig.maxCouponsPerCustomer);
  }
}