1ea70ef53SMark Nunberg/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
2ea70ef53SMark Nunberg/*
3ea70ef53SMark Nunberg*     Copyright 2015 Couchbase, Inc
4ea70ef53SMark Nunberg*
5ea70ef53SMark Nunberg*   Licensed under the Apache License, Version 2.0 (the "License");
6ea70ef53SMark Nunberg*   you may not use this file except in compliance with the License.
7ea70ef53SMark Nunberg*   You may obtain a copy of the License at
8ea70ef53SMark Nunberg*
9ea70ef53SMark Nunberg*       http://www.apache.org/licenses/LICENSE-2.0
10ea70ef53SMark Nunberg*
11ea70ef53SMark Nunberg*   Unless required by applicable law or agreed to in writing, software
12ea70ef53SMark Nunberg*   distributed under the License is distributed on an "AS IS" BASIS,
13ea70ef53SMark Nunberg*   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14ea70ef53SMark Nunberg*   See the License for the specific language governing permissions and
15ea70ef53SMark Nunberg*   limitations under the License.
16ea70ef53SMark Nunberg*/
17ea70ef53SMark Nunberg
18e649f75aSMark Nunberg#include <stdio.h>
19e649f75aSMark Nunberg#include <stdlib.h>
20e649f75aSMark Nunberg#include <assert.h>
21e649f75aSMark Nunberg#include <string.h>
22e649f75aSMark Nunberg#include <limits.h>
23e649f75aSMark Nunberg#include <errno.h>
2497b6ad02SMark Nunberg#include <stdexcept>
25e649f75aSMark Nunberg
26e649f75aSMark Nunberg#include <map>
2797b6ad02SMark Nunberg#include <vector>
28e649f75aSMark Nunberg#include <string>
29e649f75aSMark Nunberg#include <fstream>
30e649f75aSMark Nunberg#include <iostream>
31e649f75aSMark Nunberg#include <istream>
32e649f75aSMark Nunberg#include <ostream>
33e649f75aSMark Nunberg#include <sstream>
344a6d1501SMark Nunberg#include <chrono>
35e649f75aSMark Nunberg
36e649f75aSMark Nunberg#define CLIOPTS_ENABLE_CXX
378593b9a5SMark Nunberg#define INCLUDE_SUBDOC_NTOHLL
38e649f75aSMark Nunberg#include "subdoc/subdoc-api.h"
39e649f75aSMark Nunberg#include "subdoc/path.h"
40e649f75aSMark Nunberg#include "subdoc/match.h"
41e649f75aSMark Nunberg#include "subdoc/operations.h"
42e649f75aSMark Nunberg#include "contrib/cliopts/cliopts.h"
43e649f75aSMark Nunberg
44e649f75aSMark Nunbergusing std::string;
4597b6ad02SMark Nunbergusing std::vector;
46e649f75aSMark Nunbergusing std::map;
47e649f75aSMark Nunbergusing namespace cliopts;
48e649f75aSMark Nunberg
49614ae2faSMark Nunbergusing Subdoc::Path;
50614ae2faSMark Nunbergusing Subdoc::Operation;
51c8b99cceSMark Nunbergusing Subdoc::Command;
52c8b99cceSMark Nunbergusing Subdoc::Error;
53cefdbdb2SMark Nunbergusing Subdoc::Result;
54614ae2faSMark Nunberg
55a80d9184SMark Nunbergstruct OpEntry {
56a80d9184SMark Nunberg    uint8_t opcode;
57a80d9184SMark Nunberg    const char *description;
58a80d9184SMark Nunberg    OpEntry(uint8_t opcode = 0, const char *description = NULL) :
59a80d9184SMark Nunberg        opcode(opcode),
60a80d9184SMark Nunberg        description(description) {}
61a80d9184SMark Nunberg
62a80d9184SMark Nunberg    operator uint8_t() const { return opcode; }
63a80d9184SMark Nunberg};
64a80d9184SMark Nunberg
65e649f75aSMark Nunbergclass Options {
66e649f75aSMark Nunbergpublic:
67e649f75aSMark Nunberg    Options() :
68e649f75aSMark Nunberg        o_iter('i', "iterations", 1000),
69e649f75aSMark Nunberg        o_path('p', "docpath"),
70e649f75aSMark Nunberg        o_value('v', "value"),
71e649f75aSMark Nunberg        o_jsfile('f', "json"),
72e649f75aSMark Nunberg        o_cmd('c', "command"),
73f2ea93beSMark Nunberg        o_mkdirp('M', "create-intermediate"),
74ed9a8654SMark Nunberg        o_txtscan('T', "text-scan"),
75e649f75aSMark Nunberg        parser("subdoc-bench")
76e649f75aSMark Nunberg    {
77e649f75aSMark Nunberg        o_iter.description("Number of iterations to run");
78a80d9184SMark Nunberg        o_path.description("Document path to manipulate");
79e649f75aSMark Nunberg        o_value.description("Document value to insert");
80a80d9184SMark Nunberg        o_jsfile.description("JSON files to operate on. If passing multiple files, each file should be delimited by a comma");
81a80d9184SMark Nunberg        o_cmd.description("Command to use. Use -c help to show all the commands").mandatory();
82a80d9184SMark Nunberg        o_mkdirp.description("Create intermediate paths for mutation operations");
83ed9a8654SMark Nunberg        o_txtscan.description("Simply scan the text using a naive approach. Used to see how much actual overhead jsonsl induces");
84e649f75aSMark Nunberg
85e649f75aSMark Nunberg        parser.addOption(o_iter);
86e649f75aSMark Nunberg        parser.addOption(o_path);
87e649f75aSMark Nunberg        parser.addOption(o_value);
88e649f75aSMark Nunberg        parser.addOption(o_jsfile);
89e649f75aSMark Nunberg        parser.addOption(o_cmd);
90a80d9184SMark Nunberg        parser.addOption(o_mkdirp);
91ed9a8654SMark Nunberg        parser.addOption(o_txtscan);
92e649f75aSMark Nunberg
9397b6ad02SMark Nunberg        totalBytes = 0;
94e649f75aSMark Nunberg        // Set the opmap
95e649f75aSMark Nunberg        initOpmap();
96e649f75aSMark Nunberg    }
97e649f75aSMark Nunberg
98e649f75aSMark Nunberg    void initOpmap() {
99e649f75aSMark Nunberg        // generic ops:
100c8b99cceSMark Nunberg        opmap["replace"] = OpEntry(Command::REPLACE, "Replace a value");
101c8b99cceSMark Nunberg        opmap["delete"] = OpEntry(Command::REMOVE, "Delete a value");
102c8b99cceSMark Nunberg        opmap["get"] = OpEntry(Command::GET, "Retrieve a value");
103c8b99cceSMark Nunberg        opmap["exists"] = OpEntry(Command::EXISTS, "Check if a value exists");
104e649f75aSMark Nunberg
105e649f75aSMark Nunberg        // dict ops
106c8b99cceSMark Nunberg        opmap["add"] = OpEntry(Command::DICT_ADD, "Create a new dictionary value");
107c8b99cceSMark Nunberg        opmap["upsert"] = OpEntry(Command::DICT_UPSERT, "Create or replace a dictionary value");
108e649f75aSMark Nunberg
109e649f75aSMark Nunberg        // list ops
110c8b99cceSMark Nunberg        opmap["append"] = OpEntry(Command::ARRAY_APPEND, "Insert values to the end of an array");
111c8b99cceSMark Nunberg        opmap["prepend"] = OpEntry(Command::ARRAY_PREPEND, "Insert values to the beginning of an array");
112c8b99cceSMark Nunberg        opmap["addunique"] = OpEntry(Command::ARRAY_ADD_UNIQUE, "Add a unique value to an array");
113a9ab72ffSMark Nunberg        opmap["insert"] = OpEntry(Command::ARRAY_INSERT, "Insert value at given array index");
114e51529efSMark Nunberg        opmap["size"] = OpEntry(Command::GET_COUNT, "Count the number of items in an array or dict");
115e649f75aSMark Nunberg
116e649f75aSMark Nunberg        // arithmetic ops
117a113d89fSMark Nunberg        opmap["counter"] = OpEntry(Command::COUNTER);
118a113d89fSMark Nunberg
119a113d89fSMark Nunberg        // Generic ops
120a80d9184SMark Nunberg        opmap["path"] = OpEntry(0xff, "Check the validity of a path");
121e649f75aSMark Nunberg    }
122e649f75aSMark Nunberg
123e649f75aSMark Nunberg    UIntOption o_iter;
124e649f75aSMark Nunberg    StringOption o_path;
125e649f75aSMark Nunberg    StringOption o_value;
126e649f75aSMark Nunberg    StringOption o_jsfile;
127e649f75aSMark Nunberg    StringOption o_cmd;
128a80d9184SMark Nunberg    BoolOption o_mkdirp;
129ed9a8654SMark Nunberg    BoolOption o_txtscan;
130a80d9184SMark Nunberg    map<string,OpEntry> opmap;
131e649f75aSMark Nunberg    Parser parser;
13297b6ad02SMark Nunberg    size_t totalBytes;
133e649f75aSMark Nunberg};
134e649f75aSMark Nunberg
135e649f75aSMark Nunbergstatic void
13697b6ad02SMark NunbergreadJsonFile(string& name, vector<string>& out)
137e649f75aSMark Nunberg{
13897b6ad02SMark Nunberg    std::ifstream input(name.c_str());
13997b6ad02SMark Nunberg    if (input.fail()) {
14097b6ad02SMark Nunberg        throw name + ": " + strerror(errno);
141e649f75aSMark Nunberg    }
14297b6ad02SMark Nunberg    fprintf(stderr, "Reading %s\n", name.c_str());
143e649f75aSMark Nunberg    std::stringstream ss;
144e649f75aSMark Nunberg    ss << input.rdbuf();
14597b6ad02SMark Nunberg    out.push_back(ss.str());
146e649f75aSMark Nunberg    input.close();
14797b6ad02SMark Nunberg}
14897b6ad02SMark Nunberg
14997b6ad02SMark Nunbergstatic void
15097b6ad02SMark NunbergexecOperation(Options& o)
15197b6ad02SMark Nunberg{
15297b6ad02SMark Nunberg    vector<string> fileNames;
15397b6ad02SMark Nunberg    vector<string> inputStrs;
15497b6ad02SMark Nunberg    string flist = o.o_jsfile.const_result();
15597b6ad02SMark Nunberg    if (flist.find(',') == string::npos) {
15697b6ad02SMark Nunberg        fileNames.push_back(flist);
15797b6ad02SMark Nunberg    } else {
15897b6ad02SMark Nunberg        while (true) {
15997b6ad02SMark Nunberg            size_t pos = flist.find_first_of(',');
16097b6ad02SMark Nunberg            if (pos == string::npos) {
16197b6ad02SMark Nunberg                fileNames.push_back(flist);
16297b6ad02SMark Nunberg                break;
16397b6ad02SMark Nunberg            } else {
16497b6ad02SMark Nunberg                string curName = flist.substr(0, pos);
16597b6ad02SMark Nunberg                fileNames.push_back(curName);
16697b6ad02SMark Nunberg                flist = flist.substr(pos+1);
16797b6ad02SMark Nunberg            }
16897b6ad02SMark Nunberg        }
16997b6ad02SMark Nunberg    }
17097b6ad02SMark Nunberg    if (fileNames.empty()) {
17197b6ad02SMark Nunberg        throw string("At least one file must be passed!");
17297b6ad02SMark Nunberg    }
17397b6ad02SMark Nunberg    for (size_t ii = 0; ii < fileNames.size(); ii++) {
17497b6ad02SMark Nunberg        readJsonFile(fileNames[ii], inputStrs);
17597b6ad02SMark Nunberg        o.totalBytes += inputStrs.back().length();
17697b6ad02SMark Nunberg    }
177e649f75aSMark Nunberg
178a80d9184SMark Nunberg    uint8_t opcode = o.opmap[o.o_cmd.result()];
179a80d9184SMark Nunberg    if (o.o_mkdirp.passed()) {
180a80d9184SMark Nunberg        opcode |= 0x80;
181a80d9184SMark Nunberg    }
182a80d9184SMark Nunberg
183e649f75aSMark Nunberg    string value = o.o_value.const_result();
184e649f75aSMark Nunberg    string path = o.o_path.const_result();
1858188fb0bSMark Nunberg    Operation op;
186cefdbdb2SMark Nunberg    Result res;
187e649f75aSMark Nunberg
188a113d89fSMark Nunberg
189ed9a8654SMark Nunberg    size_t nquotes = 0;
190ed9a8654SMark Nunberg    bool is_txtscan = o.o_txtscan.result();
191e649f75aSMark Nunberg    size_t itermax = o.o_iter.result();
192ed9a8654SMark Nunberg
193ed9a8654SMark Nunberg    int char_table[256] = { 0 };
19439aac313STrond Norbye    char_table[static_cast<uint8_t >('"')] = 1;
19539aac313STrond Norbye    char_table[static_cast<uint8_t >('\\')] = 1;
19639aac313STrond Norbye    char_table[static_cast<uint8_t >('!')] = 1;
197ed9a8654SMark Nunberg
198e649f75aSMark Nunberg    for (size_t ii = 0; ii < itermax; ii++) {
19997b6ad02SMark Nunberg        const string& curInput = inputStrs[ii % inputStrs.size()];
200a113d89fSMark Nunberg
201ed9a8654SMark Nunberg        if (is_txtscan) {
202ed9a8654SMark Nunberg            const char *buf = curInput.c_str();
203ed9a8654SMark Nunberg            size_t nbytes = curInput.size();
204ed9a8654SMark Nunberg            for (; nbytes; buf++, nbytes--) {
205ed9a8654SMark Nunberg                if (char_table[static_cast<unsigned char>(*buf)]) {
206ed9a8654SMark Nunberg                    nquotes++;
207ed9a8654SMark Nunberg                }
208ed9a8654SMark Nunberg            }
209ed9a8654SMark Nunberg            continue;
210ed9a8654SMark Nunberg        }
211ed9a8654SMark Nunberg
212ed9a8654SMark Nunberg        op.clear();
213ed9a8654SMark Nunberg        res.clear();
214a113d89fSMark Nunberg        op.set_value(value);
2158188fb0bSMark Nunberg        op.set_code(opcode);
2168188fb0bSMark Nunberg        op.set_doc(curInput);
217cefdbdb2SMark Nunberg        op.set_result_buf(&res);
218e649f75aSMark Nunberg
2198188fb0bSMark Nunberg        Error rv = op.op_exec(path);
220c8b99cceSMark Nunberg        if (!rv.success()) {
2217bf46751SDave Rigby            throw rv;
222e649f75aSMark Nunberg        }
223e649f75aSMark Nunberg    }
224e649f75aSMark Nunberg
225ed9a8654SMark Nunberg    if (nquotes) {
226ed9a8654SMark Nunberg        printf("Found %lu quotes!\n", nquotes);
227ed9a8654SMark Nunberg    }
228ed9a8654SMark Nunberg
229e649f75aSMark Nunberg    // Print the result.
230e51529efSMark Nunberg    if (opcode == Command::GET ||
231e51529efSMark Nunberg            opcode == Command::EXISTS ||
232e51529efSMark Nunberg            opcode == Command::GET_COUNT) {
233296596ccSMark Nunberg        string match = res.matchloc().to_string();
2344a1393c4SMark Nunberg        printf("%s\n", match.c_str());
235e649f75aSMark Nunberg    } else {
236e649f75aSMark Nunberg        string newdoc;
237296596ccSMark Nunberg        for (auto ii : res.newdoc()) {
23810cb006bSMark Nunberg            newdoc.append(ii.at, ii.length);
239e649f75aSMark Nunberg        }
2404a1393c4SMark Nunberg        printf("%s\n", newdoc.c_str());
241e649f75aSMark Nunberg    }
242e649f75aSMark Nunberg}
243e649f75aSMark Nunberg
244e649f75aSMark Nunbergstatic void
245e649f75aSMark NunbergexecPathParse(Options& o)
246e649f75aSMark Nunberg{
247e649f75aSMark Nunberg    size_t itermax = o.o_iter.result();
248e649f75aSMark Nunberg    string path = o.o_path.const_result();
249614ae2faSMark Nunberg    Path pth;
250e649f75aSMark Nunberg
251e649f75aSMark Nunberg    for (size_t ii = 0; ii < itermax; ii++) {
252614ae2faSMark Nunberg        pth.clear();
253614ae2faSMark Nunberg        int rv = pth.parse(path);
254e649f75aSMark Nunberg
255e649f75aSMark Nunberg        if (rv != 0) {
256e649f75aSMark Nunberg            throw string("Failed to parse path!");
257e649f75aSMark Nunberg        }
258e649f75aSMark Nunberg    }
259e649f75aSMark Nunberg}
260e649f75aSMark Nunberg
261e649f75aSMark Nunbergvoid runMain(int argc, char **argv)
262e649f75aSMark Nunberg{
2634a6d1501SMark Nunberg    using namespace std::chrono;
2644a6d1501SMark Nunberg
265e649f75aSMark Nunberg    Options o;
266e649f75aSMark Nunberg    if (!o.parser.parse(argc, argv)) {
267e649f75aSMark Nunberg        throw string("Bad options!");
268e649f75aSMark Nunberg    }
269e649f75aSMark Nunberg    // Determine the command
270e649f75aSMark Nunberg    string cmdStr = o.o_cmd.const_result();
2714a6d1501SMark Nunberg    auto t_begin = steady_clock::now();
272e649f75aSMark Nunberg
273a80d9184SMark Nunberg    if (cmdStr == "help") {
274a80d9184SMark Nunberg        map<string,OpEntry>::const_iterator iter = o.opmap.begin();
275a80d9184SMark Nunberg        for (; iter != o.opmap.end(); ++iter) {
276a80d9184SMark Nunberg            const OpEntry& ent = iter->second;
277a80d9184SMark Nunberg            fprintf(stderr, "%s (0x%x): ", iter->first.c_str(), ent.opcode);
278a80d9184SMark Nunberg            fprintf(stderr, "%s\n", ent.description);
279a80d9184SMark Nunberg        }
280a80d9184SMark Nunberg        exit(EXIT_SUCCESS);
281a80d9184SMark Nunberg    }
282a80d9184SMark Nunberg    if (!o.o_path.passed()) {
283a80d9184SMark Nunberg        fprintf(stderr, "Path (-p) required\n");
284a80d9184SMark Nunberg        exit(EXIT_FAILURE);
285a80d9184SMark Nunberg    }
286a80d9184SMark Nunberg
287e649f75aSMark Nunberg    if (o.opmap.find(cmdStr) != o.opmap.end()) {
288e649f75aSMark Nunberg        if (!o.o_jsfile.passed()) {
289e649f75aSMark Nunberg            throw string("Operation must contain file!");
290e649f75aSMark Nunberg        }
291e649f75aSMark Nunberg        execOperation(o);
292e649f75aSMark Nunberg    } else if (cmdStr == "path") {
293e649f75aSMark Nunberg        execPathParse(o);
294e649f75aSMark Nunberg    } else {
295e649f75aSMark Nunberg        throw string("Unknown command!");
296e649f75aSMark Nunberg    }
297e649f75aSMark Nunberg
2984a6d1501SMark Nunberg    auto total = duration_cast<duration<double>>(steady_clock::now() - t_begin);
2994a6d1501SMark Nunberg
300e649f75aSMark Nunberg    // Get the number of seconds:
3014a6d1501SMark Nunberg    double n_seconds = total.count();
3024a6d1501SMark Nunberg    double ops_per_sec = static_cast<double>(o.o_iter.result()) / n_seconds;
3034a6d1501SMark Nunberg    double mb_per_sec = (
3044a6d1501SMark Nunberg            static_cast<double>(o.totalBytes) *
3054a6d1501SMark Nunberg            static_cast<double>(o.o_iter.result())) /
3064a6d1501SMark Nunberg                    n_seconds;
3074a6d1501SMark Nunberg
30897b6ad02SMark Nunberg    mb_per_sec /= (1024 * 1024);
309e649f75aSMark Nunberg
310e649dfa3SMark Nunberg    fprintf(stderr, "DURATION=%.2fs. OPS=%u\n", n_seconds, o.o_iter.result());
311e649dfa3SMark Nunberg    fprintf(stderr, "%.2f OPS/s\n",  ops_per_sec);
312e649dfa3SMark Nunberg    fprintf(stderr, "%.2f MB/s\n", mb_per_sec);
313e649f75aSMark Nunberg}
314e649f75aSMark Nunberg
315e649f75aSMark Nunbergint main(int argc, char **argv)
316e649f75aSMark Nunberg{
317e649f75aSMark Nunberg    try {
318e649f75aSMark Nunberg        runMain(argc, argv);
319e649f75aSMark Nunberg        return EXIT_SUCCESS;
320e649f75aSMark Nunberg    } catch (string& exc) {
3214a1393c4SMark Nunberg        fprintf(stderr, "%s\n", exc.c_str());
322e649f75aSMark Nunberg        return EXIT_FAILURE;
3238188fb0bSMark Nunberg    } catch (Error& rc) {
324c8b99cceSMark Nunberg        fprintf(stderr, "Command failed: %s\n", rc.description());
325a76064ecSMark Nunberg        return EXIT_FAILURE;
326d35adfc9SMark Nunberg    } catch (std::exception& ex) {
327d35adfc9SMark Nunberg        fprintf(stderr, "Command failed: %s\n", ex.what());
328a76064ecSMark Nunberg        return EXIT_FAILURE;
329e649f75aSMark Nunberg    }
330e649f75aSMark Nunberg}