Initial commit for cockroachdb driver

driver-cockroachdb/pom.xml

@@ -0,0 +1,35 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <parent>
+        <groupId>io.nosqlbench</groupId>
+        <artifactId>mvn-defaults</artifactId>
+        <version>4.15.8-SNAPSHOT</version>
+        <relativePath>../mvn-defaults</relativePath>
+    </parent>
+
+    <artifactId>driver-cockroachdb</artifactId>
+    <packaging>jar</packaging>
+    <name>${project.artifactId}</name>
+
+    <description>
+        A CockroachDB ActivityType driver for http://nosqlbench.io/
+    </description>
+
+    <dependencies>
+        <dependency>
+            <groupId>io.nosqlbench</groupId>
+            <artifactId>driver-jdbc</artifactId>
+            <version>4.15.8-SNAPSHOT</version>
+        </dependency>
+        <dependency>
+            <groupId>org.postgresql</groupId>
+            <artifactId>postgresql</artifactId>
+            <version>42.2.14</version>
+        </dependency>
+    </dependencies>
+
+</project>

driver-cockroachdb/src/main/java/io/nosqlbench/activitytype/cockroachdb/CockroachActivity.java

@@ -0,0 +1,56 @@
+package io.nosqlbench.activitytype.cockroachdb;
+
+import io.nosqlbench.activitytype.jdbc.api.JDBCActivity;
+import io.nosqlbench.engine.api.activityimpl.ActivityDef;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.postgresql.ds.PGSimpleDataSource;
+
+import javax.sql.DataSource;
+import java.sql.SQLException;
+import java.util.Arrays;
+
+public class CockroachActivity extends JDBCActivity {
+    private static final Logger LOGGER = LogManager.getLogger(CockroachActivity.class);
+
+    public CockroachActivity(ActivityDef activityDef) {
+        super(activityDef);
+    }
+
+    @Override
+    protected DataSource newDataSource() {
+        PGSimpleDataSource ds = new PGSimpleDataSource();
+
+        // serverName is required
+        String serverName = getParams().
+            getOptionalString("serverName").
+            orElseThrow(() -> new RuntimeException("serverName parameter required"));
+
+        // portNumber, user, password are optional
+        Integer portNumber = getParams().getOptionalInteger("portNumber").orElse(26257);
+        String user = getParams().getOptionalString("user").orElse(null);
+        String password = getParams().getOptionalString("password").orElse(null);
+
+        ds.setServerNames(new String[]{serverName});
+        ds.setPortNumbers(new int[]{portNumber});
+        if (user != null) {
+            ds.setUser(user);
+        }
+        if (password != null) {
+            ds.setPassword(password);
+        }
+
+        LOGGER.debug("Final DataSource fields"
+            + " serverNames=" + Arrays.toString(ds.getServerNames())
+            + " portNumbers=" + Arrays.toString(ds.getPortNumbers())
+            + " user=" + ds.getUser()
+            + " password=" + ds.getPassword());
+
+        return ds;
+    }
+
+    @Override
+    public boolean isRetryable(SQLException sqlException) {
+        return sqlException.getSQLState().equals("40001"); // sql state code for transaction conflict
+    }
+}

driver-cockroachdb/src/main/java/io/nosqlbench/activitytype/cockroachdb/CockroachActivityType.java

@@ -0,0 +1,25 @@
+package io.nosqlbench.activitytype.cockroachdb;
+
+import io.nosqlbench.activitytype.jdbc.api.JDBCActionDispenser;
+import io.nosqlbench.engine.api.activityapi.core.ActionDispenser;
+import io.nosqlbench.engine.api.activityapi.core.ActivityType;
+import io.nosqlbench.engine.api.activityimpl.ActivityDef;
+import io.nosqlbench.nb.annotations.Service;
+
+@Service(ActivityType.class)
+public class CockroachActivityType implements ActivityType<CockroachActivity> {
+    @Override
+    public String getName() {
+        return "cockroachdb";
+    }
+
+    @Override
+    public ActionDispenser getActionDispenser(CockroachActivity activity) {
+        return new JDBCActionDispenser(activity);
+    }
+
+    @Override
+    public CockroachActivity getActivity(ActivityDef activityDef) {
+        return new CockroachActivity(activityDef);
+    }
+}

driver-cockroachdb/src/main/resources/activities/cockroachdb-basic.yaml

@@ -0,0 +1,69 @@
+# java -jar nb.jar run driver=cockroachdb workload=cockroachdb-basic tags=phase:rampup cycles=10 serverName=localhost
+# java -jar nb.jar run driver=cockroachdb workload=cockroachdb-basic tags=phase:main cycles=10 serverName=localhost
+description: An example of a basic cockroach insert
+
+scenarios:
+  default:
+    - run driver=cockroachdb tags==phase:main threads==auto cycles===<<main-cycles:1000000>>
+  rampup:
+    - run driver=cockroachdb tags==phase:rampup threads==auto cycles===<<rampup-cycles:1000000>>
+  schema:
+    - run driver=cockroachdb tags==phase:schema threads==1 cycles===2
+    #- run driver=stdout tags==phase:schema threads==1 cycles===UNDEF
+
+bindings:
+  seq_key: Mod(<<keyCount:1000000>>L); ToInt()
+  seq_value: Mod(<<valueCount:1000000000>>L); <<valueSizeDist:Hash()>>; ToString() -> String
+  rw_key: <<keyDist:Uniform(0,1000000)->long>>; ToInt()
+  rw_value: <<valDist:Uniform(0,1000000000)->int>>; <<valueSizeDist:Hash()>>; ToString() -> String
+
+blocks:
+  - name: schema
+    tags:
+      phase: schema
+    params:
+    statements:
+      - create-database: |
+          CREATE DATABASE <<database:bank>>;
+        tags:
+          name: create-database
+      - create-table: |
+          CREATE TABLE IF NOT EXISTS <<database:bank>>.<<table:banktransaction>> (
+            code STRING PRIMARY KEY,
+            amount INTEGER
+          );
+        tags:
+          name: create-table
+  - name: rampup
+    tags:
+      phase: rampup
+    params:
+    statements:
+      - rampup-insert: insert into <<database:bank>>.<<table:banktransaction>> ( code, amount ) values ( '{seq_key}', {seq_value} );
+        params:
+        tags:
+          name: rampup-insert
+  - name: main-read
+    tags:
+      phase: main
+      type: read
+    params:
+      ratio: <<read_ratio:1>>
+    statements:
+      - main-find: |
+          SELECT code, amount FROM <<database:bank>>.<<table:banktransaction>> WHERE code = '{rw_key}' AND amount = {rw_value};
+        params:
+        tags:
+          name: main-find
+  - name: main-write
+    tags:
+      phase: main
+      type: write
+    params:
+      ratio: <<write_ratio:1>>
+    statements:
+      - main-insert: |
+          UPDATE <<database:bank>>.<<table:banktransaction>> SET amount = {seq_value} WHERE code = '{seq_key}';
+        params:
+        tags:
+          name: main-insert

driver-cockroachdb/src/main/resources/cockroachdb.md

@@ -0,0 +1,16 @@
+# CockroachDB Driver
+
+This is a driver for CockroachDB. It extends the generic JDBC Driver and inherits its parameters.
+
+### CockroachDB driver parameters
+
+All parameters correspond to the postgresql JDBC library parameters. See the
+[DataSource Configuration Properties](https://jdbc.postgresql.org/documentation/81/ds-ds.html)
+section for detailed parameter documentation.
+
+* **serverName** (required) - database hostname
+* **portNumber** (optional) - database listen port; defaults to 26257.
+* **user** (optional) - database account username; defaults to empty.
+* **password** (optional) - database account password; defaults to empty.
+* **connectionpool** (optional) - connection pool implementation; defaults to no connection pool. Allowed values:
+  * *hikari* - use [HikariCP](https://github.com/brettwooldridge/HikariCP)

driver-jdbc/pom.xml

@@ -0,0 +1,29 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>nosqlbench</artifactId>
+        <groupId>io.nosqlbench</groupId>
+        <version>4.15.8-SNAPSHOT</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>driver-jdbc</artifactId>
+
+    <dependencies>
+        <dependency>
+            <groupId>com.zaxxer</groupId>
+            <artifactId>HikariCP</artifactId>
+            <version>3.4.5</version>
+        </dependency>
+        <dependency>
+            <groupId>io.nosqlbench</groupId>
+            <artifactId>engine-api</artifactId>
+            <version>4.15.8-SNAPSHOT</version>
+            <scope>compile</scope>
+        </dependency>
+    </dependencies>
+
+
+</project>

driver-jdbc/src/main/java/io/nosqlbench/activitytype/jdbc/api/JDBCActionDispenser.java

@@ -0,0 +1,18 @@
+package io.nosqlbench.activitytype.jdbc.api;
+
+import io.nosqlbench.activitytype.jdbc.impl.JDBCAction;
+import io.nosqlbench.engine.api.activityapi.core.Action;
+import io.nosqlbench.engine.api.activityapi.core.ActionDispenser;
+
+public class JDBCActionDispenser implements ActionDispenser {
+    private final JDBCActivity activity;
+
+    public JDBCActionDispenser(JDBCActivity a) {
+        activity = a;
+    }
+
+    @Override
+    public Action getAction(int slot) {
+        return new JDBCAction(activity, slot);
+    }
+}

driver-jdbc/src/main/java/io/nosqlbench/activitytype/jdbc/api/JDBCActivity.java

@@ -0,0 +1,121 @@
+package io.nosqlbench.activitytype.jdbc.api;
+
+import com.codahale.metrics.Counter;
+import com.codahale.metrics.Histogram;
+import com.codahale.metrics.Timer;
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+import io.nosqlbench.activitytype.jdbc.impl.ReadyJDBCOp;
+import io.nosqlbench.engine.api.activityapi.planning.OpSequence;
+import io.nosqlbench.engine.api.activityimpl.ActivityDef;
+import io.nosqlbench.engine.api.activityimpl.SimpleActivity;
+import io.nosqlbench.engine.api.metrics.ActivityMetrics;
+import io.nosqlbench.engine.api.metrics.ExceptionCountMetrics;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import javax.sql.DataSource;
+import java.sql.SQLException;
+import java.util.HashMap;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+public abstract class JDBCActivity extends SimpleActivity {
+    private final static Logger LOGGER = LogManager.getLogger(JDBCActivity.class);
+    private Timer bindTimer;
+    private Timer resultTimer;
+    private Timer resultSuccessTimer;
+    private Histogram triesHisto;
+    private ExceptionCountMetrics exceptionCount;
+    private SQLExceptionCountMetrics sqlExceptionCount;
+
+    protected DataSource dataSource;
+    protected OpSequence<ReadyJDBCOp> opSequence;
+
+    public JDBCActivity(ActivityDef activityDef) {
+        super(activityDef);
+    }
+
+    /*
+    Subclasses construct a DataSource object. Concrete type should *not* be a pooled DataSource,
+    as this class implements wrapping with HikariDataSource if required.
+     */
+    protected abstract DataSource newDataSource();
+
+    @Override
+    public synchronized void onActivityDefUpdate(ActivityDef activityDef) {
+        super.onActivityDefUpdate(activityDef);
+
+        LOGGER.debug("initializing data source");
+        dataSource = newDataSource();
+
+        String connectionPool = getParams().getOptionalString("connectionpool").orElse("");
+        if (!connectionPool.isEmpty()) {
+            LOGGER.debug("initializing connectionpool " + connectionPool);
+            if (connectionPool.equals("hikari")) {
+                HikariConfig config = new HikariConfig();
+                config.setDataSource(dataSource);
+                dataSource = new HikariDataSource(config);
+            } else {
+                throw new RuntimeException("unknown connectionpool parameter value " + connectionPool);
+            }
+        }
+    }
+
+    @Override
+    public void initActivity() {
+        LOGGER.debug("initializing activity: " + getActivityDef().getAlias());
+        bindTimer = ActivityMetrics.timer(getActivityDef(), "bind");
+        resultTimer = ActivityMetrics.timer(getActivityDef(), "result");
+        resultSuccessTimer = ActivityMetrics.timer(getActivityDef(), "result-success");
+        triesHisto = ActivityMetrics.histogram(getActivityDef(), "tries");
+        exceptionCount = new ExceptionCountMetrics(getActivityDef());
+        sqlExceptionCount = new SQLExceptionCountMetrics(getActivityDef());
+
+        opSequence = createOpSequence(ReadyJDBCOp::new);
+        setDefaultsFromOpSequence(opSequence);
+
+        onActivityDefUpdate(getActivityDef());
+    }
+
+    public int getMaxTries() { return 3; }
+
+    public boolean isRetryable(SQLException sqlException) {
+        return true;
+    }
+
+    public DataSource getDataSource() {
+        return dataSource;
+    }
+
+    public OpSequence<ReadyJDBCOp> getOpSequence() { return opSequence; }
+
+    public Timer getBindTimer() { return bindTimer; }
+
+    public Timer getResultTimer() { return resultTimer; }
+
+    public Timer getResultSuccessTimer() { return resultSuccessTimer; }
+
+    public Histogram getTriesHisto() { return triesHisto; }
+
+    public ExceptionCountMetrics getExceptionCount() { return exceptionCount; }
+
+    public SQLExceptionCountMetrics getSQLExceptionCount() { return sqlExceptionCount; }
+
+    public static class SQLExceptionCountMetrics {
+        private final ConcurrentHashMap<Integer, Counter> counters = new ConcurrentHashMap<>();
+        private final ActivityDef activityDef;
+
+        private SQLExceptionCountMetrics(ActivityDef activityDef) {
+            this.activityDef = activityDef;
+        }
+
+        public void inc(SQLException e) {
+            Counter c = counters.computeIfAbsent(
+                e.getErrorCode(),
+                k -> ActivityMetrics.counter(activityDef, "errorcodecounts." + k)
+            );
+            c.inc();
+        }
+    }
+}

driver-jdbc/src/main/java/io/nosqlbench/activitytype/jdbc/impl/JDBCAction.java

@@ -0,0 +1,94 @@
+package io.nosqlbench.activitytype.jdbc.impl;
+
+import com.codahale.metrics.Timer;
+import io.nosqlbench.activitytype.jdbc.api.JDBCActivity;
+import io.nosqlbench.engine.api.activityapi.core.SyncAction;
+import io.nosqlbench.engine.api.activityapi.planning.OpSequence;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.sql.*;
+import java.util.concurrent.TimeUnit;
+
+public class JDBCAction implements SyncAction {
+    private static final Logger LOGGER = LogManager.getLogger(JDBCAction.class);
+
+    private final JDBCActivity activity;
+    private OpSequence<ReadyJDBCOp> sequencer;
+
+    public JDBCAction(JDBCActivity a, int slot) {
+        activity = a;
+    }
+
+    @Override
+    public void init() {
+        sequencer = activity.getOpSequence();
+    }
+
+    @Override
+    public int runCycle(long cycle) {
+        String boundStmt;
+
+        ReadyJDBCOp unboundStmt = sequencer.get(cycle);
+
+        try (Timer.Context bindTime = activity.getBindTimer().time()) {
+            boundStmt = unboundStmt.apply(cycle);
+        }
+
+        int maxTries = activity.getMaxTries();
+        int errorCode = 0;
+
+        for (int tries = 1; tries <= maxTries; tries++) {
+            errorCode = execute(boundStmt, tries);
+            if (errorCode == 0) return 0;
+        }
+
+        LOGGER.debug("Max tries " + maxTries + " exceeded for executing statement " + boundStmt);
+        return errorCode;
+    }
+
+    private int execute(String sql, int tries) {
+            long startTimeNano = System.nanoTime();
+            Long resultTime = null;
+
+            try (Connection conn = activity.getDataSource().getConnection()) {
+                Statement jdbcStmt = conn.createStatement();
+                jdbcStmt.execute(sql);
+
+                resultTime = System.nanoTime() - startTimeNano;
+                activity.getResultSuccessTimer().update(resultTime, TimeUnit.NANOSECONDS);
+
+            } catch (Exception e) {
+                LOGGER.debug("Try " + tries + ": failed to execute statement: " + sql, e);
+
+                activity.getExceptionCount().count(e);
+
+                if (e instanceof SQLException) {
+                    SQLException sqle = (SQLException) e;
+
+                    activity.getSQLExceptionCount().inc(sqle);
+
+                    // TODO non-retryable exception should return its non-zero error code to runCycle() caller
+                    if (!activity.isRetryable(sqle)) {
+                        return 0;
+                    }
+
+                    return sqle.getErrorCode();
+                }
+
+                return 1;
+
+            } finally {
+                if (resultTime == null) {
+                    resultTime = System.nanoTime() - startTimeNano;
+                }
+
+                activity.getResultTimer().update(resultTime, TimeUnit.NANOSECONDS);
+                activity.getTriesHisto().update(tries);
+            }
+
+            LOGGER.trace("Try " + tries + ": successfully executed statement: " + sql);
+            return 0;
+    }
+
+}

driver-jdbc/src/main/java/io/nosqlbench/activitytype/jdbc/impl/ReadyJDBCOp.java

@@ -0,0 +1,26 @@
+package io.nosqlbench.activitytype.jdbc.impl;
+
+import io.nosqlbench.engine.api.activityconfig.yaml.OpTemplate;
+import io.nosqlbench.virtdata.core.bindings.BindingsTemplate;
+import io.nosqlbench.virtdata.core.templates.ParsedTemplate;
+import io.nosqlbench.virtdata.core.templates.StringBindings;
+import io.nosqlbench.virtdata.core.templates.StringBindingsTemplate;
+
+import java.util.function.LongFunction;
+
+public class ReadyJDBCOp implements LongFunction<String> {
+    private final StringBindings bindings;
+
+    public ReadyJDBCOp(OpTemplate stmtDef) {
+        ParsedTemplate paramTemplate = new ParsedTemplate(stmtDef.getStmt(), stmtDef.getBindings());
+        BindingsTemplate paramBindings = new BindingsTemplate(paramTemplate.getBindPoints());
+        StringBindingsTemplate template = new StringBindingsTemplate(stmtDef.getStmt(), paramBindings);
+
+        bindings = template.resolve();
+    }
+
+    @Override
+    public String apply(long cycle) {
+        return bindings.bind(cycle);
+    }
+}

nb/pom.xml

@@ -129,6 +129,12 @@
     <!--            <version>2.12.51-SNAPSHOT</version>-->
     <!--        </dependency>-->
 
+      <dependency>
+          <groupId>io.nosqlbench</groupId>
+          <artifactId>driver-cockroachdb</artifactId>
+          <version>4.15.8-SNAPSHOT</version>
+      </dependency>
+
     <dependency>
       <groupId>javax.activation</groupId>
       <artifactId>activation</artifactId>

pom.xml

@@ -53,6 +53,7 @@
         <module>driver-kafka</module>
         <module>driver-mongodb</module>
         <module>driver-jmx</module>
+        <module>driver-jdbc</module>
 
         <!-- VIRTDATA MODULES -->
 
@@ -65,6 +66,7 @@
         <module>virtdata-lib-realer</module>
 
         <module>virtdata-userlibs</module>
+        <module>driver-cockroachdb</module>
 
     </modules>