Merge pull request #8 from moxious/record-consumer

Batch data faking utilities for load testing

Author: moxious

README.md

@@ -20,6 +20,9 @@ Usage: run-workload.js -p password
    [--ms milliseconds] how many milliseconds to test for
    [--workload /path/to/workload.json] probability table spec
    [--query CYPHER_QUERY] single cypher query to run
+   [--schema /path/to/schema.json] schema for generated records (only used with
+   --query)
+   [--batchsize [1000]] number of records from schema to generate per batch
    [--concurrency c] how many concurrent queries to run (default: 10)
    [--checkpoint cn] how often to print results in milliseconds (default: 5000)
    [--fail-fast] if specified, the work will stop after encountering one
@@ -35,6 +38,9 @@ Options:
   -a             address to connect to                    [default: "localhost"]
   -u             username                                     [default: "neo4j"]
   -p             password                                             [required]
+  -d             database
+  --schema       batch schema file
+  --batchsize    number of records per batch, usable only with schema
   -n             number of hits on the database
   --ms           number of milliseconds to execute
   --workload     absolute path to JSON probability table/workload
@@ -69,12 +75,28 @@ program arguments, as in, `npm run graph-workload -- --n 20`
 
 # Examples
 
-Create a lot of nodes as fast as possible:
+## Create a lot of nodes as fast as possible:
 
 ```
 npm run graph-workload -- -a localhost -u neo4j -p admin --query 'Unwind range(1,1000000) as id create (n);' -n 50 --concurrency 4
 ```
 
+## Write custom data generated in batches 
+
+Fake/mock data can be generated with functions from [fakerjs](https://www.npmjs.com/package/faker).
+
+Using this technique you can generate your own data and create
+custom load patterns.  Similar to other Neo4j utilities, the batch will be present in the query form: "UNWIND batch AS event".
+
+```
+npm run graph-workload -- -a localhost -u neo4j -p admin \
+  --query 'CREATE (t:Test) SET t += event' \
+  --batchsize 1000 \
+  --schema /absolute/path/to/schemas/myschema.json
+```
+
+See `src/schemas/user.json` as an example of a schema you can use in this way.  Keys are field names to generate, values are the faker functions used to populate that field.
+
 # Neo4j 4.0 / Multidatabase
 
 As of Neo4j 4.0, sessions support multi-database.  Use the `-d` or `--database` argument to direct

package.json

@@ -1,6 +1,6 @@
 {
   "name": "graph-workload",
-  "version": "0.4.1",
+  "version": "0.5.0",
   "description": "Tool for generating workloads running on Neo4j",
   "main": "src/run-workload.js",
   "scripts": {
@@ -18,14 +18,14 @@
   },
   "homepage": "https://github.com/moxious/graph-workload",
   "dependencies": {
-    "bluebird": "^3.5.3",
+    "bluebird": "^3.7.2",
     "dimsum": "^0.2.2",
     "es6-promise-pool": "^2.5.0",
     "faker": "^4.1.0",
     "generic-pool": "^3.6.1",
-    "lodash": "^4.17.15",
+    "lodash": "^4.17.19",
     "moment": "^2.24.0",
-    "neo4j-driver": "^4.0.2",
+    "neo4j-driver": "^4.1.1",
     "randomstring": "^1.1.5",
     "uuid": "^3.3.2",
     "yargs": "^12.0.5"

src/SimpleQueryStrategy.js

@@ -1,4 +1,5 @@
 const Strategy = require('./Strategy');
+const _ = require('lodash');
 
 /**
  * Represents a container class for a strategy that is just running some
@@ -10,6 +11,8 @@ module.exports = class SimpleQueryStrategy extends Strategy {
         this.name = props.name || 'SimpleQuery';
         this.query = props.query;
         this.params = props.params || {};
+        this.generator = props.generator;
+        this.batchSize = props.batchSize;
 
         if (!(props.mode === 'READ') && !(props.mode === 'WRITE')) {
             throw new Error('Mode must be READ or WRITE');
@@ -20,7 +23,17 @@ module.exports = class SimpleQueryStrategy extends Strategy {
 
     run() {
         const f = (s) => {
-            const txRunner = tx => tx.run(this.query, this.params);
+            const txRunner = tx => {
+                if (this.generator) {
+                    const batch = this.generator.generate(this.batchSize);
+                    console.log('batchrun with ', batch.length, 'elements');
+                    return tx.run(
+                        `UNWIND $batch AS event ${this.query}`,
+                        _.merge({ batch }, this.params)
+                    );
+                }
+                return tx.run(this.query, this.params);
+            };
 
             if (this.props.mode === 'READ') {
                 return s.readTransaction(txRunner);

src/datasource/MockData.js

@@ -0,0 +1,54 @@
+const _ = require('lodash');
+const faker = require('faker');
+
+class MockData {
+    /**
+     * Create a mock data source.
+     * @param {Object} schema a map of field names to data types
+     */
+    constructor(schema) {
+        this.schema = _.cloneDeep(schema);
+        this.validateSchema();
+        this.columns = Object.keys(schema);
+    }
+
+    generate(batchSize=1) {
+        const r = [];
+        for (let i=0; i<batchSize; i++) {
+            const record = _.zipObject(this.columns, _.map(this.columns, col => this.schema[col]()));
+            r.push(record);
+        }
+        return r;
+    }
+
+    /**
+     * Validate that for each schema field, there is a corresponding faker function.
+     * For example, there is a faker.address.city() function.  So if the user specifies
+     * 'address.city' we will find that function.  
+     * @throws Error when an unidentified function is specified.
+     */
+    validateSchema() {
+        if (_.isNil(this.schema) || _.isEmpty(this.schema)) {
+            throw new Error('Empty or invalid schema specified');
+        }
+
+        Object.keys(this.schema).forEach(fieldName => {
+            const val = this.schema[fieldName];
+            if (!val) { throw new Error(`Field name ${fieldName} in schema has no specification`); }
+            const parts = val.split('.');
+
+            let f = faker;
+            parts.forEach(part => {
+                f = f[part];
+                if (!f) {
+                    throw new Error(`Invalid schema: Unknown function part '${part}' in field type ${fieldName}`);
+                }
+            });
+
+            // If we've worked through all parts we have our terminal function.
+            this.schema[fieldName] = f;
+        });
+    }
+}
+
+module.exports = MockData;

src/run-configuration.js

@@ -1,6 +1,7 @@
 const terminateAfter = require('./termination-condition');
 const ProbabilityTable = require('./stats/ProbabilityTable');
 const _ = require('lodash');
+const MockData = require('./datasource/MockData');
 
 const usageStr = `
 Usage: run-workload.js -p password
@@ -10,6 +11,8 @@ Usage: run-workload.js -p password
    [--ms milliseconds] how many milliseconds to test for
    [--workload /path/to/workload.json] probability table spec
    [--query CYPHER_QUERY] single cypher query to run
+   [--schema /path/to/schema.json] schema for generated records (only used with --query)
+   [--batchsize [1000]] number of records from schema to generate per batch
    [--concurrency c] how many concurrent queries to run (default: 10)
    [--checkpoint cn] how often to print results in milliseconds (default: 5000)
    [--fail-fast] if specified, the work will stop after encountering one failure.
@@ -106,7 +109,9 @@ const generateFromArgs = (args) => {
         // We don't know where to connect...
         (!process.env.NEO4J_URI && !args.a) ||
         // Don't know what password to use...
-        (!process.env.NEO4J_PASSWORD && !args.p)
+        (!process.env.NEO4J_PASSWORD && !args.p) ||
+        // You can't specify a schema if you don't have a query to consume it.
+        (args.schema && !args.query)
     );
 
     if (badlyConfigured) {
@@ -118,10 +123,15 @@ const generateFromArgs = (args) => {
     const concurrency = chooseConcurrency(args);
     const probabilityTable = chooseProbabilityTable(args);
 
+    const schema = args.schema ? require(args.schema) : null;
+    const batchSize = args.batchSize ? args.batchSize : 1000;
     const failFast = ('fail-fast' in args) ? args['fail-fast'] : false;
 
     // Merge sub-objects.
     const obj = _.merge({
+        generator: schema ? new MockData(schema) : null,
+        batchSize,
+        runcheckpoint: args.runcheckpoint,
         checkpointFreq: args.checkpoint || process.env.CHECKPOINT_FREQUENCY || 5000,
         failFast,
         phase: 'NOT_STARTED',
@@ -149,12 +159,16 @@ module.exports = {
             .describe('u', 'username')
             .describe('p', 'password')
             .describe('d', 'database')
+            .describe('runcheckpoint', 'whether to run db checkpointing or not')
+            .describe('schema', 'batch schema file')
+            .describe('batchsize', 'number of records per batch, usable only with schema')
             .describe('n', 'number of hits on the database')
             .describe('ms', 'number of milliseconds to execute')
             .describe('workload', 'absolute path to JSON probability table/workload')
             .describe('query', 'Cypher query to run')
             .default('concurrency', 10)
             .default('checkpoint', 5000)
+            .default('runcheckpoint', false)
             .demandOption(['p'])
             .argv;
     },

src/schemas/users.json

@@ -0,0 +1,9 @@
+{
+    "id": "random.uuid",
+    "firstName": "name.firstName",
+    "lastName": "name.lastName",
+    "email": "internet.email",
+    "job": "name.jobTitle",
+    "phone": "phone.phoneNumber",
+    "favoriteColor": "commerce.color"
+}

src/strategies.js

@@ -55,6 +55,8 @@ const builder = (sessionPool, runConfig) => {
             if (label === 'custom') {
                 chosenTable[label] = new SimpleQueryStrategy({
                     sessionPool,
+                    generator: runConfig.generator,
+                    batchSize: runConfig.batchSize,
                     query: runConfig.query,
                     mode: runConfig.mode || 'WRITE',
                 });

src/workload.js

@@ -76,7 +76,19 @@ class Workload {
     const runStrategySetupPromises = Object.keys(this.strategyTable)
       .map(stratName => this.strategyTable[stratName].setup(this.driver));
 
-    return Promise.all(runStrategySetupPromises);
+    return Promise.all(runStrategySetupPromises)
+      .then(results => {
+        if (this.runConfig.runcheckpoint) {
+          const opts = {};
+          if (this.runConfig.database) { opts.database = this.runConfig.database; }
+          const session = this.driver.session(opts);
+          console.log('Calling db.checkpoint method to flush prior to run');
+          return session.run('CALL db.checkpoint()')
+            .then(session.close);
+        }
+
+        return results;
+      });
   }
 
   getRunConfiguration() { return this.runConfig; }

test/datasource/MockData_test.js

@@ -0,0 +1,46 @@
+const expect = require('chai').expect;
+const mocks = require('../mocks');
+const chai = require('chai');
+const chaiAsPromised = require("chai-as-promised");
+chai.use(chaiAsPromised);
+
+const MockData = require('../../src/datasource/MockData');
+
+describe('Mock Data', function() {
+    const goodSchema = {
+        streetAddress: 'address.streetAddress',
+        city: 'address.city',
+        state: 'address.state',
+        country: 'address.country',
+    };
+
+    it('can create something with a good schema', () => {
+        const md = new MockData(goodSchema);
+        expect(md).to.be.ok;
+    });
+
+    it('refuses empty schema', () => {
+        expect(() => new MockData()).to.throw(Error);
+    });
+
+    it('refuses wrong schema', () => {
+        expect(() => new MockData({
+            foo: 'i.do.not.exist',
+        })).to.throw(Error);
+    });
+
+    it('can generate a batch', () => {
+        const md = new MockData(goodSchema);
+
+        console.log('generating');
+        const r = md.generate(10);
+        console.log(r);
+        // expect(r).to.be.an(Array);
+        expect(r.length).to.equal(10);
+        for (let i=0; i<r.length; i++) {
+            Object.keys(goodSchema).forEach(field => {
+                expect(r[i][field]).to.be.ok;
+            });
+        }
+    });
+});