Auto run benchmarks

.github/workflows/benchmarks.yml

@@ -0,0 +1,28 @@
+name: Benchmarks - stargate-mongoose
+on:
+  pull_request:
+  push:
+jobs:
+  benchmarks:
+    runs-on: ubuntu-latest
+    steps:      
+      - name: disable and stop mono-xsp4.service
+        run: |
+          sudo kill -9 $(sudo lsof -t -i:8084)
+      - uses: actions/checkout@v3
+      - name: Set up JSON API
+        run: |
+          chmod +x bin/start_json_api.sh
+          bin/start_json_api.sh
+      - name: Wait for services
+        run: |
+          while ! nc -z localhost 8181; do sleep 1; done
+          while ! nc -z localhost 8081; do sleep 1; done
+      - name: Setting up the node version
+        uses: actions/setup-node@v3
+        with:
+          node-version: 16.15.0
+      - name: setup project
+        run: npm i   
+      - name: benchmarks
+        run: npm run benchmarks

.gitignore

@@ -5,4 +5,6 @@ node_modules
 .vscode
 dist
 package-lock.json
-.idea
+.idea
+*.out
+*.tgz

.npmignore

@@ -7,4 +7,6 @@ tests
 .env.example
 .github
 
-*.tgz
+*.tgz
+
+benchmarks

benchmarks/.env.benchmark

@@ -0,0 +1,4 @@
+JSON_API_URI=http://127.0.0.1:8181/v1/benchmark
+JSON_API_USERNAME=cassandra
+JSON_API_PASSWORD=cassandra
+JSON_API_AUTH_URL=http://127.0.0.1:8081/v1/auth

benchmarks/README.md

@@ -0,0 +1,95 @@
+# Benchmarks
+
+Each `benchmark-*.ts` file is a separate benchmark that tests different functionality.
+Benchmarks use either stargate-mongoose or the Axios HTTP client directly; this is to check whether stargate-mongoose is adding significant overhead or not.
+
+## Running Benchmarks
+
+The easiest way to run all benchmarks is using the following npm command:
+
+```
+npm run benchmarks
+```
+
+This will run all benchmarks, and output results to both stdout and `.out` files in the `benchmarks` directory.
+
+You can also run benchmarks individually using `ts-node` as follows:
+
+```
+$ ./node_modules/.bin/ts-node -r ./benchmarks/setup.ts ./benchmarks/benchmark-findone-mongoose.ts
+{
+  "name": "benchmark-findone-mongoose",
+  "totalTimeMS": 14944
+}
+```
+
+The `benchmarks/.env.benchmark` contains environment variables that configure the connection details for the benchmarks.
+The `benchmarks/setup.ts` file sets up the environment using [dotenv](https://npmjs.com/package/dotenv).
+
+## Astra Benchmarks
+
+There are 3 benchmarks that connect to Astra, rather than local JSON API:
+
+* `benchmark-findone-axios-vector-astra.ts`
+* `benchmark-findone-mongoose-vector-astra.ts`
+* `benchmark-findone-mongoose-vector-astra-http2.ts`
+
+By default, these benchmarks are skipped:
+
+```
+$ ./node_modules/.bin/ts-node -r ./benchmarks/setup.ts ./benchmarks/benchmark-findone-mongoose-vector-astra
+{"name":"benchmark-findone-mongoose-vector-astra"}
+```
+
+To run these benchmarks, you need to add the following environment variables to `.env.benchmark`:
+
+```
+ASTRA_CONNECTION_STRING=
+ASTRA_COLLECTION_NAME=
+```
+
+`ASTRA_CONNECTION_STRING` contains the connection string that will be passed to `mongoose.connect()`, like `https://abc-us-east-2.apps.astra.datastax.com/api/json/v1/mykeyspace?applicationToken=AstraCS:...`, and `ASTRA_COLLECTION_NAME` is the name of the collection that the Mongoose model will connect to.
+The Astra benchmarks assume that there is already a collection with data set up in Astra.
+
+With these environment variables set up, you can run the Astra benchmarks:
+
+```
+$ ./node_modules/.bin/ts-node -r ./benchmarks/setup.ts ./benchmarks/benchmark-findone-mongoose-vector-astra
+{
+  "name": "benchmark-findone-mongoose-vector-astra",
+  "totalTimeMS": 20966
+}
+```
+
+If you add `ASTRA_CONNECTION_STRING` to `.env.benchmark`, be careful not to accidentally commit your changes to git.
+An alternative approach is to set the environment variables in your CLI command as follows:
+
+```
+$ env ASTRA_CONNECTION_STRING="https://abc-us-east-2.apps.astra.datastax.com/..." env ASTRA_COLLECTION_NAME=mycollection ./node_modules/.bin/ts-node -r ./benchmarks/setup.ts ./benchmarks/benchmark-findone-mongoose-vector-astra
+{
+  "name": "benchmark-findone-mongoose-vector-astra",
+  "totalTimeMS": 20436
+}
+```
+
+## Pinecone Benchmarks
+
+First, install the Pinecone client:
+
+```
+npm install @pinecone-database/pinecone --no-save
+```
+
+Then set up sample data:
+
+```
+```
+env PINECONE_INDEX=your-index-here env PINECONE_API_KEY=your-api-key-here ./node_modules/.bin/ts-node benchmarks/setup-pinecone-vector-example.ts
+```
+```
+
+Then run Pinecone benchmarks:
+
+```
+env PINECONE_INDEX=your-index-here env PINECONE_API_KEY=your-api-key-here ./node_modules/.bin/ts-node benchmarks/benchmark-findone-pinecone.ts
+```

benchmarks/benchmark-findone-axios-vector-astra.ts

@@ -0,0 +1,76 @@
+import axios from 'axios';
+import { driver } from '../';
+import mongoose from 'mongoose';
+
+mongoose.set('autoCreate', false);
+mongoose.set('autoIndex', false);
+
+mongoose.setDriver(driver);
+
+main().then(
+    () => { process.exit(0); },
+    err => {
+        console.error(err);
+        process.exit(-1);
+    }
+);
+
+async function main() {
+    if (!process.env.ASTRA_CONNECTION_STRING) {
+        console.log('{"name":"benchmark-findone-axios-vector-astra"}');
+        return;
+    }
+    await mongoose.connect(process.env.ASTRA_CONNECTION_STRING, {
+        isAstra: true
+    } as mongoose.ConnectOptions);
+
+    const Vector = mongoose.model(
+        'Vector',
+        new mongoose.Schema({
+            $vector: {
+                type: [Number]
+            },
+            prompt: {
+                type: String,
+                required: true
+            }
+        }, {
+            autoCreate: false
+        }),
+        process.env.ASTRA_COLLECTION_NAME
+    );
+
+    // @ts-ignore
+    const dbName = Vector.db.db.name;
+
+    const collectionName = process.env.ASTRA_COLLECTION_NAME ?? '';
+
+    // @ts-ignore
+    const baseUrl = mongoose.connection.getClient().httpClient.baseUrl;
+
+    // @ts-ignore
+    const token = mongoose.connection.getClient().httpClient.applicationToken;
+    const $meta = [1, ...Array(1535).fill(0)];
+    const start = Date.now();
+    for (let i = 0; i < 100; ++i) {
+        await axios.post(
+            `${baseUrl}/${dbName}/${collectionName}`,
+            {
+                findOne: {
+                    filter: {},
+                    sort: { $vector: $meta }
+                }
+            },
+            {
+                headers: {
+                    Token: token
+                }
+            }
+        );
+    }
+    const results = {
+        name: 'benchmark-findone-axios-vector-astra',
+        totalTimeMS: Date.now() - start
+    };
+    console.log(JSON.stringify(results, null, '  '));
+}

benchmarks/benchmark-findone-axios-vector.ts

@@ -0,0 +1,81 @@
+import axios from 'axios';
+import { driver } from '../';
+import mongoose from 'mongoose';
+
+mongoose.set('autoCreate', false);
+mongoose.set('autoIndex', false);
+
+mongoose.setDriver(driver);
+
+main().then(
+    () => { process.exit(0); },
+    err => {
+        console.error(err);
+        process.exit(-1);
+    }
+);
+
+async function main() {
+    await mongoose.connect(process.env.JSON_API_URI ?? '', {
+        username: process.env.JSON_API_USERNAME,
+        password: process.env.JSON_API_PASSWORD,
+        authUrl: process.env.JSON_API_AUTH_URL
+    } as mongoose.ConnectOptions);
+    const Vector = mongoose.model(
+        'Vector',
+        new mongoose.Schema({
+            $vector: {
+                type: [Number]
+            },
+            prompt: {
+                type: String,
+                required: true
+            }
+        }, {
+            collectionOptions: { vector: { size: 1536, function: 'cosine' } },
+            autoCreate: false
+        }),
+        'vectors'
+    );
+
+    // @ts-ignore
+    const dbName = Vector.db.db.name;
+
+    await Vector.db.dropCollection('vectors').catch(() => {});
+    await Vector.createCollection();
+
+    const numVectors = 1000;
+    for (let i = 0; i < numVectors; ++i) {
+        const $vector = Array(1536).fill(0);
+        $vector[i] = 1;
+        await Vector.create({ $vector, prompt: `Test ${i}` });
+    }
+
+    // @ts-ignore
+    const baseUrl = mongoose.connection.getClient().httpClient.baseUrl;
+    // @ts-ignore
+    const token = mongoose.connection.getClient().httpClient.applicationToken;
+    const $meta = [1, ...Array(1535).fill(0)];
+    const start = Date.now();
+    for (let i = 0; i < 100; ++i) {
+        await axios.post(
+            `${baseUrl}/${dbName}/vectors`,
+            {
+                findOne: {
+                    filter: {},
+                    sort: { $vector: $meta }
+                }
+            },
+            {
+                headers: {
+                    Token: token
+                }
+            }
+        );
+    }
+    const results = {
+        name: 'benchmark-findone-axios-vector',
+        totalTimeMS: Date.now() - start
+    };
+    console.log(JSON.stringify(results, null, '  '));
+}

benchmarks/benchmark-findone-axios.ts

@@ -0,0 +1,75 @@
+import axios from 'axios';
+import { driver } from '../';
+import mongoose from 'mongoose';
+
+mongoose.set('autoCreate', false);
+mongoose.set('autoIndex', false);
+
+mongoose.setDriver(driver);
+
+main().then(
+    () => { process.exit(0); },
+    err => {
+        console.error(err);
+        process.exit(-1);
+    }
+);
+
+async function main() {
+    const uri = process.env.JSON_API_URI ?? '';
+    await mongoose.connect(uri, {
+        username: process.env.JSON_API_USERNAME,
+        password: process.env.JSON_API_PASSWORD,
+        authUrl: process.env.JSON_API_AUTH_URL
+    } as mongoose.ConnectOptions);
+    const Test = mongoose.model('Test', new mongoose.Schema({
+        name: {
+            type: String,
+            required: true
+        },
+        email: {
+            type: String,
+            required: true
+        },
+        age: {
+            type: Number
+        }
+    }));
+
+    await Test.db.dropCollection('tests').catch(() => {});
+    await Test.createCollection();
+
+    await Test.create({
+        name: 'John Smith',
+        email: '[email protected]',
+        age: 30
+    });
+
+    // @ts-ignore
+    const token = mongoose.connection.getClient().httpClient.applicationToken;
+
+    const start = Date.now();
+    for (let i = 0; i < 10000; ++i) {
+        await axios.post(
+            `${uri}/tests`,
+            {
+                findOne: {
+                    filter: {
+                        name: 'John Smith'
+                    },
+                    options: {}
+                }
+            },
+            {
+                headers: {
+                    Token: token
+                }
+            }
+        );
+    }
+    const results = {
+        name: 'benchmark-findone-axios',
+        totalTimeMS: Date.now() - start
+    };
+    console.log(JSON.stringify(results, null, '  '));
+}